From f53e7fb3c3f46c80cefa7b89424f5f2244fe770d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Fri, 6 Oct 2023 14:52:01 +0200
Subject: [PATCH] Extract out network_table

---
 common/wireguard/src/lib.rs           | 33 ++++++++++++---------------
 common/wireguard/src/network_table.rs | 12 ++++++++++
 common/wireguard/src/wg_tunnel.rs     |  6 ++---
 3 files changed, 29 insertions(+), 22 deletions(-)
 create mode 100644 common/wireguard/src/network_table.rs

diff --git a/common/wireguard/src/lib.rs b/common/wireguard/src/lib.rs
index 19d8a03a63e..11e5caad0c7 100644
--- a/common/wireguard/src/lib.rs
+++ b/common/wireguard/src/lib.rs
@@ -1,43 +1,38 @@
 #![cfg_attr(not(target_os = "linux"), allow(dead_code))]
 
+use std::net::SocketAddr;
+
+use dashmap::DashMap;
+use network_table::NetworkTable;
 use nym_task::TaskClient;
+use tokio::sync::mpsc;
 
 mod error;
 mod event;
+mod network_table;
 mod platform;
 mod setup;
 mod udp_listener;
 mod wg_tunnel;
 
+use crate::event::Event;
+
 // Currently the module related to setting up the virtual network device is platform specific.
 #[cfg(target_os = "linux")]
 use platform::linux::tun_device;
 
-#[derive(Default)]
-struct AllowedIps<T> {
-    pub ips: ip_network_table::IpNetworkTable<T>,
-}
-
-impl<T> AllowedIps<T> {
-    fn new() -> Self {
-        Self {
-            ips: ip_network_table::IpNetworkTable::new(),
-        }
-    }
-}
-
-type ActivePeers =
-    dashmap::DashMap<std::net::SocketAddr, tokio::sync::mpsc::UnboundedSender<crate::event::Event>>;
-
-type PeersByIp = AllowedIps<tokio::sync::mpsc::UnboundedSender<crate::event::Event>>;
+type ActivePeers = DashMap<SocketAddr, mpsc::UnboundedSender<Event>>;
+type PeersByIp = NetworkTable<mpsc::UnboundedSender<Event>>;
 
 #[cfg(target_os = "linux")]
 pub async fn start_wireguard(
     task_client: TaskClient,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+    use std::sync::Arc;
+
     // The set of active tunnels indexed by the peer's address
-    let active_peers = std::sync::Arc::new(ActivePeers::new());
-    let peers_by_ip = std::sync::Arc::new(std::sync::Mutex::new(AllowedIps::new()));
+    let active_peers = Arc::new(ActivePeers::new());
+    let peers_by_ip = Arc::new(std::sync::Mutex::new(NetworkTable::new()));
 
     // Start the tun device that is used to relay traffic outbound
     let tun_task_tx = tun_device::start_tun_device(peers_by_ip.clone());
diff --git a/common/wireguard/src/network_table.rs b/common/wireguard/src/network_table.rs
new file mode 100644
index 00000000000..362d79eea37
--- /dev/null
+++ b/common/wireguard/src/network_table.rs
@@ -0,0 +1,12 @@
+#[derive(Default)]
+pub(crate) struct NetworkTable<T> {
+    pub ips: ip_network_table::IpNetworkTable<T>,
+}
+
+impl<T> NetworkTable<T> {
+    pub(crate) fn new() -> Self {
+        Self {
+            ips: ip_network_table::IpNetworkTable::new(),
+        }
+    }
+}
diff --git a/common/wireguard/src/wg_tunnel.rs b/common/wireguard/src/wg_tunnel.rs
index b68d33f58b1..b485ab7bc94 100644
--- a/common/wireguard/src/wg_tunnel.rs
+++ b/common/wireguard/src/wg_tunnel.rs
@@ -14,7 +14,7 @@ use tokio::{
     time::timeout,
 };
 
-use crate::{error::WgError, event::Event, AllowedIps};
+use crate::{error::WgError, event::Event, NetworkTable};
 
 const MAX_PACKET: usize = 65535;
 
@@ -29,7 +29,7 @@ pub struct WireGuardTunnel {
     endpoint: Arc<tokio::sync::RwLock<SocketAddr>>,
 
     // AllowedIPs for this peer
-    allowed_ips: AllowedIps<()>,
+    allowed_ips: NetworkTable<()>,
 
     // `boringtun` tunnel, used for crypto & WG protocol
     wg_tunnel: Arc<tokio::sync::Mutex<Tunn>>,
@@ -85,7 +85,7 @@ impl WireGuardTunnel {
         // Signal close tunnel
         let (close_tx, close_rx) = broadcast::channel(1);
 
-        let mut allowed_ips = AllowedIps::new();
+        let mut allowed_ips = NetworkTable::new();
         allowed_ips.ips.insert(peer_allowed_ips, ());
 
         let tunnel = WireGuardTunnel {