wrap tcp listener

Author: VendettaReborn

examples/forward.rs

@@ -1,12 +1,9 @@
-use std::{
-    net::{Ipv4Addr, SocketAddr},
-    sync::Arc,
-};
+use std::net::{Ipv4Addr, SocketAddr};
 
 use futures::{SinkExt, StreamExt};
-use netstack_system::{Nat, UdpSocket};
+use netstack_system::UdpSocket;
 use structopt::StructOpt;
-use tokio::net::{TcpListener, TcpSocket, TcpStream};
+use tokio::net::{TcpSocket, TcpStream};
 use tracing::{error, info, warn};
 use tun2::AbstractDevice;
 
@@ -99,7 +96,6 @@ async fn main_exec(opt: Opt) {
         .inet4_addr(addr)
         .build()
         .await;
-    let nat = stack.nat();
     let listener = listener.unwrap();
     let udp_socket = udp_socket.unwrap();
 
@@ -117,7 +113,7 @@ async fn main_exec(opt: Opt) {
     futs.push(tokio_spawn!({
         let interface = opt.interface.clone();
         async move {
-            handle_inbound_stream(nat, listener, interface).await;
+            handle_inbound_stream(listener, interface).await;
         }
     }));
     futs.push(tokio_spawn!({
@@ -137,25 +133,19 @@ async fn main_exec(opt: Opt) {
 }
 
 /// simply forward tcp stream
-async fn handle_inbound_stream(nat: Arc<Nat>, listener: TcpListener, interface: String) {
-    while let Ok((mut stream, _addr)) = listener.accept().await {
+async fn handle_inbound_stream(listener: netstack_system::TcpListener, interface: String) {
+    while let Ok((mut stream, target)) = listener.accept().await {
         let interface: String = interface.clone();
-        let nat = nat.clone();
+        // let nat = nat.clone();
         tokio::spawn(async move {
             let remote = stream.peer_addr().unwrap();
             let local = stream.local_addr().unwrap();
-            info!("new tcp connection: {:?} => {:?}", remote, local);
-            let nat_port = remote.port();
-            let nat_session = match nat.look_back(nat_port).await {
-                Some(session) => session,
-                None => {
-                    warn!("session not found, allocated port: {:?}", remote.port());
-                    return;
-                }
-            };
-            let actual_remote = (nat_session.dst, nat_session.dport);
+            info!(
+                "new tcp connection: {:?} => {:?}, actual remote: {}",
+                remote, local, target
+            );
 
-            match new_tcp_stream(actual_remote.into(), &interface).await {
+            match new_tcp_stream(target.into(), &interface).await {
                 Ok(mut remote_stream) => {
                     // pipe between two tcp stream
                     match tokio::io::copy_bidirectional(&mut stream, &mut remote_stream).await {

src/lib.rs

@@ -4,11 +4,13 @@ mod checksum;
 mod error;
 mod nat;
 mod packet;
+mod tcp;
 mod udp;
 
 use checksum::UpdateCsum;
 pub use error::StackError;
 pub use nat::*;
+pub use tcp::TcpListener;
 pub use udp::UdpSocket;
 pub type Result<T> = std::result::Result<T, StackError>;
 
@@ -19,7 +21,6 @@ use std::{
     sync::Arc,
 };
 use tokio::{
-    net::TcpListener,
     sync::mpsc::{Receiver, Sender, channel},
     task::JoinHandle,
 };
@@ -72,7 +73,7 @@ impl StackBuilder {
         self,
     ) -> (
         SystemStack,
-        Option<TcpListener>,
+        Option<crate::TcpListener>,
         Option<UdpSocket>,
         Sender<Vec<u8>>,
     ) {
@@ -88,7 +89,9 @@ impl StackBuilder {
         };
 
         let tcp_listener = if self.enable_tcp {
-            TcpListener::bind((self.inet4_server_addr, 0)).await.ok()
+            tokio::net::TcpListener::bind((self.inet4_server_addr, 0))
+                .await
+                .ok()
         } else {
             None
         };
@@ -109,6 +112,8 @@ impl StackBuilder {
 
         let udp_socket = udp_rx.map(|rx| UdpSocket::new(rx, udp_writeback_tx));
 
+        let tcp_listener = tcp_listener.map(|l| crate::TcpListener::new(l, stack.nat()));
+
         (stack, tcp_listener, udp_socket, stack_tx)
     }
 }

src/tcp.rs

@@ -0,0 +1,35 @@
+use std::{net::SocketAddr, sync::Arc};
+
+use crate::Nat;
+
+pub struct TcpListener {
+    listener: tokio::net::TcpListener,
+    nat: Arc<Nat>,
+}
+
+impl TcpListener {
+    pub fn new(listener: tokio::net::TcpListener, nat: Arc<Nat>) -> Self {
+        Self { listener, nat }
+    }
+
+    pub fn inner(&self) -> &tokio::net::TcpListener {
+        &self.listener
+    }
+
+    pub async fn accept(&self) -> std::io::Result<(tokio::net::TcpStream, std::net::SocketAddr)> {
+        let (stream, remote) = self.listener.accept().await?;
+        let remote_port = remote.port();
+        let nat_session = match self.nat.look_back(remote_port).await {
+            Some(session) => session,
+            None => {
+                tracing::warn!("session not found, allocated port: {:?}", remote.port());
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::Other,
+                    "session not found",
+                ));
+            }
+        };
+        let actual_remote = SocketAddr::new(nat_session.dst, nat_session.dport);
+        return Ok((stream, actual_remote));
+    }
+}