End-to-end UDP implementation

Port re-use still needs to be implemented to prevent exhaustion over
time, and flooding.

src/virtual_iface/udp.rs

@@ -4,9 +4,8 @@ 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::socket::{SocketHandle, SocketSet, UdpPacketMetadata, UdpSocket, UdpSocketBuffer};
 use smoltcp::wire::{IpAddress, IpCidr};
 
 use crate::config::PortForwardConfig;
@@ -14,6 +13,8 @@ use crate::virtual_device::VirtualIpDevice;
 use crate::virtual_iface::{VirtualInterfacePoll, VirtualPort};
 use crate::wg::{WireGuardTunnel, DISPATCH_CAPACITY};
 
+const MAX_PACKET: usize = 65536;
+
 pub struct UdpVirtualInterface {
     port_forward: PortForwardConfig,
     wg: Arc<WireGuardTunnel>,
@@ -89,31 +90,88 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
         let mut socket_set = SocketSet::new(vec![]);
         let _server_handle = socket_set.add(server_socket?);
 
+        // A map of virtual port to client socket.
+        let mut client_sockets: HashMap<VirtualPort, SocketHandle> = HashMap::new();
+
         loop {
-            let _loop_start = smoltcp::time::Instant::now();
+            let loop_start = smoltcp::time::Instant::now();
             let wg = self.wg.clone();
-            // TODO: smoltcp UDP
+
+            match virtual_interface.poll(&mut socket_set, loop_start) {
+                Ok(processed) if processed => {
+                    trace!("UDP virtual interface polled some packets to be processed");
+                }
+                Err(e) => error!("UDP virtual interface poll error: {:?}", e),
+                _ => {}
+            }
+
+            // Loop through each client socket and check if there is any data to send back
+            // to the real client.
+            for (virtual_port, client_socket_handle) in client_sockets.iter() {
+                let mut client_socket = socket_set.get::<UdpSocket>(*client_socket_handle);
+                match client_socket.recv() {
+                    Ok((data, _peer)) => {
+                        // Send the data back to the real client using MPSC channel
+                        self.data_to_real_client_tx
+                            .send((*virtual_port, data.to_vec()))
+                            .await
+                            .unwrap_or_else(|e| {
+                                error!(
+                                    "[{}] Failed to dispatch data from virtual client to real client: {:?}",
+                                    virtual_port, e
+                                );
+                            });
+                    }
+                    Err(smoltcp::Error::Exhausted) => {}
+                    Err(e) => {
+                        error!(
+                            "[{}] Failed to read from virtual client socket: {:?}",
+                            virtual_port, e
+                        );
+                    }
+                }
+            }
 
             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
-                    .send((client_port, data))
-                    .await
+                // Register the socket in WireGuard Tunnel (overrides any previous registration as well)
+                wg.register_virtual_interface(client_port, base_ip_dispatch_tx.clone())
                     .unwrap_or_else(|e| {
                         error!(
-                            "[{}] Failed to dispatch data from virtual client to real client: {:?}",
+                            "[{}] Failed to register UDP socket in WireGuard tunnel: {:?}",
+                            client_port, e
+                        );
+                    });
+
+                let client_socket_handle = client_sockets.entry(client_port).or_insert_with(|| {
+                    let rx_meta = vec![UdpPacketMetadata::EMPTY; MAX_PACKET];
+                    let tx_meta = vec![UdpPacketMetadata::EMPTY; MAX_PACKET];
+                    let rx_data = vec![0u8; MAX_PACKET];
+                    let tx_data = vec![0u8; MAX_PACKET];
+                    let udp_rx_buffer = UdpSocketBuffer::new(rx_meta, rx_data);
+                    let udp_tx_buffer = UdpSocketBuffer::new(tx_meta, tx_data);
+                    let mut socket = UdpSocket::new(udp_rx_buffer, udp_tx_buffer);
+
+                    socket
+                        .bind((IpAddress::from(wg.source_peer_ip), client_port.0))
+                        .unwrap_or_else(|e| {
+                            error!(
+                                "[{}] UDP virtual client socket failed to bind: {:?}",
+                                client_port, e
+                            );
+                        });
+
+                    socket_set.add(socket)
+                });
+
+                let mut client_socket = socket_set.get::<UdpSocket>(*client_socket_handle);
+                client_socket
+                    .send_slice(
+                        &data,
+                        (IpAddress::from(destination.ip()), destination.port()).into(),
+                    )
+                    .unwrap_or_else(|e| {
+                        error!(
+                            "[{}] Failed to send data to virtual server: {:?}",
                             client_port, e
                         );
                     });
@@ -121,7 +179,5 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
 
             tokio::time::sleep(Duration::from_millis(1)).await;
         }
-
-        // Ok(())
     }
 }