wg: all channels strongly typed (#4035)

* Add PeerEventSender/Receiver

* Create strong types

* Create PacketRelaySender/Receiver

* Strongly typed

common/wireguard/src/active_peers.rs

@@ -9,8 +9,30 @@ use tokio::sync::mpsc::{self};
 
 use crate::event::Event;
 
-pub(crate) type PeersByKey = DashMap<x25519::PublicKey, mpsc::UnboundedSender<Event>>;
-pub(crate) type PeersByAddr = DashMap<SocketAddr, mpsc::UnboundedSender<Event>>;
+// Channels that are used to communicate with the various tunnels
+#[derive(Clone)]
+pub struct PeerEventSender(mpsc::UnboundedSender<Event>);
+pub(crate) struct PeerEventReceiver(mpsc::UnboundedReceiver<Event>);
+
+impl PeerEventSender {
+    pub(crate) fn send(&self, event: Event) -> Result<(), mpsc::error::SendError<Event>> {
+        self.0.send(event)
+    }
+}
+
+impl PeerEventReceiver {
+    pub(crate) async fn recv(&mut self) -> Option<Event> {
+        self.0.recv().await
+    }
+}
+
+pub(crate) fn peer_event_channel() -> (PeerEventSender, PeerEventReceiver) {
+    let (tx, rx) = mpsc::unbounded_channel();
+    (PeerEventSender(tx), PeerEventReceiver(rx))
+}
+
+pub(crate) type PeersByKey = DashMap<x25519::PublicKey, PeerEventSender>;
+pub(crate) type PeersByAddr = DashMap<SocketAddr, PeerEventSender>;
 
 #[derive(Default)]
 pub(crate) struct ActivePeers {
@@ -30,7 +52,7 @@ impl ActivePeers {
         &self,
         public_key: x25519::PublicKey,
         addr: SocketAddr,
-        peer_tx: mpsc::UnboundedSender<Event>,
+        peer_tx: PeerEventSender,
     ) {
         self.active_peers.insert(public_key, peer_tx.clone());
         self.active_peers_by_addr.insert(addr, peer_tx);
@@ -39,14 +61,14 @@ impl ActivePeers {
     pub(crate) fn get_by_key_mut(
         &self,
         public_key: &x25519::PublicKey,
-    ) -> Option<RefMut<'_, x25519::PublicKey, mpsc::UnboundedSender<Event>>> {
+    ) -> Option<RefMut<'_, x25519::PublicKey, PeerEventSender>> {
         self.active_peers.get_mut(public_key)
     }
 
     pub(crate) fn get_by_addr(
         &self,
         addr: &SocketAddr,
-    ) -> Option<Ref<'_, SocketAddr, mpsc::UnboundedSender<Event>>> {
+    ) -> Option<Ref<'_, SocketAddr, PeerEventSender>> {
         self.active_peers_by_addr.get(addr)
     }
 }

common/wireguard/src/packet_relayer.rs

@@ -4,16 +4,26 @@ use tap::TapFallible;
 use tokio::sync::mpsc::{self};
 
 use crate::{
+    active_peers::PeerEventSender,
     event::Event,
     tun_task_channel::{TunTaskResponseRx, TunTaskTx},
 };
 
+#[derive(Clone)]
+pub struct PacketRelaySender(pub(crate) mpsc::Sender<(u64, Vec<u8>)>);
+pub(crate) struct PacketRelayReceiver(pub(crate) mpsc::Receiver<(u64, Vec<u8>)>);
+
+pub(crate) fn packet_relay_channel() -> (PacketRelaySender, PacketRelayReceiver) {
+    let (tx, rx) = mpsc::channel(16);
+    (PacketRelaySender(tx), PacketRelayReceiver(rx))
+}
+
 // The tunnels send packets to the packet relayer, which then relays it to the tun device. And
 // conversely, it's where the tun device send responses to, which are relayed back to the correct
 // tunnel.
 pub(crate) struct PacketRelayer {
     // Receive packets from the various tunnels
-    packet_rx: mpsc::Receiver<(u64, Vec<u8>)>,
+    packet_rx: PacketRelayReceiver,
 
     // After receive from tunnels, send to the tun device
     tun_task_tx: TunTaskTx,
@@ -22,16 +32,16 @@ pub(crate) struct PacketRelayer {
     tun_task_response_rx: TunTaskResponseRx,
 
     // After receiving from the tun device, relay back to the correct tunnel
-    peers_by_tag: Arc<std::sync::Mutex<HashMap<u64, mpsc::UnboundedSender<Event>>>>,
+    peers_by_tag: Arc<std::sync::Mutex<HashMap<u64, PeerEventSender>>>,
 }
 
 impl PacketRelayer {
     pub(crate) fn new(
         tun_task_tx: TunTaskTx,
         tun_task_response_rx: TunTaskResponseRx,
-        peers_by_tag: Arc<std::sync::Mutex<HashMap<u64, mpsc::UnboundedSender<Event>>>>,
-    ) -> (Self, mpsc::Sender<(u64, Vec<u8>)>) {
-        let (packet_tx, packet_rx) = mpsc::channel(16);
+        peers_by_tag: Arc<std::sync::Mutex<HashMap<u64, PeerEventSender>>>,
+    ) -> (Self, PacketRelaySender) {
+        let (packet_tx, packet_rx) = packet_relay_channel();
         (
             Self {
                 packet_rx,
@@ -46,7 +56,7 @@ impl PacketRelayer {
     pub(crate) async fn run(mut self) {
         loop {
             tokio::select! {
-                Some((tag, packet)) = self.packet_rx.recv() => {
+                Some((tag, packet)) = self.packet_rx.0.recv() => {
                     log::info!("Sent packet to tun device with tag: {tag}");
                     self.tun_task_tx.send((tag, packet)).unwrap();
                 },

common/wireguard/src/udp_listener.rs

@@ -9,19 +9,14 @@ use log::error;
 use nym_task::TaskClient;
 use nym_wireguard_types::{registration::GatewayClientRegistry, PeerPublicKey};
 use tap::TapFallible;
-use tokio::{
-    net::UdpSocket,
-    sync::{
-        mpsc::{self},
-        Mutex,
-    },
-};
+use tokio::{net::UdpSocket, sync::Mutex};
 
 use crate::{
-    active_peers::ActivePeers,
+    active_peers::{ActivePeers, PeerEventSender},
     error::WgError,
     event::Event,
     network_table::NetworkTable,
+    packet_relayer::PacketRelaySender,
     registered_peers::{RegisteredPeer, RegisteredPeers},
     setup::{self, WG_ADDRESS, WG_PORT},
     wg_tunnel::PeersByTag,
@@ -30,7 +25,7 @@ use crate::{
 const MAX_PACKET: usize = 65535;
 
 // Registered peers
-pub(crate) type PeersByIp = NetworkTable<mpsc::UnboundedSender<Event>>;
+pub(crate) type PeersByIp = NetworkTable<PeerEventSender>;
 
 async fn add_test_peer(registered_peers: &mut RegisteredPeers) {
     let peer_static_public = PeerPublicKey::new(setup::peer_static_public_key());
@@ -64,8 +59,7 @@ pub struct WgUdpListener {
     udp: Arc<UdpSocket>,
 
     // Send data to the TUN device for sending
-    // tun_task_tx: TunTaskTx,
-    packet_tx: mpsc::Sender<(u64, Vec<u8>)>,
+    packet_tx: PacketRelaySender,
 
     // Wireguard rate limiter
     rate_limiter: RateLimiter,
@@ -75,7 +69,7 @@ pub struct WgUdpListener {
 
 impl WgUdpListener {
     pub async fn new(
-        packet_tx: mpsc::Sender<(u64, Vec<u8>)>,
+        packet_tx: PacketRelaySender,
         peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>,
         peers_by_tag: Arc<std::sync::Mutex<PeersByTag>>,
         gateway_client_registry: Arc<GatewayClientRegistry>,

common/wireguard/src/wg_tunnel.rs

@@ -9,24 +9,27 @@ use bytes::Bytes;
 use log::{debug, error, info, warn};
 use rand::RngCore;
 use tap::TapFallible;
-use tokio::{
-    net::UdpSocket,
-    sync::{broadcast, mpsc},
-    time::timeout,
-};
+use tokio::{net::UdpSocket, sync::broadcast, time::timeout};
 
-use crate::{error::WgError, event::Event, network_table::NetworkTable, registered_peers::PeerIdx};
+use crate::{
+    active_peers::{peer_event_channel, PeerEventReceiver, PeerEventSender},
+    error::WgError,
+    event::Event,
+    network_table::NetworkTable,
+    packet_relayer::PacketRelaySender,
+    registered_peers::PeerIdx,
+};
 
 const HANDSHAKE_MAX_RATE: u64 = 10;
 
 const MAX_PACKET: usize = 65535;
 
 // We index the tunnels by tag
-pub(crate) type PeersByTag = HashMap<u64, mpsc::UnboundedSender<Event>>;
+pub(crate) type PeersByTag = HashMap<u64, PeerEventSender>;
 
 pub struct WireGuardTunnel {
     // Incoming data from the UDP socket received in the main event loop
-    peer_rx: mpsc::UnboundedReceiver<Event>,
+    peer_rx: PeerEventReceiver,
 
     // UDP socket used for sending data
     udp: Arc<UdpSocket>,
@@ -45,8 +48,7 @@ pub struct WireGuardTunnel {
     close_rx: broadcast::Receiver<()>,
 
     // Send data to the task that handles sending data through the tun device
-    // tun_task_tx: TunTaskTx,
-    packet_tx: mpsc::Sender<(u64, Vec<u8>)>,
+    packet_tx: PacketRelaySender,
 
     tag: u64,
 }
@@ -67,8 +69,8 @@ impl WireGuardTunnel {
         index: PeerIdx,
         peer_allowed_ips: ip_network::IpNetwork,
         // rate_limiter: Option<RateLimiter>,
-        packet_tx: mpsc::Sender<(u64, Vec<u8>)>,
-    ) -> (Self, mpsc::UnboundedSender<Event>, u64) {
+        packet_tx: PacketRelaySender,
+    ) -> (Self, PeerEventSender, u64) {
         let local_addr = udp.local_addr().unwrap();
         let peer_addr = udp.peer_addr();
         log::info!("New wg tunnel: endpoint: {endpoint}, local_addr: {local_addr}, peer_addr: {peer_addr:?}");
@@ -95,7 +97,7 @@ impl WireGuardTunnel {
         ));
 
         // Channels with incoming data that is received by the main event loop
-        let (peer_tx, peer_rx) = mpsc::unbounded_channel();
+        let (peer_tx, peer_rx) = peer_event_channel();
 
         // Signal close tunnel
         let (close_tx, close_rx) = broadcast::channel(1);
@@ -212,6 +214,7 @@ impl WireGuardTunnel {
             TunnResult::WriteToTunnelV4(packet, addr) => {
                 if self.allowed_ips.longest_match(addr).is_some() {
                     self.packet_tx
+                        .0
                         .send((self.tag, packet.to_vec()))
                         .await
                         .unwrap();
@@ -222,6 +225,7 @@ impl WireGuardTunnel {
             TunnResult::WriteToTunnelV6(packet, addr) => {
                 if self.allowed_ips.longest_match(addr).is_some() {
                     self.packet_tx
+                        .0
                         .send((self.tag, packet.to_vec()))
                         .await
                         .unwrap();
@@ -326,10 +330,10 @@ pub(crate) fn start_wg_tunnel(
     peer_static_public: x25519::PublicKey,
     peer_index: PeerIdx,
     peer_allowed_ips: ip_network::IpNetwork,
-    packet_tx: mpsc::Sender<(u64, Vec<u8>)>,
+    packet_tx: PacketRelaySender,
 ) -> (
     tokio::task::JoinHandle<x25519::PublicKey>,
-    mpsc::UnboundedSender<Event>,
+    PeerEventSender,
     u64,
 ) {
     let (mut tunnel, peer_tx, tag) = WireGuardTunnel::new(