diff --git a/common/wireguard/src/platform/linux/tun_device.rs b/common/wireguard/src/platform/linux/tun_device.rs index 1bf84eb52d..196102394f 100644 --- a/common/wireguard/src/platform/linux/tun_device.rs +++ b/common/wireguard/src/platform/linux/tun_device.rs @@ -42,11 +42,35 @@ pub struct TunDevice { // And when we get replies, this is where we should send it tun_task_response_tx: TunTaskResponseTx, + routing_mode: RoutingMode, +} + +enum RoutingMode { // The routing table, as how wireguard does it - peers_by_ip: Option>>, + AllowedIps(AllowedIpsInner), // This is an alternative to the routing table, where we just match outgoing source IP with // incoming destination IP. + Nat(NatInner), +} + +impl RoutingMode { + fn new_nat() -> Self { + RoutingMode::Nat(NatInner { + nat_table: HashMap::new(), + }) + } + + fn new_allowed_ips(peers_by_ip: Arc>) -> Self { + RoutingMode::AllowedIps(AllowedIpsInner { peers_by_ip }) + } +} + +struct AllowedIpsInner { + peers_by_ip: Arc>, +} + +struct NatInner { nat_table: HashMap, } @@ -65,12 +89,16 @@ impl TunDevice { let (tun_task_tx, tun_task_rx) = tun_task_channel(); let (tun_task_response_tx, tun_task_response_rx) = tun_task_response_channel(); + let routing_mode = match peers_by_ip { + Some(peers_by_ip) => RoutingMode::new_allowed_ips(peers_by_ip), + None => RoutingMode::new_nat(), + }; + let tun_device = TunDevice { tun_task_rx, tun_task_response_tx, tun, - peers_by_ip, - nat_table: HashMap::new(), + routing_mode, }; (tun_device, tun_task_tx, tun_task_response_rx) @@ -93,7 +121,9 @@ impl TunDevice { ); // TODO: expire old entries - self.nat_table.insert(src_addr, tag); + if let RoutingMode::Nat(nat_table) = &mut self.routing_mode { + nat_table.nat_table.insert(src_addr, tag); + } self.tun .write_all(&packet) @@ -121,30 +151,32 @@ impl TunDevice { // Route packet to the correct peer. - // This is how wireguard does it, by consulting the AllowedIPs table. - if false { - let peers = self.peers_by_ip.as_ref().unwrap().lock().await; - 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())) - .await - .tap_err(|err| log::error!("{err}")) - .ok(); - return; + match self.routing_mode { + // This is how wireguard does it, by consulting the AllowedIPs table. + RoutingMode::AllowedIps(ref peers_by_ip) => { + let peers = peers_by_ip.peers_by_ip.as_ref().lock().await; + 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())) + .await + .tap_err(|err| log::error!("{err}")) + .ok(); + return; + } } - } - // But we do it by consulting the NAT table. - { - if let Some(tag) = self.nat_table.get(&dst_addr) { - log::info!("Forward packet to wg tunnel with tag: {tag}"); - self.tun_task_response_tx - .send((*tag, packet.to_vec())) - .await - .tap_err(|err| log::error!("{err}")) - .ok(); - return; + // But we do it by consulting the NAT table. + RoutingMode::Nat(ref nat_table) => { + if let Some(tag) = nat_table.nat_table.get(&dst_addr) { + log::info!("Forward packet to wg tunnel with tag: {tag}"); + self.tun_task_response_tx + .send((*tag, packet.to_vec())) + .await + .tap_err(|err| log::error!("{err}")) + .ok(); + return; + } } }