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

wireguard: create structs for udp handler and tun device (#4007)

Browse Source 
* Extract out parse_peer

* wip: handle_packet extract

* Extract out active_peers.rs

* wip: rework to struct from free function

* udp_listener working

* wip

* more udp_listener

* tun_device

* wip

* tun_device

* Remove some old commented out stuff

* tidy

* Remove commented out line
This commit is contained in:
Jon Häggblad
2023-10-17 12:01:11 +02:00
committed by GitHub
parent 53ab49cdd0
commit 96a925c040
7 changed files with 274 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/wireguard/src/active_peers.rs
+52
View File
@@ -0,0 +1,52 @@
use std::net::SocketAddr;
use boringtun::x25519;
use dashmap::{
mapref::one::{Ref, RefMut},
DashMap,
};
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>>;
#[derive(Default)]
pub(crate) struct ActivePeers {
active_peers: PeersByKey,
active_peers_by_addr: PeersByAddr,
}
impl ActivePeers {
pub(crate) fn remove(&self, public_key: &x25519::PublicKey) {
log::info!("Removing peer: {public_key:?}");
self.active_peers.remove(public_key);
log::warn!("TODO: remove from peers_by_ip?");
log::warn!("TODO: remove from peers_by_addr");
}
pub(crate) fn insert(
&self,
public_key: x25519::PublicKey,
addr: SocketAddr,
peer_tx: mpsc::UnboundedSender<Event>,
) {
self.active_peers.insert(public_key, peer_tx.clone());
self.active_peers_by_addr.insert(addr, peer_tx);
}
pub(crate) fn get_by_key_mut(
&self,
public_key: &x25519::PublicKey,
) -> Option<RefMut<'_, x25519::PublicKey, mpsc::UnboundedSender<Event>>> {
self.active_peers.get_mut(public_key)
}
pub(crate) fn get_by_addr(
&self,
addr: &SocketAddr,
) -> Option<Ref<'_, SocketAddr, mpsc::UnboundedSender<Event>>> {
self.active_peers_by_addr.get(addr)
}
}
common/wireguard/src/error.rs
+2
View File
@@ -4,4 +4,6 @@ use thiserror::Error;
pub enum WgError {
#[error("unable to get tunnel")]
UnableToGetTunnel,
#[error("handshake failed")]
HandshakeFailed,
}
common/wireguard/src/lib.rs
+6 -3
View File
@@ -3,6 +3,7 @@
// #![warn(clippy::expect_used)]
// #![warn(clippy::unwrap_used)]
mod active_peers;
mod error;
mod event;
mod network_table;
@@ -17,7 +18,7 @@ mod wg_tunnel;
use platform::linux::tun_device;
#[derive(Clone)]
struct TunTaskTx(tokio::sync::mpsc::UnboundedSender<Vec<u8>>);
pub struct TunTaskTx(tokio::sync::mpsc::UnboundedSender<Vec<u8>>);
impl TunTaskTx {
fn send(&self, packet: Vec<u8>) -> Result<(), tokio::sync::mpsc::error::SendError<Vec<u8>>> {
@@ -39,10 +40,12 @@ pub async fn start_wireguard(
let peers_by_ip = Arc::new(std::sync::Mutex::new(network_table::NetworkTable::new()));
// Start the tun device that is used to relay traffic outbound
let tun_task_tx = tun_device::start_tun_device(peers_by_ip.clone());
let (tun, tun_task_tx) = tun_device::TunDevice::new(peers_by_ip.clone());
tun.start();
// Start the UDP listener that clients connect to
udp_listener::start_udp_listener(tun_task_tx, peers_by_ip, task_client).await?;
let udp_listener = udp_listener::WgUdpListener::new(tun_task_tx, peers_by_ip).await?;
udp_listener.start(task_client);
Ok(())
}
common/wireguard/src/network_table.rs
+1 -1
View File
@@ -4,7 +4,7 @@ use ip_network::IpNetwork;
use ip_network_table::IpNetworkTable;
#[derive(Default)]
pub(crate) struct NetworkTable<T> {
pub struct NetworkTable<T> {
ips: IpNetworkTable<T>,
}
common/wireguard/src/platform/linux/tun_device.rs
+87 -48
View File
@@ -28,72 +28,111 @@ fn setup_tokio_tun_device(name: &str, address: Ipv4Addr, netmask: Ipv4Addr) -> t
.expect("Failed to setup tun device, do you have permission?")
}
pub(crate) fn start_tun_device(peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>) -> TunTaskTx {
let tun = setup_tokio_tun_device(
format!("{TUN_BASE_NAME}%d").as_str(),
TUN_DEVICE_ADDRESS.parse().unwrap(),
TUN_DEVICE_NETMASK.parse().unwrap(),
);
log::info!("Created TUN device: {}", tun.name());
pub struct TunDevice {
// The TUN device that we read/write to, to send/receive packets
tun: tokio_tun::Tun,
let (mut tun_device_rx, mut tun_device_tx) = tokio::io::split(tun);
// Incoming data that we should send
tun_task_rx: mpsc::UnboundedReceiver<Vec<u8>>,
// Channels to communicate with the other tasks
let (tun_task_tx, mut tun_task_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let tun_task_tx = TunTaskTx(tun_task_tx);
// The routing table.
// An alternative would be to do NAT by just matching incoming with outgoing.
peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>,
}
tokio::spawn(async move {
impl TunDevice {
pub fn new(peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>) -> (Self, TunTaskTx) {
let tun = setup_tokio_tun_device(
format!("{TUN_BASE_NAME}%d").as_str(),
TUN_DEVICE_ADDRESS.parse().unwrap(),
TUN_DEVICE_NETMASK.parse().unwrap(),
);
log::info!("Created TUN device: {}", tun.name());
// Channels to communicate with the other tasks
let (tun_task_tx, tun_task_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let tun_task_tx = TunTaskTx(tun_task_tx);
let tun_device = TunDevice {
tun_task_rx,
tun,
peers_by_ip,
};
(tun_device, tun_task_tx)
}
fn handle_tun_read(&self, packet: &[u8]) {
let dst_addr = boringtun::noise::Tunn::dst_address(packet).unwrap();
let headers = SlicedPacket::from_ip(packet).unwrap();
let src_addr = match headers.ip.unwrap() {
InternetSlice::Ipv4(ip, _) => ip.source_addr().to_string(),
InternetSlice::Ipv6(ip, _) => ip.source_addr().to_string(),
};
log::info!(
"iface: read Packet({src_addr} -> {dst_addr}, {} bytes)",
packet.len(),
);
// Route packet to the correct peer.
if let Some(peer_tx) = self
.peers_by_ip
.lock()
.unwrap()
.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}"))
.unwrap();
} else {
log::info!("No peer found, packet dropped");
}
}
async fn handle_tun_write(&mut self, data: Vec<u8>) {
let headers = SlicedPacket::from_ip(&data).unwrap();
let (source_addr, destination_addr) = match headers.ip.unwrap() {
InternetSlice::Ipv4(ip, _) => (ip.source_addr(), ip.destination_addr()),
InternetSlice::Ipv6(_, _) => unimplemented!(),
};
log::info!(
"iface: write Packet({source_addr} -> {destination_addr}, {} bytes)",
data.len()
);
self.tun.write_all(&data).await.unwrap();
}
pub async fn run(mut self) {
let mut buf = [0u8; 1024];
loop {
tokio::select! {
// Reading from the TUN device
len = tun_device_rx.read(&mut buf) => match len {
len = self.tun.read(&mut buf) => match len {
Ok(len) => {
let packet = &buf[..len];
let dst_addr = boringtun::noise::Tunn::dst_address(packet).unwrap();
let headers = SlicedPacket::from_ip(packet).unwrap();
let src_addr = match headers.ip.unwrap() {
InternetSlice::Ipv4(ip, _) => ip.source_addr().to_string(),
InternetSlice::Ipv6(ip, _) => ip.source_addr().to_string(),
};
log::info!("iface: read Packet({src_addr} -> {dst_addr}, {len} bytes)");
// Route packet to the correct peer.
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::Ip(packet.to_vec().into()))
.tap_err(|err| log::error!("{err}"))
.unwrap();
} else {
log::info!("No peer found, packet dropped");
}
self.handle_tun_read(packet);
},
Err(err) => {
log::info!("iface: read error: {err}");
break;
}
},
// Writing to the TUN device
Some(data) = tun_task_rx.recv() => {
let headers = SlicedPacket::from_ip(&data).unwrap();
let (source_addr, destination_addr) = match headers.ip.unwrap() {
InternetSlice::Ipv4(ip, _) => (ip.source_addr(), ip.destination_addr()),
InternetSlice::Ipv6(_, _) => unimplemented!(),
};
log::info!(
"iface: write Packet({source_addr} -> {destination_addr}, {} bytes)",
data.len()
);
// log::info!("iface: writing {} bytes", data.len());
tun_device_tx.write_all(&data).await.unwrap();
Some(data) = self.tun_task_rx.recv() => {
self.handle_tun_write(data).await;
}
}
}
log::info!("TUN device shutting down");
});
tun_task_tx
}
pub fn start(self) {
tokio::spawn(async move { self.run().await });
}
}
common/wireguard/src/udp_listener.rs
+113 -71
View File
@@ -4,7 +4,6 @@ use boringtun::{
noise::{self, handshake::parse_handshake_anon, rate_limiter::RateLimiter, TunnResult},
x25519,
};
use dashmap::DashMap;
use futures::StreamExt;
use log::error;
use nym_task::TaskClient;
@@ -15,6 +14,8 @@ use tokio::{
};
use crate::{
active_peers::ActivePeers,
error::WgError,
event::Event,
network_table::NetworkTable,
registered_peers::{RegisteredPeer, RegisteredPeers},
@@ -27,10 +28,6 @@ const MAX_PACKET: usize = 65535;
// Registered peers
pub(crate) type PeersByIp = NetworkTable<mpsc::UnboundedSender<Event>>;
// Active peers
pub(crate) type ActivePeers = DashMap<x25519::PublicKey, mpsc::UnboundedSender<Event>>;
pub(crate) type PeersByAddr = DashMap<SocketAddr, mpsc::UnboundedSender<Event>>;
async fn add_test_peer(registered_peers: &mut RegisteredPeers) {
let peer_static_public = setup::peer_static_public_key();
let peer_index = 0;
@@ -43,29 +40,62 @@ async fn add_test_peer(registered_peers: &mut RegisteredPeers) {
registered_peers.insert(peer_static_public, test_peer).await;
}
pub(crate) async fn start_udp_listener(
tun_task_tx: TunTaskTx,
pub struct WgUdpListener {
// Our private key
static_private: x25519::StaticSecret,
// Our public key
static_public: x25519::PublicKey,
// The list of registered peers that we allow
registered_peers: RegisteredPeers,
// The routing table, as defined by wireguard
peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>,
mut task_client: TaskClient,
) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
let wg_address = SocketAddr::new(WG_ADDRESS.parse().unwrap(), WG_PORT);
log::info!("Starting wireguard UDP listener on {wg_address}");
let udp = Arc::new(UdpSocket::bind(wg_address).await?);
// 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);
// The UDP socket to the peer
udp: Arc<UdpSocket>,
// Create a test peer for dev
let mut registered_peers = RegisteredPeers::default();
add_test_peer(&mut registered_peers).await;
// Send data to the TUN device for sending
tun_task_tx: TunTaskTx,
tokio::spawn(async move {
// The set of active tunnels indexed by the peer's address
let active_peers = Arc::new(ActivePeers::new());
let active_peers_by_addr = PeersByAddr::new();
// Wireguard rate limiter
rate_limiter: RateLimiter,
}
impl WgUdpListener {
pub async fn new(
tun_task_tx: TunTaskTx,
peers_by_ip: Arc<std::sync::Mutex<PeersByIp>>,
) -> Result<Self, Box<dyn std::error::Error + Send + Sync + 'static>> {
let wg_address = SocketAddr::new(WG_ADDRESS.parse().unwrap(), WG_PORT);
log::info!("Starting wireguard UDP listener on {wg_address}");
let udp = Arc::new(UdpSocket::bind(wg_address).await?);
// 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);
// Create a test peer for dev
let mut registered_peers = RegisteredPeers::default();
add_test_peer(&mut registered_peers).await;
Ok(Self {
static_private,
static_public,
registered_peers,
peers_by_ip,
udp,
tun_task_tx,
rate_limiter,
})
}
pub async fn run(self, mut task_client: TaskClient) {
// The set of active tunnels
let active_peers = ActivePeers::default();
// 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();
@@ -81,16 +111,13 @@ pub(crate) async fn start_udp_listener(
}
// Reset the rate limiter every 1 sec
() = tokio::time::sleep(Duration::from_secs(1)) => {
rate_limiter.reset_count();
self.rate_limiter.reset_count();
},
// Handle tunnel closing
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);
log::warn!("TODO: remove from peers_by_ip?");
log::warn!("TODO: remove from peers_by_addr");
}
Err(err) => {
error!("WireGuard UDP listener: error receiving shutdown from peer: {err}");
@@ -98,13 +125,13 @@ pub(crate) async fn start_udp_listener(
}
},
// Handle incoming packets
Ok((len, addr)) = udp.recv_from(&mut buf) => {
Ok((len, addr)) = self.udp.recv_from(&mut buf) => {
log::trace!("udp: received {} bytes from {}", len, addr);
// If this addr has already been encountered, send directly to tunnel
// TODO: optimization opportunity to instead create a connected UDP socket
// inside the wg tunnel, where you can recv the data directly.
if let Some(peer_tx) = active_peers_by_addr.get(&addr) {
if let Some(peer_tx) = active_peers.get_by_addr(&addr) {
log::info!("udp: received {len} bytes from {addr} from known peer");
peer_tx
.send(Event::Wg(buf[..len].to_vec().into()))
@@ -114,11 +141,11 @@ pub(crate) async fn start_udp_listener(
}
// Verify the incoming packet
let verified_packet = match rate_limiter.verify_packet(Some(addr.ip()), &buf[..len], &mut dst_buf) {
let verified_packet = match self.rate_limiter.verify_packet(Some(addr.ip()), &buf[..len], &mut dst_buf) {
Ok(packet) => packet,
Err(TunnResult::WriteToNetwork(cookie)) => {
log::info!("Send back cookie to: {addr}");
udp.send_to(cookie, addr).await.tap_err(|e| log::error!("{e}")).ok();
self.udp.send_to(cookie, addr).await.tap_err(|e| log::error!("{e}")).ok();
continue;
}
Err(err) => {
@@ -128,40 +155,25 @@ pub(crate) async fn start_udp_listener(
};
// Check if this is a registered peer, if not, just skip
let registered_peer = {
let reg_peer = match verified_packet {
noise::Packet::HandshakeInit(ref packet) => {
let Ok(handshake) = parse_handshake_anon(&static_private, &static_public, packet) else {
log::warn!("Handshake failed: {addr}");
continue;
};
registered_peers.get_by_key(&x25519::PublicKey::from(handshake.peer_static_public))
},
noise::Packet::HandshakeResponse(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
},
noise::Packet::PacketCookieReply(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
},
noise::Packet::PacketData(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
},
};
match reg_peer {
Some(reg_peer) => reg_peer.lock().await,
None => {
log::warn!("Peer not registered: {addr}");
continue;
}
let registered_peer = match parse_peer(
verified_packet,
&self.registered_peers,
&self.static_private,
&self.static_public
) {
Ok(Some(peer)) => peer.lock().await,
Ok(None) => {
log::warn!("Peer not registered: {addr}");
continue;
}
Err(err) => {
log::error!("{err}");
continue;
},
};
// Look up if the peer is already connected
if let Some(peer_tx) = active_peers.get_mut(&registered_peer.public_key) {
if let Some(peer_tx) = active_peers.get_by_key_mut(&registered_peer.public_key) {
// We found the peer as connected, even though the addr was not known
log::info!("udp: received {len} bytes from {addr} which is a known peer with unknown addr");
peer_tx.send(Event::WgVerified(buf[..len].to_vec().into()))
@@ -175,30 +187,60 @@ pub(crate) async fn start_udp_listener(
log::warn!("Creating new rate limiter, consider re-using?");
let (join_handle, peer_tx) = crate::wg_tunnel::start_wg_tunnel(
addr,
udp.clone(),
static_private.clone(),
self.udp.clone(),
self.static_private.clone(),
registered_peer.public_key,
registered_peer.index,
registered_peer.allowed_ips,
tun_task_tx.clone(),
self.tun_task_tx.clone(),
);
peers_by_ip.lock().unwrap().insert(registered_peer.allowed_ips, peer_tx.clone());
active_peers_by_addr.insert(addr, peer_tx.clone());
self.peers_by_ip.lock().unwrap().insert(registered_peer.allowed_ips, peer_tx.clone());
peer_tx.send(Event::Wg(buf[..len].to_vec().into()))
.tap_err(|e| log::error!("{e}"))
.ok();
log::info!("Adding peer: {addr}");
active_peers.insert(registered_peer.public_key, peer_tx);
log::info!("Adding peer: {:?}: {addr}", registered_peer.public_key);
active_peers.insert(registered_peer.public_key, addr, peer_tx);
active_peers_task_handles.push(join_handle);
}
},
}
}
log::info!("WireGuard listener: shutting down");
});
}
Ok(())
pub fn start(self, task_client: TaskClient) {
tokio::spawn(async move { self.run(task_client).await });
}
}
fn parse_peer<'a>(
verified_packet: noise::Packet,
registered_peers: &'a RegisteredPeers,
static_private: &x25519::StaticSecret,
static_public: &x25519::PublicKey,
) -> Result<Option<&'a Arc<tokio::sync::Mutex<RegisteredPeer>>>, WgError> {
let registered_peer = match verified_packet {
noise::Packet::HandshakeInit(ref packet) => {
let Ok(handshake) = parse_handshake_anon(static_private, static_public, packet) else {
return Err(WgError::HandshakeFailed);
};
registered_peers.get_by_key(&x25519::PublicKey::from(handshake.peer_static_public))
}
noise::Packet::HandshakeResponse(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
}
noise::Packet::PacketCookieReply(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
}
noise::Packet::PacketData(packet) => {
let peer_idx = packet.receiver_idx >> 8;
registered_peers.get_by_idx(peer_idx)
}
};
Ok(registered_peer)
}
gateway/src/node/mod.rs
+13 -1
View File
@@ -156,6 +156,15 @@ impl<St> Gateway<St> {
mixnet_handling::Listener::new(listening_address, shutdown).start(connection_handler);
}
#[cfg(feature = "wireguard")]
async fn start_wireguard(
&self,
shutdown: TaskClient,
) -> Result<(), Box<dyn Error + Send + Sync>> {
// TODO: possibly we should start the UDP listener and TUN device explicitly here
nym_wireguard::start_wireguard(shutdown).await
}
fn start_client_websocket_listener(
&self,
forwarding_channel: MixForwardingSender,
@@ -378,7 +387,10 @@ impl<St> Gateway<St> {
// Once this is a bit more mature, make this a commandline flag instead of a compile time
// flag
#[cfg(feature = "wireguard")]
if let Err(err) = nym_wireguard::start_wireguard(shutdown.subscribe()).await {
if let Err(err) = self
.start_wireguard(shutdown.subscribe().named("wireguard"))
.await
{
// that's a nasty workaround, but anyhow errors are generally nicer, especially on exit
bail!("{err}")
}