GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wireguard: parse the public key up front in the UDP handler (#3977)

Browse Source 
* wireguard: try to have a flow where we parse the public key up front

* Fix bug with continue instead of return in loop

* fix clippy::enum-variant-names
This commit is contained in:
Jon Häggblad
2023-10-16 09:44:10 +02:00
committed by GitHub
parent 82c107f7ad
commit 2c8187eb6c
4 changed files with 122 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/wireguard/src/event.rs
+11 -5
View File
@@ -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} }}")
}
common/wireguard/src/platform/linux/tun_device.rs
+1 -1
View File
@@ -63,7 +63,7 @@ pub(crate) fn start_tun_device(peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>) ->
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 {
common/wireguard/src/udp_listener.rs
+93 -23
View File
@@ -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<SocketAddr, mpsc::UnboundedSender<Event>>;
pub(crate) type ActivePeers = DashMap<x25519::PublicKey, mpsc::UnboundedSender<Event>>;
pub(crate) type PeersByIp = NetworkTable<mpsc::UnboundedSender<Event>>;
#[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<ActivePeers>,
@@ -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<x25519::PublicKey, Arc<tokio::sync::Mutex<RegisteredPeer>>> =
HashMap::new();
let mut registered_peers_by_idx: HashMap<PeerIdx, Arc<tokio::sync::Mutex<RegisteredPeer>>> =
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(&registered_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);
}
},
common/wireguard/src/wg_tunnel.rs
+17 -8
View File
@@ -61,6 +61,7 @@ impl WireGuardTunnel {
peer_static_public: x25519::PublicKey,
peer_allowed_ips: ip_network::IpNetwork,
index: PeerIdx,
// rate_limiter: Option<RateLimiter>,
tunnel_tx: TunTaskTx,
) -> (Self, mpsc::UnboundedSender<Event>) {
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<SocketAddr>,
tokio::task::JoinHandle<x25519::PublicKey>,
mpsc::UnboundedSender<Event>,
) {
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)
}