From d9f088f36efdea101ba6a309bf565127bc748e5f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Fri, 20 Oct 2023 15:17:03 +0200
Subject: [PATCH] Fully wrap tun task channel in strong type (#4023)

---
 common/wireguard/src/lib.rs                   | 23 +---------------
 .../src/platform/linux/tun_device.rs          | 12 +++------
 common/wireguard/src/tun_task_channel.rs      | 26 +++++++++++++++++++
 common/wireguard/src/udp_listener.rs          |  2 +-
 common/wireguard/src/wg_tunnel.rs             |  3 ++-
 5 files changed, 33 insertions(+), 33 deletions(-)
 create mode 100644 common/wireguard/src/tun_task_channel.rs

diff --git a/common/wireguard/src/lib.rs b/common/wireguard/src/lib.rs
index 0df9e8ce1a..813ba4718f 100644
--- a/common/wireguard/src/lib.rs
+++ b/common/wireguard/src/lib.rs
@@ -10,6 +10,7 @@ mod network_table;
 mod platform;
 mod registered_peers;
 mod setup;
+mod tun_task_channel;
 mod udp_listener;
 mod wg_tunnel;
 
@@ -20,28 +21,6 @@ use std::sync::Arc;
 #[cfg(target_os = "linux")]
 use platform::linux::tun_device;
 
-type TunTaskPayload = (u64, Vec<u8>);
-
-#[derive(Clone)]
-pub struct TunTaskTx(tokio::sync::mpsc::UnboundedSender<TunTaskPayload>);
-
-impl TunTaskTx {
-    fn send(
-        &self,
-        data: TunTaskPayload,
-    ) -> Result<(), tokio::sync::mpsc::error::SendError<TunTaskPayload>> {
-        self.0.send(data)
-    }
-}
-
-pub struct TunTaskRx(tokio::sync::mpsc::UnboundedReceiver<TunTaskPayload>);
-
-impl TunTaskRx {
-    async fn recv(&mut self) -> Option<TunTaskPayload> {
-        self.0.recv().await
-    }
-}
-
 /// Start wireguard UDP listener and TUN device
 ///
 /// # Errors
diff --git a/common/wireguard/src/platform/linux/tun_device.rs b/common/wireguard/src/platform/linux/tun_device.rs
index deea1937c1..c714eb6cdf 100644
--- a/common/wireguard/src/platform/linux/tun_device.rs
+++ b/common/wireguard/src/platform/linux/tun_device.rs
@@ -6,17 +6,14 @@ use std::{
 
 use etherparse::{InternetSlice, SlicedPacket};
 use tap::TapFallible;
-use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
-    sync::mpsc::{self},
-};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 use crate::{
     event::Event,
     setup::{TUN_BASE_NAME, TUN_DEVICE_ADDRESS, TUN_DEVICE_NETMASK},
+    tun_task_channel::{tun_task_channel, TunTaskPayload, TunTaskRx, TunTaskTx},
     udp_listener::PeersByIp,
     wg_tunnel::PeersByTag,
-    TunTaskPayload, TunTaskRx, TunTaskTx,
 };
 
 fn setup_tokio_tun_device(name: &str, address: Ipv4Addr, netmask: Ipv4Addr) -> tokio_tun::Tun {
@@ -38,7 +35,6 @@ pub struct TunDevice {
     tun: tokio_tun::Tun,
 
     // Incoming data that we should send
-    // tun_task_rx: mpsc::UnboundedReceiver<Vec<u8>>,
     tun_task_rx: TunTaskRx,
 
     // The routing table.
@@ -62,9 +58,7 @@ impl TunDevice {
         log::info!("Created TUN device: {}", tun.name());
 
         // Channels to communicate with the other tasks
-        let (tun_task_tx, tun_task_rx) = mpsc::unbounded_channel();
-        let tun_task_tx = TunTaskTx(tun_task_tx);
-        let tun_task_rx = TunTaskRx(tun_task_rx);
+        let (tun_task_tx, tun_task_rx) = tun_task_channel();
 
         let tun_device = TunDevice {
             tun_task_rx,
diff --git a/common/wireguard/src/tun_task_channel.rs b/common/wireguard/src/tun_task_channel.rs
new file mode 100644
index 0000000000..423243d27e
--- /dev/null
+++ b/common/wireguard/src/tun_task_channel.rs
@@ -0,0 +1,26 @@
+pub(crate) type TunTaskPayload = (u64, Vec<u8>);
+
+#[derive(Clone)]
+pub struct TunTaskTx(tokio::sync::mpsc::UnboundedSender<TunTaskPayload>);
+
+pub(crate) struct TunTaskRx(tokio::sync::mpsc::UnboundedReceiver<TunTaskPayload>);
+
+impl TunTaskTx {
+    pub(crate) fn send(
+        &self,
+        data: TunTaskPayload,
+    ) -> Result<(), tokio::sync::mpsc::error::SendError<TunTaskPayload>> {
+        self.0.send(data)
+    }
+}
+
+impl TunTaskRx {
+    pub(crate) async fn recv(&mut self) -> Option<TunTaskPayload> {
+        self.0.recv().await
+    }
+}
+
+pub(crate) fn tun_task_channel() -> (TunTaskTx, TunTaskRx) {
+    let (tun_task_tx, tun_task_rx) = tokio::sync::mpsc::unbounded_channel();
+    (TunTaskTx(tun_task_tx), TunTaskRx(tun_task_rx))
+}
diff --git a/common/wireguard/src/udp_listener.rs b/common/wireguard/src/udp_listener.rs
index fa68962bf1..e06e0feb20 100644
--- a/common/wireguard/src/udp_listener.rs
+++ b/common/wireguard/src/udp_listener.rs
@@ -24,8 +24,8 @@ use crate::{
     network_table::NetworkTable,
     registered_peers::{RegisteredPeer, RegisteredPeers},
     setup::{self, WG_ADDRESS, WG_PORT},
+    tun_task_channel::TunTaskTx,
     wg_tunnel::PeersByTag,
-    TunTaskTx,
 };
 
 const MAX_PACKET: usize = 65535;
diff --git a/common/wireguard/src/wg_tunnel.rs b/common/wireguard/src/wg_tunnel.rs
index cfc10d5057..f9308fbe2c 100644
--- a/common/wireguard/src/wg_tunnel.rs
+++ b/common/wireguard/src/wg_tunnel.rs
@@ -15,7 +15,8 @@ use tokio::{
 };
 
 use crate::{
-    error::WgError, event::Event, network_table::NetworkTable, registered_peers::PeerIdx, TunTaskTx,
+    error::WgError, event::Event, network_table::NetworkTable, registered_peers::PeerIdx,
+    tun_task_channel::TunTaskTx,
 };
 
 const HANDSHAKE_MAX_RATE: u64 = 10;