Checkpoint

src/virtual_iface/tcp.rs

@@ -75,7 +75,7 @@ impl VirtualInterfacePoll for TcpVirtualInterface {
         // Consumer for IP packets to send through the virtual interface
         // Initialize the interface
         let device =
-            VirtualIpDevice::new(VirtualPort(self.virtual_port, PortProtocol::Tcp), self.wg)
+            VirtualIpDevice::new_direct(VirtualPort(self.virtual_port, PortProtocol::Tcp), self.wg)
                 .with_context(|| "Failed to initialize TCP VirtualIpDevice")?;
         let mut virtual_interface = InterfaceBuilder::new(device)
             .ip_addrs([

src/virtual_iface/udp.rs

@@ -1,11 +1,18 @@
+use anyhow::Context;
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
+use dashmap::DashMap;
+use smoltcp::iface::InterfaceBuilder;
+use smoltcp::socket::{SocketSet, UdpPacketMetadata, UdpSocket, UdpSocketBuffer};
+use smoltcp::wire::{IpAddress, IpCidr};
 
 use crate::config::PortForwardConfig;
+use crate::virtual_device::VirtualIpDevice;
 use crate::virtual_iface::{VirtualInterfacePoll, VirtualPort};
-use crate::wg::WireGuardTunnel;
+use crate::wg::{WireGuardTunnel, DISPATCH_CAPACITY};
 
 pub struct UdpVirtualInterface {
     port_forward: PortForwardConfig,
@@ -37,16 +44,68 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
         let mut data_to_virtual_server_rx = self.data_to_virtual_server_rx;
 
         // The IP to bind client sockets to
-        let _source_peer_ip = self.wg.source_peer_ip;
+        let source_peer_ip = self.wg.source_peer_ip;
 
         // The IP/port to bind the server socket to
-        let _destination = self.port_forward.destination;
+        let destination = self.port_forward.destination;
+
+        // Initialize a channel for IP packets.
+        // The "base transmitted" is cloned so that each virtual port can register a sender in the tunnel.
+        // The receiver is given to the device so that the Virtual Interface can process incoming IP packets from the tunnel.
+        let (base_ip_dispatch_tx, ip_dispatch_rx) = tokio::sync::mpsc::channel(DISPATCH_CAPACITY);
+
+        let device = VirtualIpDevice::new(self.wg.clone(), ip_dispatch_rx);
+        let mut virtual_interface = InterfaceBuilder::new(device)
+            .ip_addrs([
+                // Interface handles IP packets for the sender and recipient
+                IpCidr::new(source_peer_ip.into(), 32),
+                IpCidr::new(destination.ip().into(), 32),
+            ])
+            .finalize();
+
+        // Server socket: this is a placeholder for the interface.
+        let server_socket: anyhow::Result<UdpSocket> = {
+            static mut UDP_SERVER_RX_META: [UdpPacketMetadata; 0] = [];
+            static mut UDP_SERVER_RX_DATA: [u8; 0] = [];
+            static mut UDP_SERVER_TX_META: [UdpPacketMetadata; 0] = [];
+            static mut UDP_SERVER_TX_DATA: [u8; 0] = [];
+            let udp_rx_buffer =
+                UdpSocketBuffer::new(unsafe { &mut UDP_SERVER_RX_META[..] }, unsafe {
+                    &mut UDP_SERVER_RX_DATA[..]
+                });
+            let udp_tx_buffer =
+                UdpSocketBuffer::new(unsafe { &mut UDP_SERVER_TX_META[..] }, unsafe {
+                    &mut UDP_SERVER_TX_DATA[..]
+                });
+            let mut socket = UdpSocket::new(udp_rx_buffer, udp_tx_buffer);
+
+            socket
+                .bind((IpAddress::from(destination.ip()), destination.port()))
+                .with_context(|| "UDP virtual server socket failed to listen")?;
+
+            Ok(socket)
+        };
+
+        let mut socket_set = SocketSet::new(vec![]);
+        let _server_handle = socket_set.add(server_socket?);
 
         loop {
             let _loop_start = smoltcp::time::Instant::now();
+            let wg = self.wg.clone();
             // TODO: smoltcp UDP
 
             if let Ok((client_port, data)) = data_to_virtual_server_rx.try_recv() {
+                // Register the socket in WireGuard Tunnel if not already
+                if !wg.is_registered(client_port) {
+                    wg.register_virtual_interface(client_port, base_ip_dispatch_tx.clone())
+                        .unwrap_or_else(|e| {
+                            error!(
+                                "[{}] Failed to register UDP socket in WireGuard tunnel",
+                                client_port
+                            );
+                        });
+                }
+
                 // TODO: Find the matching client socket and send
                 // Echo for now
                 self.data_to_real_client_tx