diff --git a/common/wireguard/src/event.rs b/common/wireguard/src/event.rs index 7d425cd5d6..9b62bf2b4e 100644 --- a/common/wireguard/src/event.rs +++ b/common/wireguard/src/event.rs @@ -3,23 +3,29 @@ use std::fmt::{Display, Formatter}; use bytes::Bytes; #[allow(unused)] -#[derive(Debug, Clone)] +#[derive(Debug)] pub enum Event { /// IP packet received from the WireGuard tunnel that should be passed through to the corresponding virtual device/internet. /// Original implementation also has protocol here since it understands it, but we'll have to infer it downstream - WgPacket(Bytes), + Wg(Bytes), + /// IP packet received from the UDP listener that was verified as part of the handshake + WgVerified(Bytes), /// IP packet to be sent through the WireGuard tunnel as crafted by the virtual device. - IpPacket(Bytes), + Ip(Bytes), } impl Display for Event { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { - Event::WgPacket(data) => { + Event::Wg(data) => { let size = data.len(); write!(f, "WgPacket{{ size={size} }}") } - Event::IpPacket(data) => { + Event::WgVerified(data) => { + let size = data.len(); + write!(f, "WgVerifiedPacket{{ size={size} }}") + } + Event::Ip(data) => { let size = data.len(); write!(f, "IpPacket{{ size={size} }}") } diff --git a/common/wireguard/src/platform/linux/tun_device.rs b/common/wireguard/src/platform/linux/tun_device.rs index 0e0b71da28..96b8ffeaa1 100644 --- a/common/wireguard/src/platform/linux/tun_device.rs +++ b/common/wireguard/src/platform/linux/tun_device.rs @@ -63,7 +63,7 @@ pub(crate) fn start_tun_device(peers_by_ip: Arc>) -> if let Some(peer_tx) = peers_by_ip.lock().unwrap().longest_match(dst_addr).map(|(_, tx)| tx) { log::info!("Forward packet to wg tunnel"); peer_tx - .send(Event::IpPacket(packet.to_vec().into())) + .send(Event::Ip(packet.to_vec().into())) .tap_err(|err| log::error!("{err}")) .unwrap(); } else { diff --git a/common/wireguard/src/udp_listener.rs b/common/wireguard/src/udp_listener.rs index 8435955670..7539b343a7 100644 --- a/common/wireguard/src/udp_listener.rs +++ b/common/wireguard/src/udp_listener.rs @@ -1,7 +1,12 @@ -use std::{net::SocketAddr, sync::Arc}; +use std::{collections::HashMap, net::SocketAddr, sync::Arc}; +use boringtun::{ + noise::{handshake::parse_handshake_anon, rate_limiter::RateLimiter, TunnResult}, + x25519, +}; use dashmap::DashMap; use futures::StreamExt; +use ip_network::IpNetwork; use log::error; use nym_task::TaskClient; use tap::TapFallible; @@ -20,9 +25,16 @@ use crate::{ const MAX_PACKET: usize = 65535; pub(crate) type PeerIdx = u32; -pub(crate) type ActivePeers = DashMap>; +pub(crate) type ActivePeers = DashMap>; pub(crate) type PeersByIp = NetworkTable>; +#[derive(Debug)] +struct RegisteredPeer { + public_key: x25519::PublicKey, + allowed_ips: IpNetwork, + // endpoint: SocketAddr, +} + pub(crate) async fn start_udp_listener( tun_task_tx: TunTaskTx, active_peers: Arc, @@ -33,19 +45,37 @@ pub(crate) async fn start_udp_listener( log::info!("Starting wireguard UDP listener on {wg_address}"); let udp_socket = Arc::new(UdpSocket::bind(wg_address).await?); - // Setup static key for development + // Setup our own keys let static_private = setup::server_static_private_key(); + let static_public = x25519::PublicKey::from(&static_private); + let handshake_max_rate = 100u64; + let rate_limiter = RateLimiter::new(&static_public, handshake_max_rate); - // A single hardcoded peer + // Test peer let peer_static_public = setup::peer_static_public_key(); let peer_allowed_ips = setup::peer_allowed_ips(); let peer_index = 0; + let test_peer = Arc::new(tokio::sync::Mutex::new(RegisteredPeer { + public_key: peer_static_public, + allowed_ips: peer_allowed_ips, + })); + + type PeerIdx = u32; + let mut registered_peers: HashMap>> = + HashMap::new(); + let mut registered_peers_by_idx: HashMap>> = + HashMap::new(); + + registered_peers.insert(peer_static_public, Arc::clone(&test_peer)); + registered_peers_by_idx.insert(0, test_peer); tokio::spawn(async move { // Each tunnel is run in its own task, and the task handle is stored here so we can remove // it from `active_peers` when the tunnel is closed let mut active_peers_task_handles = futures::stream::FuturesUnordered::new(); + let mut buf = [0u8; MAX_PACKET]; + let mut dst_buf = [0u8; MAX_PACKET]; while !task_client.is_shutdown() { tokio::select! { @@ -54,11 +84,11 @@ pub(crate) async fn start_udp_listener( break; } // Handle tunnel closing - Some(addr) = active_peers_task_handles.next() => { - match addr { - Ok(addr) => { - log::info!("Removing peer: {addr:?}"); - active_peers.remove(&addr); + Some(public_key) = active_peers_task_handles.next() => { + match public_key { + Ok(public_key) => { + log::info!("Removing peer: {public_key:?}"); + active_peers.remove(&public_key); // TODO: remove from peers_by_ip } Err(err) => { @@ -70,38 +100,78 @@ pub(crate) async fn start_udp_listener( Ok((len, addr)) = udp_socket.recv_from(&mut buf) => { log::trace!("udp: received {} bytes from {}", len, addr); - if let Some(peer_tx) = active_peers.get_mut(&addr) { + // Verify the incoming packet + let verified_packet = match rate_limiter.verify_packet(Some(addr.ip()), &buf[..len], &mut dst_buf) { + Ok(packet) => packet, + Err(TunnResult::WriteToNetwork(cookie)) => { + log::info!("WireGuard UDP listener: send back cookie"); + udp_socket.send_to(cookie, addr).await.unwrap(); + continue; + } + Err(err) => { + log::warn!("{err:?}"); + continue; + } + }; + + // Check if this is a registered peer, if not, just skip + let registered_peer = match verified_packet { + boringtun::noise::Packet::HandshakeInit(ref packet) => { + let Ok(handshake) = parse_handshake_anon(&static_private, &static_public, packet) else { + log::warn!("Handshake failed"); + continue; + }; + registered_peers.get(&x25519::PublicKey::from(handshake.peer_static_public)) + }, + boringtun::noise::Packet::HandshakeResponse(packet) => { + let peer_idx = packet.receiver_idx >> 8; + registered_peers_by_idx.get(&peer_idx) + }, + boringtun::noise::Packet::PacketCookieReply(packet) => { + let peer_idx = packet.receiver_idx >> 8; + registered_peers_by_idx.get(&peer_idx) + }, + boringtun::noise::Packet::PacketData(packet) => { + let peer_idx = packet.receiver_idx >> 8; + registered_peers_by_idx.get(&peer_idx) + }, + }; + + let registered_peer = if let Some(registered_peer) = registered_peer { + registered_peer.lock().await + } else { + log::warn!("Peer not registered"); + continue; + }; + + // Look up if the peer is already connected + if let Some(peer_tx) = active_peers.get_mut(®istered_peer.public_key) { + // If it is, send it the packet to deal with log::info!("udp: received {len} bytes from {addr} from known peer"); - peer_tx.send(Event::WgPacket(buf[..len].to_vec().into())) + peer_tx.send(Event::WgVerified(buf[..len].to_vec().into())) .tap_err(|err| log::error!("{err}")) .unwrap(); } else { + // If it isn't, start a new tunnel log::info!("udp: received {len} bytes from {addr} from unknown peer, starting tunnel"); - // TODO: this is a temporary solution for development since this - // assumes we know the peer_static_public this corresponds to. - // TODO: rework this before production! This is likely not secure! - log::warn!("Assuming peer_static_public is known"); - log::warn!("SECURITY: Rework me to do proper handshake before creating the tunnel!"); let (join_handle, peer_tx) = crate::wg_tunnel::start_wg_tunnel( addr, udp_socket.clone(), static_private.clone(), - peer_static_public, - peer_allowed_ips, + registered_peer.public_key, + registered_peer.allowed_ips, peer_index, tun_task_tx.clone(), ); - peers_by_ip.lock().unwrap().insert(peer_allowed_ips, peer_tx.clone()); + peers_by_ip.lock().unwrap().insert(registered_peer.allowed_ips, peer_tx.clone()); - peer_tx.send(Event::WgPacket(buf[..len].to_vec().into())) + peer_tx.send(Event::Wg(buf[..len].to_vec().into())) .tap_err(|err| log::error!("{err}")) .unwrap(); - // WIP(JON): active peers should probably be keyed by peer_static_public - // instead. Does this current setup lead to any issues? log::info!("Adding peer: {addr}"); - active_peers.insert(addr, peer_tx); + active_peers.insert(registered_peer.public_key, peer_tx); active_peers_task_handles.push(join_handle); } }, diff --git a/common/wireguard/src/wg_tunnel.rs b/common/wireguard/src/wg_tunnel.rs index 2d44a03fb1..11f8485cf3 100644 --- a/common/wireguard/src/wg_tunnel.rs +++ b/common/wireguard/src/wg_tunnel.rs @@ -61,6 +61,7 @@ impl WireGuardTunnel { peer_static_public: x25519::PublicKey, peer_allowed_ips: ip_network::IpNetwork, index: PeerIdx, + // rate_limiter: Option, tunnel_tx: TunTaskTx, ) -> (Self, mpsc::UnboundedSender) { let local_addr = udp.local_addr().unwrap(); @@ -126,12 +127,17 @@ impl WireGuardTunnel { Some(packet) => { info!("event loop: {packet}"); match packet { - Event::WgPacket(data) => { + Event::Wg(data) => { let _ = self.consume_wg(&data) .await .tap_err(|err| error!("WireGuard tunnel: consume_wg error: {err}")); }, - Event::IpPacket(data) => self.consume_eth(&data).await, + Event::WgVerified(data) => { + let _ = self.consume_verified_wg(&data) + .await + .tap_err(|err| error!("WireGuard tunnel: consume_verified_wg error: {err}")); + } + Event::Ip(data) => self.consume_eth(&data).await, } }, None => { @@ -191,8 +197,6 @@ impl WireGuardTunnel { } } TunnResult::WriteToTunnelV4(packet, addr) => { - // TODO: once the flow is redone, we should add updating the endpoint dynamically - // self.set_endpoint(addr); if self.allowed_ips.longest_match(addr).is_some() { self.tun_task_tx.send(packet.to_vec()).unwrap(); } else { @@ -200,8 +204,6 @@ impl WireGuardTunnel { } } TunnResult::WriteToTunnelV6(packet, addr) => { - // TODO: once the flow is redone, we should add updating the endpoint dynamically - // self.set_endpoint(addr); if self.allowed_ips.longest_match(addr).is_some() { self.tun_task_tx.send(packet.to_vec()).unwrap(); } else { @@ -218,6 +220,13 @@ impl WireGuardTunnel { Ok(()) } + async fn consume_verified_wg(&mut self, data: &[u8]) -> Result<(), WgError> { + // Potentially we could take some shortcuts here in the name of performance, but currently + // I don't see that the needed functions in boringtun is exposed in the public API. + // TODO: make sure we don't put double pressure on the rate limiter! + self.consume_wg(data).await + } + async fn consume_eth(&self, data: &Bytes) { info!("consume_eth: raw packet size: {}", data.len()); let encapsulated_packet = self.encapsulate_packet(data).await; @@ -300,7 +309,7 @@ pub(crate) fn start_wg_tunnel( peer_index: PeerIdx, tunnel_tx: TunTaskTx, ) -> ( - tokio::task::JoinHandle, + tokio::task::JoinHandle, mpsc::UnboundedSender, ) { let (mut tunnel, peer_tx) = WireGuardTunnel::new( @@ -314,7 +323,7 @@ pub(crate) fn start_wg_tunnel( ); let join_handle = tokio::spawn(async move { tunnel.spin_off().await; - endpoint + peer_static_public }); (join_handle, peer_tx) }