Clean dependencies and imports

4 files changed, 21 insertions, 15 deletions

diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 7adcf8a..34273d5 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -1,6 +1,4 @@
 use std::collections::HashMap;
-
-use std::net::{Ipv4Addr, Ipv6Addr};
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::mpsc::sync_channel;
 use std::sync::mpsc::SyncSender;
@@ -10,7 +8,6 @@ use std::time::Instant;
 
 use log::debug;
 use spin::{Mutex, RwLock};
-use treebitmap::IpLookupTable;
 use zerocopy::LayoutVerified;
 
 use super::anti_replay::AntiReplay;
diff --git a/src/wireguard/router/route.rs b/src/wireguard/router/route.rs
index e5f5955..1c93009 100644
--- a/src/wireguard/router/route.rs
+++ b/src/wireguard/router/route.rs
@@ -25,6 +25,7 @@ impl<T> RoutingTable<T> {
         }
     }
 
+    // collect keys mapping to the given value
     fn collect<A>(table: &IpLookupTable<A, Arc<T>>, value: &Arc<T>) -> Vec<(A, u32)>
     where
         A: Address,
@@ -38,6 +39,13 @@ impl<T> RoutingTable<T> {
         res
     }
 
+    pub fn insert(&self, ip: IpAddr, cidr: u32, value: Arc<T>) {
+        match ip {
+            IpAddr::V4(v4) => self.ipv4.write().insert(v4.mask(cidr), cidr, value),
+            IpAddr::V6(v6) => self.ipv6.write().insert(v6.mask(cidr), cidr, value),
+        };
+    }
+
     pub fn list(&self, value: &Arc<T>) -> Vec<(IpAddr, u32)> {
         let mut res = vec![];
         res.extend(
@@ -55,10 +63,11 @@ impl<T> RoutingTable<T> {
 
     pub fn remove(&self, value: &Arc<T>) {
         let mut v4 = self.ipv4.write();
-        let mut v6 = self.ipv6.write();
         for (ip, cidr) in Self::collect(&*v4, value) {
             v4.remove(ip, cidr);
         }
+
+        let mut v6 = self.ipv6.write();
         for (ip, cidr) in Self::collect(&*v6, value) {
             v6.remove(ip, cidr);
         }
@@ -153,11 +162,4 @@ impl<T> RoutingTable<T> {
             _ => None,
         }
     }
-
-    pub fn insert(&self, ip: IpAddr, cidr: u32, value: Arc<T>) {
-        match ip {
-            IpAddr::V4(v4) => self.ipv4.write().insert(v4.mask(cidr), cidr, value),
-            IpAddr::V6(v6) => self.ipv6.write().insert(v6.mask(cidr), cidr, value),
-        };
-    }
 }
diff --git a/src/wireguard/timers.rs b/src/wireguard/timers.rs
index 8f6b3ee..bfd2583 100644
--- a/src/wireguard/timers.rs
+++ b/src/wireguard/timers.rs
@@ -36,6 +36,11 @@ impl Timers {
 }
 
 impl<T: tun::Tun, B: bind::Bind> PeerInner<T, B> {
+
+    pub fn get_keepalive_interval(&self) -> u64 {
+        self.timers().keepalive_interval
+    }
+
     pub fn stop_timers(&self) {
         // take a write lock preventing simultaneous timer events or "start_timers" call
         let mut timers = self.timers_mut();
@@ -191,7 +196,6 @@ impl<T: tun::Tun, B: bind::Bind> PeerInner<T, B> {
         self.timers_any_authenticated_packet_sent();
     } 
 
-
     pub fn set_persistent_keepalive_interval(&self, secs: u64) {
         let mut timers = self.timers_mut();
 
@@ -405,6 +409,6 @@ impl<T: tun::Tun, B: bind::Bind> Callbacks for Events<T, B> {
 
     #[inline(always)]
     fn key_confirmed(peer: &Self::Opaque) {
-        peer.timers().retransmit_handshake.stop();
+        peer.timers_handshake_complete();
     }
 }
diff --git a/src/wireguard/wireguard.rs b/src/wireguard/wireguard.rs
index 00ee053..e548c8a 100644
--- a/src/wireguard/wireguard.rs
+++ b/src/wireguard/wireguard.rs
@@ -368,7 +368,6 @@ impl<T: Tun, B: Bind> Wireguard<T, B> {
                     wg.pending.fetch_sub(1, Ordering::SeqCst);
 
                     let device = wg.handshake.read();
-
                     match job {
                         HandshakeJob::Message(msg, src) => {
                             // feed message to handshake device
@@ -418,10 +417,14 @@ impl<T: Tun, B: Bind> Wireguard<T, B> {
                                             // update endpoint
                                             peer.router.set_endpoint(src);
 
-                                            // update timers after sending handshake response
                                             if resp_len > 0 {
+                                                // update timers after sending handshake response
                                                 debug!("{} : handshake worker, handshake response sent", wg);
                                                 peer.state.sent_handshake_response();
+                                            } else {
+                                                // update timers after receiving handshake response
+                                                debug!("{} : handshake worker, handshake response was received", wg);
+                                                peer.state.timers_handshake_complete();
                                             }
 
                                             // add any new keypair to peer