network: implement TCP proxy NAT via smoltcp

network/src/proxynat/icmp.rs

@@ -205,6 +205,8 @@ impl ProxyIcmpHandler {
             }
         }
 
+        context.reclaim().await?;
+
         Ok(())
     }
 }

network/src/proxynat/mod.rs

@@ -10,8 +10,10 @@ use crate::proxynat::udp::ProxyUdpHandler;
 use crate::nat::{NatHandler, NatHandlerFactory, NatKeyProtocol};
 
 use self::icmp::ProxyIcmpHandler;
+use self::tcp::ProxyTcpHandler;
 
 mod icmp;
+mod tcp;
 mod udp;
 
 pub struct ProxyNatHandlerFactory {}
@@ -56,7 +58,17 @@ impl NatHandlerFactory for ProxyNatHandlerFactory {
                 }
             }
 
-            _ => None,
+            NatKeyProtocol::Tcp => {
+                let (rx_sender, rx_receiver) = channel::<Vec<u8>>(4);
+                let mut handler = ProxyTcpHandler::new(rx_sender);
+
+                if let Err(error) = handler.spawn(context, rx_receiver).await {
+                    warn!("unable to spawn tcp proxy handler: {}", error);
+                    None
+                } else {
+                    Some(Box::new(handler))
+                }
+            }
         }
     }
 }

network/src/proxynat/tcp.rs

@@ -0,0 +1,396 @@
+use std::{
+    net::{IpAddr, SocketAddr},
+    time::Duration,
+};
+
+use anyhow::Result;
+use async_trait::async_trait;
+use etherparse::{EtherType, Ethernet2Header};
+use log::{debug, warn};
+use smoltcp::{
+    iface::{Config, Interface, SocketSet, SocketStorage},
+    phy::Medium,
+    socket::tcp::{self, SocketBuffer, State},
+    time::Instant,
+    wire::{HardwareAddress, IpAddress, IpCidr},
+};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
+    select,
+    sync::mpsc::channel,
+};
+use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
+
+use crate::{
+    chandev::ChannelDevice,
+    nat::{NatHandler, NatHandlerContext},
+};
+
+const TCP_BUFFER_SIZE: usize = 65535;
+const TCP_ACCEPT_TIMEOUT_SECS: u64 = 120;
+
+pub struct ProxyTcpHandler {
+    rx_sender: Sender<Vec<u8>>,
+}
+
+#[async_trait]
+impl NatHandler for ProxyTcpHandler {
+    async fn receive(&self, data: &[u8]) -> Result<()> {
+        self.rx_sender.try_send(data.to_vec())?;
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+enum ProxyTcpAcceptSelect {
+    Internal(Vec<u8>),
+    TxIpPacket(Vec<u8>),
+    TimePassed,
+    DoNothing,
+    Close,
+}
+
+#[derive(Debug)]
+enum ProxyTcpDataSelect {
+    ExternalRecv(usize),
+    ExternalSent(usize),
+    InternalRecv(Vec<u8>),
+    TxIpPacket(Vec<u8>),
+    TimePassed,
+    DoNothing,
+    Close,
+}
+
+impl ProxyTcpHandler {
+    pub fn new(rx_sender: Sender<Vec<u8>>) -> Self {
+        ProxyTcpHandler { rx_sender }
+    }
+
+    pub async fn spawn(
+        &mut self,
+        context: NatHandlerContext,
+        rx_receiver: Receiver<Vec<u8>>,
+    ) -> Result<()> {
+        let external_addr = match context.key.external_ip.addr {
+            IpAddress::Ipv4(addr) => {
+                SocketAddr::new(IpAddr::V4(addr.0.into()), context.key.external_ip.port)
+            }
+            IpAddress::Ipv6(addr) => {
+                SocketAddr::new(IpAddr::V6(addr.0.into()), context.key.external_ip.port)
+            }
+        };
+
+        let socket = TcpStream::connect(external_addr).await?;
+        tokio::spawn(async move {
+            if let Err(error) = ProxyTcpHandler::process(context, socket, rx_receiver).await {
+                warn!("processing of tcp proxy failed: {}", error);
+            }
+        });
+        Ok(())
+    }
+
+    async fn process(
+        context: NatHandlerContext,
+        mut external_socket: TcpStream,
+        mut rx_receiver: Receiver<Vec<u8>>,
+    ) -> Result<()> {
+        let (ip_sender, mut ip_receiver) = channel::<Vec<u8>>(4);
+        let mut external_buffer = vec![0u8; TCP_BUFFER_SIZE];
+
+        let mut device = ChannelDevice::new(context.mtu, Medium::Ip, ip_sender.clone());
+        let config = Config::new(HardwareAddress::Ip);
+
+        let tcp_rx_buffer = SocketBuffer::new(vec![0; TCP_BUFFER_SIZE]);
+        let tcp_tx_buffer = SocketBuffer::new(vec![0; TCP_BUFFER_SIZE]);
+        let internal_socket = tcp::Socket::new(tcp_rx_buffer, tcp_tx_buffer);
+        let mut iface = Interface::new(config, &mut device, Instant::now());
+
+        iface.update_ip_addrs(|addrs| {
+            let _ = addrs.push(IpCidr::new(context.key.external_ip.addr, 0));
+        });
+
+        let mut sockets = SocketSet::new([SocketStorage::EMPTY]);
+        let internal_socket_handle = sockets.add(internal_socket);
+        let (mut external_r, mut external_w) = external_socket.split();
+
+        {
+            let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
+            socket.connect(
+                iface.context(),
+                context.key.client_ip,
+                context.key.external_ip,
+            )?;
+        }
+
+        iface.poll(Instant::now(), &mut device, &mut sockets);
+
+        let mut sleeper: Option<tokio::time::Sleep> = None;
+        loop {
+            let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
+            if socket.is_active() && socket.state() != State::SynSent {
+                break;
+            }
+
+            if socket.state() == State::Closed {
+                break;
+            }
+
+            let deadline = tokio::time::sleep(Duration::from_secs(TCP_ACCEPT_TIMEOUT_SECS));
+            let selection = if let Some(sleep) = sleeper.take() {
+                select! {
+                    biased;
+                    x = rx_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpAcceptSelect::Internal(data)
+                    } else {
+                        ProxyTcpAcceptSelect::Close
+                    },
+                    x = ip_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpAcceptSelect::TxIpPacket(data)
+                    } else {
+                        ProxyTcpAcceptSelect::Close
+                    },
+                    _ = sleep => ProxyTcpAcceptSelect::TimePassed,
+                    _ = deadline => ProxyTcpAcceptSelect::Close,
+                }
+            } else {
+                select! {
+                    biased;
+                    x = rx_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpAcceptSelect::Internal(data)
+                    } else {
+                        ProxyTcpAcceptSelect::Close
+                    },
+                    x = ip_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpAcceptSelect::TxIpPacket(data)
+                    } else {
+                        ProxyTcpAcceptSelect::Close
+                    },
+                    _ = std::future::ready(()) => ProxyTcpAcceptSelect::DoNothing,
+                    _ = deadline => ProxyTcpAcceptSelect::Close,
+                }
+            };
+            match selection {
+                ProxyTcpAcceptSelect::TimePassed => {
+                    iface.poll(Instant::now(), &mut device, &mut sockets);
+                }
+
+                ProxyTcpAcceptSelect::DoNothing => {
+                    sleeper = Some(tokio::time::sleep(Duration::from_millis(50)));
+                }
+
+                ProxyTcpAcceptSelect::Internal(data) => {
+                    let (_, payload) = Ethernet2Header::from_slice(&data)?;
+                    device.rx = Some(payload.to_vec());
+                    iface.poll(Instant::now(), &mut device, &mut sockets);
+                }
+
+                ProxyTcpAcceptSelect::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);
+                    }
+                }
+
+                ProxyTcpAcceptSelect::Close => {
+                    break;
+                }
+            }
+        }
+
+        let accepted = if sockets
+            .get_mut::<tcp::Socket>(internal_socket_handle)
+            .is_active()
+        {
+            debug!("failed to accept tcp connection from client");
+            true
+        } else {
+            true
+        };
+
+        let mut already_shutdown = false;
+        let mut sleeper: Option<tokio::time::Sleep> = None;
+        loop {
+            if !accepted {
+                break;
+            }
+
+            let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
+
+            match socket.state() {
+                State::Closed
+                | State::Listen
+                | State::Closing
+                | State::LastAck
+                | State::TimeWait => {
+                    break;
+                }
+                State::FinWait1
+                | State::SynSent
+                | State::CloseWait
+                | State::FinWait2
+                | State::SynReceived
+                | State::Established => {}
+            }
+
+            let bytes_to_client = if socket.can_send() {
+                socket.send_capacity() - socket.send_queue()
+            } else {
+                0
+            };
+
+            let (bytes_to_external, do_shutdown) = if socket.may_recv() {
+                if let Ok(data) = socket.peek(TCP_BUFFER_SIZE) {
+                    if data.is_empty() {
+                        (None, false)
+                    } else {
+                        (Some(data), false)
+                    }
+                } else {
+                    (None, false)
+                }
+            } else if !already_shutdown && matches!(socket.state(), State::CloseWait) {
+                (None, true)
+            } else {
+                (None, false)
+            };
+            let selection = if let Some(sleep) = sleeper.take() {
+                if !do_shutdown {
+                    select! {
+                        biased;
+                        x = rx_receiver.recv() => if let Some(data) = x {
+                            ProxyTcpDataSelect::InternalRecv(data)
+                        } else {
+                            ProxyTcpDataSelect::Close
+                        },
+                        x = ip_receiver.recv() => if let Some(data) = x {
+                            ProxyTcpDataSelect::TxIpPacket(data)
+                        } else {
+                            ProxyTcpDataSelect::Close
+                        },
+                        x = external_w.write(bytes_to_external.unwrap_or(b"")), if bytes_to_external.is_some() => ProxyTcpDataSelect::ExternalSent(x?),
+                        x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
+                        _ = sleep => ProxyTcpDataSelect::TimePassed,
+                    }
+                } else {
+                    select! {
+                        biased;
+                        x = rx_receiver.recv() => if let Some(data) = x {
+                            ProxyTcpDataSelect::InternalRecv(data)
+                        } else {
+                            ProxyTcpDataSelect::Close
+                        },
+                        x = ip_receiver.recv() => if let Some(data) = x {
+                            ProxyTcpDataSelect::TxIpPacket(data)
+                        } else {
+                            ProxyTcpDataSelect::Close
+                        },
+                        _ = external_w.shutdown() => ProxyTcpDataSelect::ExternalSent(0),
+                        x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
+                        _ = sleep => ProxyTcpDataSelect::TimePassed,
+                    }
+                }
+            } else if !do_shutdown {
+                select! {
+                    biased;
+                    x = rx_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpDataSelect::InternalRecv(data)
+                    } else {
+                        ProxyTcpDataSelect::Close
+                    },
+                    x = ip_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpDataSelect::TxIpPacket(data)
+                    } else {
+                        ProxyTcpDataSelect::Close
+                    },
+                    x = external_w.write(bytes_to_external.unwrap_or(b"")), if bytes_to_external.is_some() => ProxyTcpDataSelect::ExternalSent(x?),
+                    x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
+                    _ = std::future::ready(()) => ProxyTcpDataSelect::DoNothing,
+                }
+            } else {
+                select! {
+                    biased;
+                    x = rx_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpDataSelect::InternalRecv(data)
+                    } else {
+                        ProxyTcpDataSelect::Close
+                    },
+                    x = ip_receiver.recv() => if let Some(data) = x {
+                        ProxyTcpDataSelect::TxIpPacket(data)
+                    } else {
+                        ProxyTcpDataSelect::Close
+                    },
+                    _ = external_w.shutdown() => ProxyTcpDataSelect::ExternalSent(0),
+                    x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
+                    _ = std::future::ready(()) => ProxyTcpDataSelect::DoNothing,
+                }
+            };
+            match selection {
+                ProxyTcpDataSelect::ExternalRecv(size) => {
+                    if size == 0 {
+                        socket.close();
+                    } else {
+                        socket.send_slice(&external_buffer[..size])?;
+                    }
+                }
+
+                ProxyTcpDataSelect::ExternalSent(size) => {
+                    if size == 0 {
+                        already_shutdown = true;
+                    } else {
+                        socket.recv(|_| (size, ()))?;
+                    }
+                }
+
+                ProxyTcpDataSelect::InternalRecv(data) => {
+                    let (_, payload) = Ethernet2Header::from_slice(&data)?;
+                    device.rx = Some(payload.to_vec());
+                    iface.poll(Instant::now(), &mut device, &mut sockets);
+                }
+
+                ProxyTcpDataSelect::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);
+                    }
+                }
+
+                ProxyTcpDataSelect::TimePassed => {
+                    iface.poll(Instant::now(), &mut device, &mut sockets);
+                }
+
+                ProxyTcpDataSelect::DoNothing => {
+                    sleeper = Some(tokio::time::sleep(Duration::from_millis(50)));
+                }
+
+                ProxyTcpDataSelect::Close => {
+                    break;
+                }
+            }
+        }
+
+        context.reclaim().await?;
+
+        Ok(())
+    }
+}

network/src/proxynat/udp.rs

@@ -128,6 +128,8 @@ impl ProxyUdpHandler {
             }
         }
 
+        context.reclaim().await?;
+
         Ok(())
     }
 }