1 files changed, 65 insertions, 14 deletions

diff --git a/src/timers.rs b/src/timers.rs
index 67ece06..23cbb87 100644
--- a/src/timers.rs
+++ b/src/timers.rs
@@ -3,11 +3,13 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
+use log::info;
+
 use hjul::{Runner, Timer};
 
 use crate::constants::*;
 use crate::router::Callbacks;
-use crate::types::{tun, bind};
+use crate::types::{bind, tun};
 use crate::wireguard::{Peer, PeerInner};
 
 pub struct Timers {
@@ -16,8 +18,17 @@ pub struct Timers {
 
     retransmit_handshake: Timer,
     send_keepalive: Timer,
+    send_persistent_keepalive: Timer,
     zero_key_material: Timer,
     new_handshake: Timer,
+    need_another_keepalive: AtomicBool,
+}
+
+impl Timers {
+    #[inline(always)]
+    fn need_another_keepalive(&self) -> bool {
+        self.need_another_keepalive.swap(false, Ordering::SeqCst)
+    }
 }
 
 impl Timers {
@@ -28,34 +39,42 @@ impl Timers {
     {
         // create a timer instance for the provided peer
         Timers {
+            need_another_keepalive: AtomicBool::new(false),
             handshake_pending: AtomicBool::new(false),
             handshake_attempts: AtomicUsize::new(0),
             retransmit_handshake: {
                 let peer = peer.clone();
                 runner.timer(move || {
-                    if peer.timers.read().handshake_retry() {
+                    if peer.timers().handshake_retry() {
+                        info!("Retransmit handshake for {}", peer);
                         peer.new_handshake();
+                    } else {
+                        info!("Failed to complete handshake for {}", peer);
+                        peer.router.purge_staged_packets();
+                        peer.timers().send_keepalive.stop();
+                        peer.timers().zero_key_material.start(REJECT_AFTER_TIME * 3);
                     }
                 })
             },
-            new_handshake: {
+            send_keepalive: {
                 let peer = peer.clone();
                 runner.timer(move || {
-                    peer.new_handshake();
-                    peer.timers.read().handshake_begun();
+                    peer.router.send_keepalive();
+                    if peer.timers().need_another_keepalive() {
+                        peer.timers().send_keepalive.start(KEEPALIVE_TIMEOUT);
+                    }
                 })
             },
-            send_keepalive: {
+            new_handshake: {
                 let peer = peer.clone();
                 runner.timer(move || {
-                    peer.router.send_keepalive();
-                    let keepalive = peer.keepalive.load(Ordering::Acquire);
-                    if keepalive > 0 {
-                        peer.timers
-                            .read()
-                            .send_keepalive
-                            .reset(Duration::from_secs(keepalive as u64))
-                    }
+                    info!(
+                        "Retrying handshake with {}, because we stopped hearing back after {} seconds", 
+                        peer, 
+                        (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT).as_secs()
+                    );
+                    peer.new_handshake();
+                    peer.timers.read().handshake_begun();
                 })
             },
             zero_key_material: {
@@ -64,6 +83,17 @@ impl Timers {
                     peer.router.zero_keys();
                 })
             },
+            send_persistent_keepalive: {
+                let peer = peer.clone();
+                runner.timer(move || {
+                    let keepalive = peer.state.keepalive.load(Ordering::Acquire);
+                    if keepalive > 0 {
+                        peer.router.send_keepalive();
+                        peer.timers().send_keepalive.stop();
+                        peer.timers().send_persistent_keepalive.start(Duration::from_secs(keepalive as u64));
+                    }
+                })
+            }
         }
     }
 
@@ -83,6 +113,12 @@ impl Timers {
         }
     }
 
+    pub fn updated_persistent_keepalive(&self, keepalive: usize) {
+        if keepalive > 0 {
+            self.send_persistent_keepalive.reset(Duration::from_secs(keepalive as u64));
+        }
+    }
+
     pub fn dummy(runner: &Runner) -> Timers {
         Timers {
             handshake_pending: AtomicBool::new(false),
@@ -90,13 +126,28 @@ impl Timers {
             retransmit_handshake: runner.timer(|| {}),
             new_handshake: runner.timer(|| {}),
             send_keepalive: runner.timer(|| {}),
+            send_persistent_keepalive: runner.timer(|| {}),
             zero_key_material: runner.timer(|| {}),
+            need_another_keepalive: AtomicBool::new(false),
         }
     }
 
     pub fn handshake_sent(&self) {
         self.send_keepalive.stop();
     }
+
+
+    pub fn any_authenticatec_packet_recieved(&self) {
+
+    }
+
+    pub fn handshake_initiated(&self) {
+
+    }
+
+    pub fn handhsake_complete(&self) {
+        
+    }
 }
 
 /* Instance of the router callbacks */