2 files changed, 10 insertions, 11 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),
-        };
-    }
 }