From 7cafd25036df945d2bc6bf14022b4074ff01f3c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Tue, 24 Oct 2023 11:22:59 +0200 Subject: [PATCH] wg: all channels strongly typed (#4035) * Add PeerEventSender/Receiver * Create strong types * Create PacketRelaySender/Receiver * Strongly typed --- common/wireguard/src/active_peers.rs | 32 ++++++++++++++++++++---- common/wireguard/src/packet_relayer.rs | 22 ++++++++++++----- common/wireguard/src/udp_listener.rs | 18 +++++--------- common/wireguard/src/wg_tunnel.rs | 34 ++++++++++++++------------ 4 files changed, 68 insertions(+), 38 deletions(-) diff --git a/common/wireguard/src/active_peers.rs b/common/wireguard/src/active_peers.rs index 43b4b43366..6973d29be1 100644 --- a/common/wireguard/src/active_peers.rs +++ b/common/wireguard/src/active_peers.rs @@ -9,8 +9,30 @@ use tokio::sync::mpsc::{self}; use crate::event::Event; -pub(crate) type PeersByKey = DashMap>; -pub(crate) type PeersByAddr = DashMap>; +// Channels that are used to communicate with the various tunnels +#[derive(Clone)] +pub struct PeerEventSender(mpsc::UnboundedSender); +pub(crate) struct PeerEventReceiver(mpsc::UnboundedReceiver); + +impl PeerEventSender { + pub(crate) fn send(&self, event: Event) -> Result<(), mpsc::error::SendError> { + self.0.send(event) + } +} + +impl PeerEventReceiver { + pub(crate) async fn recv(&mut self) -> Option { + 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; +pub(crate) type PeersByAddr = DashMap; #[derive(Default)] pub(crate) struct ActivePeers { @@ -30,7 +52,7 @@ impl ActivePeers { &self, public_key: x25519::PublicKey, addr: SocketAddr, - peer_tx: mpsc::UnboundedSender, + 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>> { + ) -> Option> { self.active_peers.get_mut(public_key) } pub(crate) fn get_by_addr( &self, addr: &SocketAddr, - ) -> Option>> { + ) -> Option> { self.active_peers_by_addr.get(addr) } } diff --git a/common/wireguard/src/packet_relayer.rs b/common/wireguard/src/packet_relayer.rs index 8f0a4ce1ac..40d49d5d37 100644 --- a/common/wireguard/src/packet_relayer.rs +++ b/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)>); +pub(crate) struct PacketRelayReceiver(pub(crate) mpsc::Receiver<(u64, Vec)>); + +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)>, + 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>>>, + peers_by_tag: Arc>>, } impl PacketRelayer { pub(crate) fn new( tun_task_tx: TunTaskTx, tun_task_response_rx: TunTaskResponseRx, - peers_by_tag: Arc>>>, - ) -> (Self, mpsc::Sender<(u64, Vec)>) { - let (packet_tx, packet_rx) = mpsc::channel(16); + peers_by_tag: Arc>>, + ) -> (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(); }, diff --git a/common/wireguard/src/udp_listener.rs b/common/wireguard/src/udp_listener.rs index c703468d97..5c04c010a4 100644 --- a/common/wireguard/src/udp_listener.rs +++ b/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>; +pub(crate) type PeersByIp = NetworkTable; 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, // Send data to the TUN device for sending - // tun_task_tx: TunTaskTx, - packet_tx: mpsc::Sender<(u64, Vec)>, + 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)>, + packet_tx: PacketRelaySender, peers_by_ip: Arc>, peers_by_tag: Arc>, gateway_client_registry: Arc, diff --git a/common/wireguard/src/wg_tunnel.rs b/common/wireguard/src/wg_tunnel.rs index 32ece1e842..f35f51e1bd 100644 --- a/common/wireguard/src/wg_tunnel.rs +++ b/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>; +pub(crate) type PeersByTag = HashMap; pub struct WireGuardTunnel { // Incoming data from the UDP socket received in the main event loop - peer_rx: mpsc::UnboundedReceiver, + peer_rx: PeerEventReceiver, // UDP socket used for sending data udp: Arc, @@ -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)>, + packet_tx: PacketRelaySender, tag: u64, } @@ -67,8 +69,8 @@ impl WireGuardTunnel { index: PeerIdx, peer_allowed_ips: ip_network::IpNetwork, // rate_limiter: Option, - packet_tx: mpsc::Sender<(u64, Vec)>, - ) -> (Self, mpsc::UnboundedSender, 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)>, + packet_tx: PacketRelaySender, ) -> ( tokio::task::JoinHandle, - mpsc::UnboundedSender, + PeerEventSender, u64, ) { let (mut tunnel, peer_tx, tag) = WireGuardTunnel::new(