diff --git a/network/src/nat.rs b/network/src/nat.rs index d2b01ca..b3185a8 100644 --- a/network/src/nat.rs +++ b/network/src/nat.rs @@ -24,7 +24,6 @@ use tokio::sync::mpsc::Sender; pub enum NatKeyProtocol { Tcp, Udp, - Ping, } #[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] diff --git a/network/src/proxynat/mod.rs b/network/src/proxynat/mod.rs new file mode 100644 index 0000000..2cddef0 --- /dev/null +++ b/network/src/proxynat/mod.rs @@ -0,0 +1,49 @@ +mod udp; + +use async_trait::async_trait; + +use log::{debug, warn}; + +use tokio::sync::mpsc::channel; +use tokio::sync::mpsc::Sender; + +use crate::proxynat::udp::ProxyUdpHandler; + +use crate::nat::{NatHandler, NatHandlerFactory, NatKey, NatKeyProtocol}; + +pub struct ProxyNatHandlerFactory {} + +impl ProxyNatHandlerFactory { + pub fn new() -> Self { + Self {} + } +} + +#[async_trait] +impl NatHandlerFactory for ProxyNatHandlerFactory { + async fn nat(&self, key: NatKey, sender: Sender>) -> Option> { + debug!("creating proxy nat entry for key: {}", key); + + match key.protocol { + NatKeyProtocol::Udp => { + let (rx_sender, rx_receiver) = channel::>(4); + let mut handler = ProxyUdpHandler::new(key, rx_sender); + + if let Err(error) = handler.spawn(rx_receiver, sender.clone()).await { + warn!("unable to spawn udp proxy handler: {}", error); + None + } else { + Some(Box::new(handler)) + } + } + + _ => None, + } + } +} + +pub enum ProxyNatSelect { + External(usize), + Internal(Vec), + Closed, +} diff --git a/network/src/proxynat.rs b/network/src/proxynat/udp.rs similarity index 70% rename from network/src/proxynat.rs rename to network/src/proxynat/udp.rs index ea4ca83..a82a335 100644 --- a/network/src/proxynat.rs +++ b/network/src/proxynat/udp.rs @@ -11,49 +11,19 @@ use smoltcp::{ use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, select, - sync::mpsc::channel, }; use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender}; use udp_stream::UdpStream; -use crate::nat::{NatHandler, NatHandlerFactory, NatKey, NatKeyProtocol}; +use crate::nat::{NatHandler, NatKey}; -pub struct ProxyNatHandlerFactory {} +use super::ProxyNatSelect; -struct ProxyUdpHandler { +pub struct ProxyUdpHandler { key: NatKey, rx_sender: Sender>, } -impl ProxyNatHandlerFactory { - pub fn new() -> Self { - Self {} - } -} - -#[async_trait] -impl NatHandlerFactory for ProxyNatHandlerFactory { - async fn nat(&self, key: NatKey, sender: Sender>) -> Option> { - debug!("creating proxy nat entry for key: {}", key); - - match key.protocol { - NatKeyProtocol::Udp => { - let (rx_sender, rx_receiver) = channel::>(4); - let mut handler = ProxyUdpHandler { key, rx_sender }; - - if let Err(error) = handler.spawn(rx_receiver, sender.clone()).await { - warn!("unable to spawn udp proxy handler: {}", error); - None - } else { - Some(Box::new(handler)) - } - } - - _ => None, - } - } -} - #[async_trait] impl NatHandler for ProxyUdpHandler { async fn receive(&self, data: &[u8]) -> Result<()> { @@ -62,14 +32,12 @@ impl NatHandler for ProxyUdpHandler { } } -enum ProxySelect { - External(usize), - Internal(Vec), - Closed, -} - impl ProxyUdpHandler { - async fn spawn( + pub fn new(key: NatKey, rx_sender: Sender>) -> Self { + ProxyUdpHandler { key, rx_sender } + } + + pub async fn spawn( &mut self, rx_receiver: Receiver>, tx_sender: Sender>, @@ -109,15 +77,15 @@ impl ProxyUdpHandler { loop { let selection = select! { x = rx_receiver.recv() => if let Some(data) = x { - ProxySelect::Internal(data) + ProxyNatSelect::Internal(data) } else { - ProxySelect::Closed + ProxyNatSelect::Closed }, - x = socket.read(&mut external_buffer) => ProxySelect::External(x?), + x = socket.read(&mut external_buffer) => ProxyNatSelect::External(x?), }; match selection { - ProxySelect::External(size) => { + ProxyNatSelect::External(size) => { let data = &external_buffer[0..size]; let packet = PacketBuilder::ethernet2(key.local_mac.0, key.client_mac.0); let packet = match (key.external_ip.addr, key.client_ip.addr) { @@ -138,8 +106,7 @@ impl ProxyUdpHandler { debug!("failed to transmit udp packet: {}", error); } } - ProxySelect::Internal(data) => { - debug!("udp socket to handle data: {:?}", data); + ProxyNatSelect::Internal(data) => { let packet = SlicedPacket::from_ethernet(&data)?; let Some(ref net) = packet.net else { continue; @@ -150,10 +117,9 @@ impl ProxyUdpHandler { }; let udp = UdpSlice::from_slice(ip.payload)?; - debug!("UDP from internal: {:?}", udp.payload()); socket.write_all(udp.payload()).await?; } - ProxySelect::Closed => warn!("UDP socket closed"), + ProxyNatSelect::Closed => warn!("UDP socket closed"), } } }