network: successfully implement NAT support for UDP. timeouts are not yet implemented.

This commit is contained in:
Alex Zenla
2024-02-10 12:00:15 +00:00
parent 360506bbdd
commit d2f86afb7c
3 changed files with 109 additions and 65 deletions

View File

@ -3,6 +3,7 @@
use anyhow::Result;
use async_trait::async_trait;
use etherparse::Ethernet2Slice;
use etherparse::IpNumber;
use etherparse::IpPayloadSlice;
use etherparse::Ipv4Slice;
@ -17,32 +18,31 @@ use smoltcp::wire::IpEndpoint;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt::Display;
use tokio::sync::mpsc::Sender;
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub enum NatKey {
Tcp {
client: IpEndpoint,
external: IpEndpoint,
},
pub enum NatKeyProtocol {
Tcp,
Udp,
Ping,
}
Udp {
client: IpEndpoint,
external: IpEndpoint,
},
Ping {
client: IpAddress,
external: IpAddress,
},
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct NatKey {
pub protocol: NatKeyProtocol,
pub client_mac: EthernetAddress,
pub local_mac: EthernetAddress,
pub client_ip: IpEndpoint,
pub external_ip: IpEndpoint,
}
impl Display for NatKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
NatKey::Tcp { client, external } => write!(f, "TCP {client} -> {external}"),
NatKey::Udp { client, external } => write!(f, "UDP {client} -> {external}"),
NatKey::Ping { client, external } => write!(f, "Ping {client} -> {external}"),
}
write!(
f,
"{} -> {} {:?} {} -> {}",
self.client_mac, self.local_mac, self.protocol, self.client_ip, self.external_ip
)
}
}
@ -51,6 +51,11 @@ pub trait NatHandler: Send {
async fn receive(&self, packet: &[u8]) -> Result<()>;
}
#[async_trait]
pub trait NatHandlerFactory: Send {
async fn nat(&self, key: NatKey, sender: Sender<Vec<u8>>) -> Option<Box<dyn NatHandler>>;
}
pub struct NatTable {
inner: HashMap<NatKey, Box<dyn NatHandler>>,
}
@ -63,23 +68,24 @@ impl NatTable {
}
}
#[async_trait]
pub trait NatHandlerFactory: Send {
async fn nat(&self, key: NatKey) -> Option<Box<dyn NatHandler>>;
}
pub struct NatRouter {
_mac: EthernetAddress,
_local_mac: EthernetAddress,
factory: Box<dyn NatHandlerFactory>,
table: NatTable,
tx_sender: Sender<Vec<u8>>,
}
impl NatRouter {
pub fn new(factory: Box<dyn NatHandlerFactory>, mac: EthernetAddress) -> Self {
pub fn new(
factory: Box<dyn NatHandlerFactory>,
mac: EthernetAddress,
tx_sender: Sender<Vec<u8>>,
) -> Self {
Self {
_mac: mac,
_local_mac: mac,
factory,
table: NatTable::new(),
tx_sender,
}
}
@ -101,8 +107,9 @@ impl NatRouter {
match net {
NetSlice::Ipv4(ipv4) => {
self.process_ipv4(data, ipv4).await?;
self.process_ipv4(data, ether, ipv4).await?;
}
_ => {
return Ok(());
}
@ -111,18 +118,22 @@ impl NatRouter {
Ok(())
}
pub async fn process_ipv4<'a>(&mut self, data: &[u8], ipv4: &Ipv4Slice<'a>) -> Result<()> {
pub async fn process_ipv4<'a>(
&mut self,
data: &[u8],
ether: &Ethernet2Slice<'a>,
ipv4: &Ipv4Slice<'a>,
) -> Result<()> {
let source_addr = IpAddress::Ipv4(ipv4.header().source_addr().into());
let dest_addr = IpAddress::Ipv4(ipv4.header().destination_addr().into());
match ipv4.header().protocol() {
IpNumber::TCP => {
self.process_tcp(data, source_addr, dest_addr, ipv4.payload())
self.process_tcp(data, ether, source_addr, dest_addr, ipv4.payload())
.await?;
}
IpNumber::UDP => {
self.process_udp(data, source_addr, dest_addr, ipv4.payload())
self.process_udp(data, ether, source_addr, dest_addr, ipv4.payload())
.await?;
}
@ -135,6 +146,7 @@ impl NatRouter {
pub async fn process_tcp<'a>(
&mut self,
data: &'a [u8],
ether: &Ethernet2Slice<'a>,
source_addr: IpAddress,
dest_addr: IpAddress,
payload: &IpPayloadSlice<'a>,
@ -142,9 +154,12 @@ impl NatRouter {
let header = TcpHeaderSlice::from_slice(payload.payload)?;
let source = IpEndpoint::new(source_addr, header.source_port());
let dest = IpEndpoint::new(dest_addr, header.destination_port());
let key = NatKey::Tcp {
client: source,
external: dest,
let key = NatKey {
protocol: NatKeyProtocol::Tcp,
client_mac: EthernetAddress(ether.destination()),
local_mac: EthernetAddress(ether.source()),
client_ip: source,
external_ip: dest,
};
self.process_nat(data, key).await?;
Ok(())
@ -153,6 +168,7 @@ impl NatRouter {
pub async fn process_udp<'a>(
&mut self,
data: &'a [u8],
ether: &Ethernet2Slice<'a>,
source_addr: IpAddress,
dest_addr: IpAddress,
payload: &IpPayloadSlice<'a>,
@ -160,9 +176,12 @@ impl NatRouter {
let header = UdpHeaderSlice::from_slice(payload.payload)?;
let source = IpEndpoint::new(source_addr, header.source_port());
let dest = IpEndpoint::new(dest_addr, header.destination_port());
let key = NatKey::Udp {
client: source,
external: dest,
let key = NatKey {
protocol: NatKeyProtocol::Udp,
client_mac: EthernetAddress(ether.destination()),
local_mac: EthernetAddress(ether.source()),
client_ip: source,
external_ip: dest,
};
self.process_nat(data, key).await?;
Ok(())
@ -172,7 +191,7 @@ impl NatRouter {
let handler: Option<&mut Box<dyn NatHandler>> = match self.table.inner.entry(key) {
Entry::Occupied(entry) => Some(entry.into_mut()),
Entry::Vacant(entry) => {
if let Some(handler) = self.factory.nat(key).await {
if let Some(handler) = self.factory.nat(key, self.tx_sender.clone()).await {
Some(entry.insert(handler))
} else {
None