Add derived_session

2 files changed, 19 insertions, 8 deletions

diff --git a/src/wireguard/timers.rs b/src/wireguard/timers.rs
index 2e9263d..485f466 100644
--- a/src/wireguard/timers.rs
+++ b/src/wireguard/timers.rs
@@ -4,7 +4,6 @@ use std::sync::Arc;
 use std::time::{Duration, Instant, SystemTime};
 
 use log::{debug, info};
-use spin::Mutex;
 use hjul::{Runner, Timer};
 
 use super::constants::*;
@@ -16,14 +15,14 @@ use super::types::KeyPair;
 
 pub struct Timers {
     handshake_attempts: AtomicUsize,
+    sent_lastminute_handshake: AtomicBool,
+    need_another_keepalive: AtomicBool,
 
     retransmit_handshake: Timer,
     send_keepalive: Timer,
     send_persistent_keepalive: Timer,
-    sent_lastminute_handshake: AtomicBool,
     zero_key_material: Timer,
     new_handshake: Timer,
-    need_another_keepalive: AtomicBool,
 }
 
 impl Timers {
@@ -82,8 +81,7 @@ impl<B: bind::Bind> PeerInner<B> {
         self.timers()
             .sent_lastminute_handshake
             .store(false, Ordering::SeqCst);
-        // TODO: Store time in peer for config
-        // self.walltime_last_handshake
+        *self.walltime_last_handshake.lock() = SystemTime::now();
     }
 
     /* Should be called after an ephemeral key is created, which is before sending a
@@ -106,6 +104,10 @@ impl<B: bind::Bind> PeerInner<B> {
         }
     }
 
+    pub fn timers_session_derieved(&self) {
+        self.timers().zero_key_material.reset(REJECT_AFTER_TIME * 3);
+    }
+
     /* Called after a handshake worker sends a handshake initiation to the peer
      */
     pub fn sent_handshake_initiation(&self) {
@@ -120,7 +122,7 @@ impl<B: bind::Bind> PeerInner<B> {
         *self.last_handshake_sent.lock() = Instant::now();
         self.timers_any_authenticated_packet_traversal();
         self.timers_any_authenticated_packet_sent();
-    }  
+    } 
 
     fn packet_send_queued_handshake_initiation(&self, is_retry: bool) {
         if !is_retry {
diff --git a/src/wireguard/wireguard.rs b/src/wireguard/wireguard.rs
index e308c50..674c96f 100644
--- a/src/wireguard/wireguard.rs
+++ b/src/wireguard/wireguard.rs
@@ -42,6 +42,7 @@ pub struct PeerInner<B: Bind> {
     pub id: u64,
 
     // handshake state
+    pub walltime_last_handshake: Mutex<SystemTime>,
     pub last_handshake_sent: Mutex<Instant>, // instant for last handshake
     pub handshake_queued: AtomicBool,        // is a handshake job currently queued for the peer?
     pub queue: Mutex<Sender<HandshakeJob<B::Endpoint>>>, // handshake queue
@@ -244,6 +245,7 @@ impl<T: Tun, B: Bind> Wireguard<T, B> {
         let state = Arc::new(PeerInner {
             id: rng.gen(),
             pk,
+            walltime_last_handshake: Mutex::new(SystemTime::UNIX_EPOCH),
             last_handshake_sent: Mutex::new(self.state.start - TIME_HORIZON),
             handshake_queued: AtomicBool::new(false),
             queue: Mutex::new(self.state.queue.lock().clone()),
@@ -443,9 +445,16 @@ impl<T: Tun, B: Bind> Wireguard<T, B> {
                                                 peer.state.sent_handshake_response();
                                             }
 
-                                            // add resulting keypair to peer
+                                            // add any new keypair to peer
                                             keypair.map(|kp| {
-                                                debug!("{} : handshake worker, new keypair", wg);
+                                                debug!(
+                                                    "{} : handshake worker, new keypair for {}",
+                                                    wg, peer
+                                                );
+
+                                                // this means that a handshake response was processed or sent
+                                                peer.timers_session_derieved();
+
                                                 // free any unused ids
                                                 for id in peer.router.add_keypair(kp) {
                                                     state.device.release(id);