network: performance improvements across NAT stack

This commit is contained in:
Alex Zenla
2024-02-11 11:53:59 +00:00
parent b9dadc6f37
commit f96f9d8abf
6 changed files with 140 additions and 56 deletions

View File

@ -27,9 +27,13 @@ pub struct ProxyIcmpHandler {
#[async_trait]
impl NatHandler for ProxyIcmpHandler {
async fn receive(&self, data: &[u8]) -> Result<()> {
self.rx_sender.try_send(data.to_vec())?;
Ok(())
async fn receive(&self, data: &[u8]) -> Result<bool> {
if self.rx_sender.is_closed() {
Ok(true)
} else {
self.rx_sender.try_send(data.to_vec())?;
Ok(true)
}
}
}

View File

@ -47,7 +47,7 @@ impl NatHandlerFactory for ProxyNatHandlerFactory {
}
NatKeyProtocol::Icmp => {
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(4);
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(300);
let mut handler = ProxyIcmpHandler::new(rx_sender);
if let Err(error) = handler.spawn(context, rx_receiver).await {
@ -59,7 +59,7 @@ impl NatHandlerFactory for ProxyNatHandlerFactory {
}
NatKeyProtocol::Tcp => {
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(4);
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(300);
let mut handler = ProxyTcpHandler::new(rx_sender);
if let Err(error) = handler.spawn(context, rx_receiver).await {

View File

@ -29,6 +29,7 @@ use crate::{
const TCP_BUFFER_SIZE: usize = 65535;
const TCP_ACCEPT_TIMEOUT_SECS: u64 = 120;
const TCP_DANGLE_TIMEOUT_SECS: u64 = 10;
pub struct ProxyTcpHandler {
rx_sender: Sender<Vec<u8>>,
@ -36,9 +37,13 @@ pub struct ProxyTcpHandler {
#[async_trait]
impl NatHandler for ProxyTcpHandler {
async fn receive(&self, data: &[u8]) -> Result<()> {
self.rx_sender.try_send(data.to_vec())?;
Ok(())
async fn receive(&self, data: &[u8]) -> Result<bool> {
if self.rx_sender.is_closed() {
Ok(false)
} else {
self.rx_sender.try_send(data.to_vec())?;
Ok(true)
}
}
}
@ -62,6 +67,13 @@ enum ProxyTcpDataSelect {
Close,
}
#[derive(Debug)]
enum ProxyTcpFinishSelect {
InternalRecv(Vec<u8>),
TxIpPacket(Vec<u8>),
Close,
}
impl ProxyTcpHandler {
pub fn new(rx_sender: Sender<Vec<u8>>) -> Self {
ProxyTcpHandler { rx_sender }
@ -95,10 +107,14 @@ impl ProxyTcpHandler {
mut external_socket: TcpStream,
mut rx_receiver: Receiver<Vec<u8>>,
) -> Result<()> {
let (ip_sender, mut ip_receiver) = channel::<Vec<u8>>(4);
let (ip_sender, mut ip_receiver) = channel::<Vec<u8>>(300);
let mut external_buffer = vec![0u8; TCP_BUFFER_SIZE];
let mut device = ChannelDevice::new(context.mtu, Medium::Ip, ip_sender.clone());
let mut device = ChannelDevice::new(
context.mtu - Ethernet2Header::LEN,
Medium::Ip,
ip_sender.clone(),
);
let config = Config::new(HardwareAddress::Ip);
let tcp_rx_buffer = SocketBuffer::new(vec![0; TCP_BUFFER_SIZE]);
@ -212,10 +228,10 @@ impl ProxyTcpHandler {
.get_mut::<tcp::Socket>(internal_socket_handle)
.is_active()
{
debug!("failed to accept tcp connection from client");
true
} else {
true
debug!("failed to accept tcp connection from client");
false
};
let mut already_shutdown = false;
@ -268,13 +284,13 @@ impl ProxyTcpHandler {
if !do_shutdown {
select! {
biased;
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
} else {
ProxyTcpDataSelect::Close
},
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
} else {
ProxyTcpDataSelect::Close
},
@ -285,13 +301,13 @@ impl ProxyTcpHandler {
} else {
select! {
biased;
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
} else {
ProxyTcpDataSelect::Close
},
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
} else {
ProxyTcpDataSelect::Close
},
@ -303,13 +319,13 @@ impl ProxyTcpHandler {
} else if !do_shutdown {
select! {
biased;
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
} else {
ProxyTcpDataSelect::Close
},
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
} else {
ProxyTcpDataSelect::Close
},
@ -320,13 +336,13 @@ impl ProxyTcpHandler {
} else {
select! {
biased;
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
} else {
ProxyTcpDataSelect::Close
},
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::TxIpPacket(data)
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpDataSelect::InternalRecv(data)
} else {
ProxyTcpDataSelect::Close
},
@ -389,6 +405,58 @@ impl ProxyTcpHandler {
}
}
let _ = external_socket.shutdown().await;
drop(external_socket);
loop {
let deadline = tokio::time::sleep(Duration::from_secs(TCP_DANGLE_TIMEOUT_SECS));
tokio::pin!(deadline);
let selection = select! {
biased;
x = ip_receiver.recv() => if let Some(data) = x {
ProxyTcpFinishSelect::TxIpPacket(data)
} else {
ProxyTcpFinishSelect::Close
},
x = rx_receiver.recv() => if let Some(data) = x {
ProxyTcpFinishSelect::InternalRecv(data)
} else {
ProxyTcpFinishSelect::Close
},
_ = deadline => ProxyTcpFinishSelect::Close,
};
match selection {
ProxyTcpFinishSelect::InternalRecv(data) => {
let (_, payload) = Ethernet2Header::from_slice(&data)?;
device.rx = Some(payload.to_vec());
iface.poll(Instant::now(), &mut device, &mut sockets);
}
ProxyTcpFinishSelect::TxIpPacket(payload) => {
let mut buffer: Vec<u8> = Vec::new();
let header = Ethernet2Header {
source: context.key.local_mac.0,
destination: context.key.client_mac.0,
ether_type: match context.key.external_ip.addr {
IpAddress::Ipv4(_) => EtherType::IPV4,
IpAddress::Ipv6(_) => EtherType::IPV6,
},
};
header.write(&mut buffer)?;
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
debug!("failed to transmit tcp packet: {}", error);
}
}
ProxyTcpFinishSelect::Close => {
break;
}
}
}
context.reclaim().await?;
Ok(())

View File

@ -25,9 +25,13 @@ pub struct ProxyUdpHandler {
#[async_trait]
impl NatHandler for ProxyUdpHandler {
async fn receive(&self, data: &[u8]) -> Result<()> {
self.rx_sender.try_send(data.to_vec())?;
Ok(())
async fn receive(&self, data: &[u8]) -> Result<bool> {
if self.rx_sender.is_closed() {
Ok(true)
} else {
self.rx_sender.try_send(data.to_vec())?;
Ok(true)
}
}
}