From 3ba0247634bbaa1da61532ca43e67fb2ad6c1106 Mon Sep 17 00:00:00 2001
From: Mathias Hall-Andersen <mathias@hall-andersen.dk>
Date: Mon, 18 Nov 2019 13:13:55 +0100
Subject: Better compartmentalization of cryptokey router

---
 src/wireguard/router/device.rs | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

(limited to 'src/wireguard/router/device.rs')

diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 04b2045..7adcf8a 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -22,7 +22,7 @@ use super::types::{Callbacks, RouterError};
 use super::workers::{worker_parallel, JobParallel};
 use super::SIZE_MESSAGE_PREFIX;
 
-use super::route::get_route;
+use super::route::RoutingTable;
 
 use super::super::{bind, tun, Endpoint, KeyPair};
 
@@ -35,8 +35,7 @@ pub struct DeviceInner<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Write
 
     // routing
     pub recv: RwLock<HashMap<u32, Arc<DecryptionState<E, C, T, B>>>>, // receiver id -> decryption state
-    pub ipv4: RwLock<IpLookupTable<Ipv4Addr, Arc<PeerInner<E, C, T, B>>>>, // ipv4 cryptkey routing
-    pub ipv6: RwLock<IpLookupTable<Ipv6Addr, Arc<PeerInner<E, C, T, B>>>>, // ipv6 cryptkey routing
+    pub table: RoutingTable<PeerInner<E, C, T, B>>,
 
     // work queues
     pub queue_next: AtomicUsize, // next round-robin index
@@ -95,8 +94,7 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Device<E, C,
             queues: Mutex::new(Vec::with_capacity(num_workers)),
             queue_next: AtomicUsize::new(0),
             recv: RwLock::new(HashMap::new()),
-            ipv4: RwLock::new(IpLookupTable::new()),
-            ipv6: RwLock::new(IpLookupTable::new()),
+            table: RoutingTable::new(),
         };
 
         // start worker threads
@@ -157,7 +155,11 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Device<E, C,
         let packet = &msg[SIZE_MESSAGE_PREFIX..];
 
         // lookup peer based on IP packet destination address
-        let peer = get_route(&self.state, packet).ok_or(RouterError::NoCryptoKeyRoute)?;
+        let peer = self
+            .state
+            .table
+            .get_route(packet)
+            .ok_or(RouterError::NoCryptoKeyRoute)?;
 
         // schedule for encryption and transmission to peer
         if let Some(job) = peer.send_job(msg, true) {
-- 
cgit v1.2.3-59-g8ed1b