From 756aca36ad74bc70bfb287bc3d2b1af922410269 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Fri, 3 Nov 2023 09:25:53 +0100
Subject: [PATCH] Take RoutingMode as argument

---
 common/wireguard/src/lib.rs                    | 17 ++++++++++++-----
 .../wireguard/src/platform/linux/tun_device.rs | 18 +++++++-----------
 service-providers/ip-packet-router/src/lib.rs  |  9 +++++----
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/common/wireguard/src/lib.rs b/common/wireguard/src/lib.rs
index bd5aa47b4f..9d734ce8d9 100644
--- a/common/wireguard/src/lib.rs
+++ b/common/wireguard/src/lib.rs
@@ -32,17 +32,24 @@ pub async fn start_wireguard(
     task_client: nym_task::TaskClient,
     gateway_client_registry: Arc<GatewayClientRegistry>,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
-    // We can either index peers by their IP like standard wireguard
+    // TODO: make this configurable
+
+    // We can optionally index peers by their IP like standard wireguard. If we don't then we do
+    // plain NAT where we match incoming destination IP with outgoing source IP.
     let peers_by_ip = Arc::new(tokio::sync::Mutex::new(network_table::NetworkTable::new()));
 
-    // ... or by their tunnel tag, which is a random number assigned to them
-    let peers_by_tag = Arc::new(tokio::sync::Mutex::new(wg_tunnel::PeersByTag::new()));
+    // Alternative 1:
+    let routing_mode = tun_device::RoutingMode::new_allowed_ips(peers_by_ip.clone());
+    // Alternative 2:
+    //let routing_mode = tun_device::RoutingMode::new_nat();
 
     // Start the tun device that is used to relay traffic outbound
-    let (tun, tun_task_tx, tun_task_response_rx) =
-        tun_device::TunDevice::new(Some(peers_by_ip.clone()));
+    let (tun, tun_task_tx, tun_task_response_rx) = tun_device::TunDevice::new(routing_mode);
     tun.start();
 
+    // We also index peers by a tag
+    let peers_by_tag = Arc::new(tokio::sync::Mutex::new(wg_tunnel::PeersByTag::new()));
+
     // If we want to have the tun device on a separate host, it's the tun_task and
     // tun_task_response channels that needs to be sent over the network to the host where the tun
     // device is running.
diff --git a/common/wireguard/src/platform/linux/tun_device.rs b/common/wireguard/src/platform/linux/tun_device.rs
index 196102394f..638c00ce3f 100644
--- a/common/wireguard/src/platform/linux/tun_device.rs
+++ b/common/wireguard/src/platform/linux/tun_device.rs
@@ -45,7 +45,7 @@ pub struct TunDevice {
     routing_mode: RoutingMode,
 }
 
-enum RoutingMode {
+pub enum RoutingMode {
     // The routing table, as how wireguard does it
     AllowedIps(AllowedIpsInner),
 
@@ -55,28 +55,29 @@ enum RoutingMode {
 }
 
 impl RoutingMode {
-    fn new_nat() -> Self {
+    pub fn new_nat() -> Self {
         RoutingMode::Nat(NatInner {
             nat_table: HashMap::new(),
         })
     }
 
-    fn new_allowed_ips(peers_by_ip: Arc<tokio::sync::Mutex<PeersByIp>>) -> Self {
+    pub fn new_allowed_ips(peers_by_ip: Arc<tokio::sync::Mutex<PeersByIp>>) -> Self {
         RoutingMode::AllowedIps(AllowedIpsInner { peers_by_ip })
     }
 }
 
-struct AllowedIpsInner {
+pub struct AllowedIpsInner {
     peers_by_ip: Arc<tokio::sync::Mutex<PeersByIp>>,
 }
 
-struct NatInner {
+pub struct NatInner {
     nat_table: HashMap<IpAddr, u64>,
 }
 
 impl TunDevice {
     pub fn new(
-        peers_by_ip: Option<Arc<tokio::sync::Mutex<PeersByIp>>>,
+        routing_mode: RoutingMode,
+        // peers_by_ip: Option<Arc<tokio::sync::Mutex<PeersByIp>>>,
     ) -> (Self, TunTaskTx, TunTaskResponseRx) {
         let tun = setup_tokio_tun_device(
             format!("{TUN_BASE_NAME}%d").as_str(),
@@ -89,11 +90,6 @@ impl TunDevice {
         let (tun_task_tx, tun_task_rx) = tun_task_channel();
         let (tun_task_response_tx, tun_task_response_rx) = tun_task_response_channel();
 
-        let routing_mode = match peers_by_ip {
-            Some(peers_by_ip) => RoutingMode::new_allowed_ips(peers_by_ip),
-            None => RoutingMode::new_nat(),
-        };
-
         let tun_device = TunDevice {
             tun_task_rx,
             tun_task_response_tx,
diff --git a/service-providers/ip-packet-router/src/lib.rs b/service-providers/ip-packet-router/src/lib.rs
index 5ef1ad41a0..2dd483817d 100644
--- a/service-providers/ip-packet-router/src/lib.rs
+++ b/service-providers/ip-packet-router/src/lib.rs
@@ -116,12 +116,13 @@ impl IpPacketRouterBuilder {
         let self_address = *mixnet_client.nym_address();
 
         // Create the TUN device that we interact with the rest of the world with
-        let (tun, tun_task_tx, tun_task_response_rx) =
-            nym_wireguard::tun_device::TunDevice::new(None);
+        let (tun, tun_task_tx, tun_task_response_rx) = nym_wireguard::tun_device::TunDevice::new(
+            nym_wireguard::tun_device::RoutingMode::new_nat(),
+        );
         tun.start();
 
         let ip_packet_router_service = IpPacketRouter {
-            config: self.config,
+            _config: self.config,
             // tun,
             tun_task_tx,
             tun_task_response_rx,
@@ -144,7 +145,7 @@ impl IpPacketRouterBuilder {
 }
 
 struct IpPacketRouter {
-    config: Config,
+    _config: Config,
     // tun: nym_wireguard::tun_device::TunDevice,
     tun_task_tx: nym_wireguard::tun_task_channel::TunTaskTx,
     tun_task_response_rx: nym_wireguard::tun_task_channel::TunTaskResponseRx,