diff --git a/Cargo.lock b/Cargo.lock index 24f0e34659..5c80148210 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7541,6 +7541,7 @@ dependencies = [ "log", "nym-task", "nym-wireguard-types", + "rand 0.8.5", "serde", "tap", "thiserror", diff --git a/common/wireguard/Cargo.toml b/common/wireguard/Cargo.toml index 25cb796cc4..bddacbca1f 100644 --- a/common/wireguard/Cargo.toml +++ b/common/wireguard/Cargo.toml @@ -26,11 +26,12 @@ ip_network = "0.4.1" ip_network_table = "0.2.0" log.workspace = true nym-task = { path = "../task" } +nym-wireguard-types = { path = "../wireguard-types" } +rand.workspace = true +serde = { workspace = true, features = ["derive"] } tap.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["rt-multi-thread", "net", "io-util"] } -serde = { workspace = true, features = ["derive"] } -nym-wireguard-types = { path = "../wireguard-types" } [target.'cfg(target_os = "linux")'.dependencies] tokio-tun = "0.9.0" diff --git a/common/wireguard/src/lib.rs b/common/wireguard/src/lib.rs index 5a4520d6c2..0df9e8ce1a 100644 --- a/common/wireguard/src/lib.rs +++ b/common/wireguard/src/lib.rs @@ -20,12 +20,25 @@ use std::sync::Arc; #[cfg(target_os = "linux")] use platform::linux::tun_device; +type TunTaskPayload = (u64, Vec); + #[derive(Clone)] -pub struct TunTaskTx(tokio::sync::mpsc::UnboundedSender>); +pub struct TunTaskTx(tokio::sync::mpsc::UnboundedSender); impl TunTaskTx { - fn send(&self, packet: Vec) -> Result<(), tokio::sync::mpsc::error::SendError>> { - self.0.send(packet) + fn send( + &self, + data: TunTaskPayload, + ) -> Result<(), tokio::sync::mpsc::error::SendError> { + self.0.send(data) + } +} + +pub struct TunTaskRx(tokio::sync::mpsc::UnboundedReceiver); + +impl TunTaskRx { + async fn recv(&mut self) -> Option { + self.0.recv().await } } @@ -39,16 +52,21 @@ pub async fn start_wireguard( task_client: nym_task::TaskClient, gateway_client_registry: Arc, ) -> Result<(), Box> { + // We can either index peers by their IP like standard wireguard let peers_by_ip = Arc::new(std::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(std::sync::Mutex::new(wg_tunnel::PeersByTag::new())); + // Start the tun device that is used to relay traffic outbound - let (tun, tun_task_tx) = tun_device::TunDevice::new(peers_by_ip.clone()); + let (tun, tun_task_tx) = tun_device::TunDevice::new(peers_by_ip.clone(), peers_by_tag.clone()); tun.start(); // Start the UDP listener that clients connect to let udp_listener = udp_listener::WgUdpListener::new( tun_task_tx, peers_by_ip, + peers_by_tag, Arc::clone(&gateway_client_registry), ) .await?; diff --git a/common/wireguard/src/platform/linux/tun_device.rs b/common/wireguard/src/platform/linux/tun_device.rs index 5486e20f03..deea1937c1 100644 --- a/common/wireguard/src/platform/linux/tun_device.rs +++ b/common/wireguard/src/platform/linux/tun_device.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, net::{IpAddr, Ipv4Addr}, sync::Arc, }; @@ -14,7 +15,8 @@ use crate::{ event::Event, setup::{TUN_BASE_NAME, TUN_DEVICE_ADDRESS, TUN_DEVICE_NETMASK}, udp_listener::PeersByIp, - TunTaskTx, + wg_tunnel::PeersByTag, + TunTaskPayload, TunTaskRx, TunTaskTx, }; fn setup_tokio_tun_device(name: &str, address: Ipv4Addr, netmask: Ipv4Addr) -> tokio_tun::Tun { @@ -36,15 +38,22 @@ pub struct TunDevice { tun: tokio_tun::Tun, // Incoming data that we should send - tun_task_rx: mpsc::UnboundedReceiver>, + // tun_task_rx: mpsc::UnboundedReceiver>, + tun_task_rx: TunTaskRx, // The routing table. // An alternative would be to do NAT by just matching incoming with outgoing. peers_by_ip: Arc>, + + nat_table: HashMap, + peers_by_tag: Arc>, } impl TunDevice { - pub fn new(peers_by_ip: Arc>) -> (Self, TunTaskTx) { + pub fn new( + peers_by_ip: Arc>, + peers_by_tag: Arc>, + ) -> (Self, TunTaskTx) { let tun = setup_tokio_tun_device( format!("{TUN_BASE_NAME}%d").as_str(), TUN_DEVICE_ADDRESS.parse().unwrap(), @@ -53,13 +62,16 @@ 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, tun_task_rx) = mpsc::unbounded_channel(); let tun_task_tx = TunTaskTx(tun_task_tx); + let tun_task_rx = TunTaskRx(tun_task_rx); let tun_device = TunDevice { tun_task_rx, tun, peers_by_ip, + nat_table: HashMap::new(), + peers_by_tag, }; (tun_device, tun_task_tx) @@ -80,36 +92,56 @@ impl TunDevice { ); // Route packet to the correct peer. - let Ok(peers) = self.peers_by_ip.lock() else { - log::error!("Failed to lock peers_by_ip, aborting tun device read"); - return; - }; - if let Some(peer_tx) = peers.longest_match(dst_addr).map(|(_, tx)| tx) { - log::info!("Forward packet to wg tunnel"); - peer_tx - .send(Event::Ip(packet.to_vec().into())) - .tap_err(|err| log::error!("{err}")) - .ok(); - } else { - log::info!("No peer found, packet dropped"); + if false { + let Ok(peers) = self.peers_by_ip.lock() else { + log::error!("Failed to lock peers_by_ip, aborting tun device read"); + return; + }; + if let Some(peer_tx) = peers.longest_match(dst_addr).map(|(_, tx)| tx) { + log::info!("Forward packet to wg tunnel"); + peer_tx + .send(Event::Ip(packet.to_vec().into())) + .tap_err(|err| log::error!("{err}")) + .ok(); + return; + } } + + { + if let Some(tag) = self.nat_table.get(&dst_addr) { + log::info!("Forward packet to wg tunnel with tag: {tag}"); + self.peers_by_tag.lock().unwrap().get(tag).and_then(|tx| { + tx.send(Event::Ip(packet.to_vec().into())) + .tap_err(|err| log::error!("{err}")) + .ok() + }); + return; + } + } + + log::info!("No peer found, packet dropped"); } - async fn handle_tun_write(&mut self, data: Vec) { - let Some(dst_addr) = boringtun::noise::Tunn::dst_address(&data) else { + async fn handle_tun_write(&mut self, data: TunTaskPayload) { + let (tag, packet) = data; + let Some(dst_addr) = boringtun::noise::Tunn::dst_address(&packet) else { log::error!("Unable to parse dst_address in packet that was supposed to be written to tun device"); return; }; - let Some(src_addr) = parse_src_address(&data) else { + let Some(src_addr) = parse_src_address(&packet) else { log::error!("Unable to parse src_address in packet that was supposed to be written to tun device"); return; }; log::info!( "iface: write Packet({src_addr} -> {dst_addr}, {} bytes)", - data.len() + packet.len() ); + + // TODO: expire old entries + self.nat_table.insert(src_addr, tag); + self.tun - .write_all(&data) + .write_all(&packet) .await .tap_err(|err| { log::error!("iface: write error: {err}"); diff --git a/common/wireguard/src/udp_listener.rs b/common/wireguard/src/udp_listener.rs index 045dd7e892..fa68962bf1 100644 --- a/common/wireguard/src/udp_listener.rs +++ b/common/wireguard/src/udp_listener.rs @@ -24,6 +24,7 @@ use crate::{ network_table::NetworkTable, registered_peers::{RegisteredPeer, RegisteredPeers}, setup::{self, WG_ADDRESS, WG_PORT}, + wg_tunnel::PeersByTag, TunTaskTx, }; @@ -57,6 +58,9 @@ pub struct WgUdpListener { // The routing table, as defined by wireguard peers_by_ip: Arc>, + // ... or alternatively we can map peers by their tag + peers_by_tag: Arc>, + // The UDP socket to the peer udp: Arc, @@ -73,6 +77,7 @@ impl WgUdpListener { pub async fn new( tun_task_tx: TunTaskTx, peers_by_ip: Arc>, + peers_by_tag: Arc>, gateway_client_registry: Arc, ) -> Result> { let wg_address = SocketAddr::new(WG_ADDRESS.parse().unwrap(), WG_PORT); @@ -94,6 +99,7 @@ impl WgUdpListener { static_public, registered_peers, peers_by_ip, + peers_by_tag, udp, tun_task_tx, rate_limiter, @@ -194,7 +200,7 @@ impl WgUdpListener { // NOTE: we are NOT passing in the existing rate_limiter. Re-visit this // choice later. log::warn!("Creating new rate limiter, consider re-using?"); - let (join_handle, peer_tx) = crate::wg_tunnel::start_wg_tunnel( + let (join_handle, peer_tx, tag) = crate::wg_tunnel::start_wg_tunnel( addr, self.udp.clone(), self.static_private.clone(), @@ -205,6 +211,7 @@ impl WgUdpListener { ); self.peers_by_ip.lock().unwrap().insert(registered_peer.allowed_ips, peer_tx.clone()); + self.peers_by_tag.lock().unwrap().insert(tag, peer_tx.clone()); peer_tx.send(Event::Wg(buf[..len].to_vec().into())) .tap_err(|e| log::error!("{e}")) diff --git a/common/wireguard/src/wg_tunnel.rs b/common/wireguard/src/wg_tunnel.rs index 5eb4d3647c..cfc10d5057 100644 --- a/common/wireguard/src/wg_tunnel.rs +++ b/common/wireguard/src/wg_tunnel.rs @@ -1,4 +1,4 @@ -use std::{net::SocketAddr, sync::Arc, time::Duration}; +use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; use async_recursion::async_recursion; use boringtun::{ @@ -22,6 +22,9 @@ const HANDSHAKE_MAX_RATE: u64 = 10; const MAX_PACKET: usize = 65535; +// We index the tunnels by tag +pub(crate) type PeersByTag = HashMap>; + pub struct WireGuardTunnel { // Incoming data from the UDP socket received in the main event loop peer_rx: mpsc::UnboundedReceiver, @@ -44,6 +47,8 @@ pub struct WireGuardTunnel { // Send data to the task that handles sending data through the tun device tun_task_tx: TunTaskTx, + + tag: u64, } impl Drop for WireGuardTunnel { @@ -63,7 +68,7 @@ impl WireGuardTunnel { peer_allowed_ips: ip_network::IpNetwork, // rate_limiter: Option, tunnel_tx: TunTaskTx, - ) -> (Self, mpsc::UnboundedSender) { + ) -> (Self, mpsc::UnboundedSender, 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:?}"); @@ -98,6 +103,11 @@ impl WireGuardTunnel { let mut allowed_ips = NetworkTable::new(); allowed_ips.insert(peer_allowed_ips, ()); + // random u64 + use rand::RngCore; + let mut rng = rand::rngs::OsRng; + let tag = rng.next_u64(); + let tunnel = WireGuardTunnel { peer_rx, udp, @@ -107,9 +117,10 @@ impl WireGuardTunnel { close_tx, close_rx, tun_task_tx: tunnel_tx, + tag, }; - (tunnel, peer_tx) + (tunnel, peer_tx, tag) } fn close(&self) { @@ -198,14 +209,14 @@ impl WireGuardTunnel { } TunnResult::WriteToTunnelV4(packet, addr) => { if self.allowed_ips.longest_match(addr).is_some() { - self.tun_task_tx.send(packet.to_vec()).unwrap(); + self.tun_task_tx.send((self.tag, packet.to_vec())).unwrap(); } else { warn!("Packet from {addr} not in allowed_ips"); } } TunnResult::WriteToTunnelV6(packet, addr) => { if self.allowed_ips.longest_match(addr).is_some() { - self.tun_task_tx.send(packet.to_vec()).unwrap(); + self.tun_task_tx.send((self.tag, packet.to_vec())).unwrap(); } else { warn!("Packet (v6) from {addr} not in allowed_ips"); } @@ -311,8 +322,9 @@ pub(crate) fn start_wg_tunnel( ) -> ( tokio::task::JoinHandle, mpsc::UnboundedSender, + u64, ) { - let (mut tunnel, peer_tx) = WireGuardTunnel::new( + let (mut tunnel, peer_tx, tag) = WireGuardTunnel::new( udp, endpoint, static_private, @@ -325,5 +337,5 @@ pub(crate) fn start_wg_tunnel( tunnel.spin_off().await; peer_static_public }); - (join_handle, peer_tx) + (join_handle, peer_tx, tag) }