support no tun mode (#141)

easytier/src/gateway/tcp_proxy.rs

@@ -1,3 +1,4 @@
+use core::panic;
 use crossbeam::atomic::AtomicCell;
 use dashmap::DashMap;
 use pnet::packet::ip::IpNextHeaderProtocols;
@@ -6,19 +7,18 @@ use pnet::packet::tcp::{ipv4_checksum, MutableTcpPacket, TcpPacket};
 use pnet::packet::MutablePacket;
 use pnet::packet::Packet;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
-use std::sync::atomic::AtomicU16;
+use std::sync::atomic::{AtomicBool, AtomicU16};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::io::copy_bidirectional;
 use tokio::net::{TcpListener, TcpSocket, TcpStream};
-use tokio::sync::Mutex;
+use tokio::sync::{mpsc, Mutex};
 use tokio::task::JoinSet;
 use tracing::Instrument;
 
 use crate::common::error::Result;
-use crate::common::global_ctx::GlobalCtx;
+use crate::common::global_ctx::{ArcGlobalCtx, GlobalCtx};
 use crate::common::join_joinset_background;
-use crate::common::netns::NetNS;
 
 use crate::peers::peer_manager::PeerManager;
 use crate::peers::{NicPacketFilter, PeerPacketFilter};
@@ -26,6 +26,9 @@ use crate::tunnel::packet_def::{PacketType, ZCPacket};
 
 use super::CidrSet;
 
+#[cfg(feature = "smoltcp")]
+use super::tokio_smoltcp::{self, channel_device, Net, NetConfig};
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 enum NatDstEntryState {
     // receive syn packet but not start connecting to dst
@@ -61,6 +64,64 @@ impl NatDstEntry {
     }
 }
 
+enum ProxyTcpStream {
+    KernelTcpStream(TcpStream),
+    #[cfg(feature = "smoltcp")]
+    SmolTcpStream(tokio_smoltcp::TcpStream),
+}
+
+impl ProxyTcpStream {
+    pub fn set_nodelay(&self, nodelay: bool) -> Result<()> {
+        match self {
+            Self::KernelTcpStream(stream) => stream.set_nodelay(nodelay).map_err(Into::into),
+            #[cfg(feature = "smoltcp")]
+            Self::SmolTcpStream(_stream) => {
+                tracing::warn!("smol tcp stream set_nodelay not implemented");
+                Ok(())
+            }
+        }
+    }
+
+    pub async fn copy_bidirectional(&mut self, dst: &mut TcpStream) -> Result<()> {
+        match self {
+            Self::KernelTcpStream(stream) => {
+                copy_bidirectional(stream, dst).await?;
+                Ok(())
+            }
+            #[cfg(feature = "smoltcp")]
+            Self::SmolTcpStream(stream) => {
+                copy_bidirectional(stream, dst).await?;
+                Ok(())
+            }
+        }
+    }
+}
+
+enum ProxyTcpListener {
+    KernelTcpListener(TcpListener),
+    #[cfg(feature = "smoltcp")]
+    SmolTcpListener(tokio_smoltcp::TcpListener),
+}
+
+impl ProxyTcpListener {
+    pub async fn accept(&mut self) -> Result<(ProxyTcpStream, SocketAddr)> {
+        match self {
+            Self::KernelTcpListener(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok((ProxyTcpStream::KernelTcpStream(stream), addr))
+            }
+            #[cfg(feature = "smoltcp")]
+            Self::SmolTcpListener(listener) => {
+                let Ok((stream, src)) = listener.accept().await else {
+                    return Err(anyhow::anyhow!("smol tcp listener closed").into());
+                };
+                tracing::info!(?src, "smol tcp listener accepted");
+                Ok((ProxyTcpStream::SmolTcpStream(stream), src))
+            }
+        }
+    }
+}
+
 type ArcNatDstEntry = Arc<NatDstEntry>;
 
 type SynSockMap = Arc<DashMap<SocketAddr, ArcNatDstEntry>>;
@@ -81,6 +142,12 @@ pub struct TcpProxy {
     addr_conn_map: AddrConnSockMap,
 
     cidr_set: CidrSet,
+
+    smoltcp_stack_sender: Option<mpsc::Sender<ZCPacket>>,
+    smoltcp_stack_receiver: Arc<Mutex<Option<mpsc::Receiver<ZCPacket>>>>,
+    #[cfg(feature = "smoltcp")]
+    smoltcp_net: Arc<Mutex<Option<Net>>>,
+    enable_smoltcp: Arc<AtomicBool>,
 }
 
 #[async_trait::async_trait]
@@ -157,6 +224,8 @@ impl NicPacketFilter for TcpProxy {
 
 impl TcpProxy {
     pub fn new(global_ctx: Arc<GlobalCtx>, peer_manager: Arc<PeerManager>) -> Arc<Self> {
+        let (smoltcp_stack_sender, smoltcp_stack_receiver) = mpsc::channel::<ZCPacket>(1000);
+
         Arc::new(Self {
             global_ctx: global_ctx.clone(),
             peer_manager,
@@ -169,6 +238,14 @@ impl TcpProxy {
             addr_conn_map: Arc::new(DashMap::new()),
 
             cidr_set: CidrSet::new(global_ctx),
+
+            smoltcp_stack_sender: Some(smoltcp_stack_sender),
+            smoltcp_stack_receiver: Arc::new(Mutex::new(Some(smoltcp_stack_receiver))),
+
+            #[cfg(feature = "smoltcp")]
+            smoltcp_net: Arc::new(Mutex::new(None)),
+
+            enable_smoltcp: Arc::new(AtomicBool::new(true)),
         })
     }
 
@@ -224,34 +301,114 @@ impl TcpProxy {
         Ok(())
     }
 
+    async fn get_proxy_listener(&self) -> Result<ProxyTcpListener> {
+        #[cfg(feature = "smoltcp")]
+        if self.global_ctx.get_flags().use_smoltcp || self.global_ctx.no_tun() {
+            // use smoltcp network stack
+            self.local_port
+                .store(8899, std::sync::atomic::Ordering::Relaxed);
+
+            let mut cap = smoltcp::phy::DeviceCapabilities::default();
+            cap.max_transmission_unit = 1280;
+            cap.medium = smoltcp::phy::Medium::Ip;
+            let (dev, stack_sink, mut stack_stream) = channel_device::ChannelDevice::new(cap);
+
+            let mut smoltcp_stack_receiver =
+                self.smoltcp_stack_receiver.lock().await.take().unwrap();
+            self.tasks.lock().unwrap().spawn(async move {
+                while let Some(packet) = smoltcp_stack_receiver.recv().await {
+                    tracing::trace!(?packet, "receive from peer send to smoltcp packet");
+                    if let Err(e) = stack_sink.send(Ok(packet.payload().to_vec())).await {
+                        tracing::error!("send to smoltcp stack failed: {:?}", e);
+                    }
+                }
+                tracing::error!("smoltcp stack sink exited");
+                panic!("smoltcp stack sink exited");
+            });
+
+            let peer_mgr = self.peer_manager.clone();
+            self.tasks.lock().unwrap().spawn(async move {
+                while let Some(data) = stack_stream.recv().await {
+                    tracing::trace!(
+                        ?data,
+                        "receive from smoltcp stack and send to peer mgr packet"
+                    );
+                    let Some(ipv4) = Ipv4Packet::new(&data) else {
+                        tracing::error!(?data, "smoltcp stack stream get non ipv4 packet");
+                        continue;
+                    };
+
+                    let dst = ipv4.get_destination();
+                    let packet = ZCPacket::new_with_payload(&data);
+                    if let Err(e) = peer_mgr.send_msg_ipv4(packet, dst).await {
+                        tracing::error!("send to peer failed in smoltcp sender: {:?}", e);
+                    }
+                }
+                tracing::error!("smoltcp stack stream exited");
+                panic!("smoltcp stack stream exited");
+            });
+
+            let interface_config = smoltcp::iface::Config::new(smoltcp::wire::HardwareAddress::Ip);
+            let net = Net::new(
+                dev,
+                NetConfig::new(
+                    interface_config,
+                    format!("{}/24", self.global_ctx.get_ipv4().unwrap())
+                        .parse()
+                        .unwrap(),
+                    vec![],
+                ),
+            );
+            net.set_any_ip(true);
+            let tcp = net.tcp_bind("0.0.0.0:8899".parse().unwrap()).await?;
+            self.smoltcp_net.lock().await.replace(net);
+
+            self.enable_smoltcp
+                .store(true, std::sync::atomic::Ordering::Relaxed);
+
+            return Ok(ProxyTcpListener::SmolTcpListener(tcp));
+        }
+
+        {
+            // use kernel network stack
+            let listen_addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0);
+            let net_ns = self.global_ctx.net_ns.clone();
+            let tcp_listener = net_ns
+                .run_async(|| async { TcpListener::bind(&listen_addr).await })
+                .await?;
+            self.local_port.store(
+                tcp_listener.local_addr()?.port(),
+                std::sync::atomic::Ordering::Relaxed,
+            );
+
+            self.enable_smoltcp
+                .store(false, std::sync::atomic::Ordering::Relaxed);
+
+            return Ok(ProxyTcpListener::KernelTcpListener(tcp_listener));
+        }
+    }
+
     async fn run_listener(&self) -> Result<()> {
         // bind on both v4 & v6
-        let listen_addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0);
-
-        let net_ns = self.global_ctx.net_ns.clone();
-        let tcp_listener = net_ns
-            .run_async(|| async { TcpListener::bind(&listen_addr).await })
-            .await?;
-
-        self.local_port.store(
-            tcp_listener.local_addr()?.port(),
-            std::sync::atomic::Ordering::Relaxed,
-        );
+        let mut tcp_listener = self.get_proxy_listener().await?;
 
+        let global_ctx = self.global_ctx.clone();
         let tasks = self.tasks.clone();
         let syn_map = self.syn_map.clone();
         let conn_map = self.conn_map.clone();
         let addr_conn_map = self.addr_conn_map.clone();
         let accept_task = async move {
-            tracing::info!(listener = ?tcp_listener, "tcp connection start accepting");
-
             let conn_map = conn_map.clone();
             while let Ok((tcp_stream, socket_addr)) = tcp_listener.accept().await {
                 let Some(entry) = syn_map.get(&socket_addr) else {
                     tracing::error!("tcp connection from unknown source: {:?}", socket_addr);
                     continue;
                 };
-                tracing::info!(?socket_addr, "tcp connection accepted for proxy");
+                tracing::info!(
+                    ?socket_addr,
+                    "tcp connection accepted for proxy, nat dst: {:?}",
+                    entry.dst
+                );
                 assert_eq!(entry.state.load(), NatDstEntryState::SynReceived);
 
                 let entry_clone = entry.clone();
@@ -265,7 +422,7 @@ impl TcpProxy {
                 assert!(old_nat_val.is_none());
 
                 tasks.lock().unwrap().spawn(Self::connect_to_nat_dst(
-                    net_ns.clone(),
+                    global_ctx.clone(),
                     tcp_stream,
                     conn_map.clone(),
                     addr_conn_map.clone(),
@@ -293,8 +450,8 @@ impl TcpProxy {
     }
 
     async fn connect_to_nat_dst(
-        net_ns: NetNS,
-        src_tcp_stream: TcpStream,
+        global_ctx: ArcGlobalCtx,
+        src_tcp_stream: ProxyTcpStream,
         conn_map: ConnSockMap,
         addr_conn_map: AddrConnSockMap,
         nat_entry: ArcNatDstEntry,
@@ -303,14 +460,24 @@ impl TcpProxy {
             tracing::warn!("set_nodelay failed, ignore it: {:?}", e);
         }
 
-        let _guard = net_ns.guard();
+        let _guard = global_ctx.net_ns.guard();
         let socket = TcpSocket::new_v4().unwrap();
         if let Err(e) = socket.set_nodelay(true) {
             tracing::warn!("set_nodelay failed, ignore it: {:?}", e);
         }
+
+        let nat_dst = if Some(nat_entry.dst.ip()) == global_ctx.get_ipv4().map(|ip| IpAddr::V4(ip))
+        {
+            format!("127.0.0.1:{}", nat_entry.dst.port())
+                .parse()
+                .unwrap()
+        } else {
+            nat_entry.dst
+        };
+
         let Ok(Ok(dst_tcp_stream)) = tokio::time::timeout(
             Duration::from_secs(10),
-            TcpSocket::new_v4().unwrap().connect(nat_entry.dst),
+            TcpSocket::new_v4().unwrap().connect(nat_dst),
         )
         .await
         else {
@@ -321,6 +488,8 @@ impl TcpProxy {
         };
         drop(_guard);
 
+        tracing::info!(?nat_entry, ?nat_dst, "tcp connection to dst established");
+
         assert_eq!(nat_entry.state.load(), NatDstEntryState::ConnectingDst);
         nat_entry.state.store(NatDstEntryState::Connected);
 
@@ -335,7 +504,7 @@ impl TcpProxy {
     }
 
     async fn handle_nat_connection(
-        mut src_tcp_stream: TcpStream,
+        mut src_tcp_stream: ProxyTcpStream,
         mut dst_tcp_stream: TcpStream,
         conn_map: ConnSockMap,
         addr_conn_map: AddrConnSockMap,
@@ -343,8 +512,8 @@ impl TcpProxy {
     ) {
         let nat_entry_clone = nat_entry.clone();
         nat_entry.tasks.lock().await.spawn(async move {
-            let ret = copy_bidirectional(&mut src_tcp_stream, &mut dst_tcp_stream).await;
-            tracing::trace!(nat_entry = ?nat_entry_clone, ret = ?ret, "nat tcp connection closed");
+            let ret = src_tcp_stream.copy_bidirectional(&mut dst_tcp_stream).await;
+            tracing::info!(nat_entry = ?nat_entry_clone, ret = ?ret, "nat tcp connection closed");
             nat_entry_clone.state.store(NatDstEntryState::Closed);
 
             Self::remove_entry_from_all_conn_map(conn_map, addr_conn_map, nat_entry_clone);
@@ -356,7 +525,10 @@ impl TcpProxy {
     }
 
     async fn try_handle_peer_packet(&self, packet: &mut ZCPacket) -> Option<()> {
-        if self.cidr_set.is_empty() && !self.global_ctx.enable_exit_node() {
+        if self.cidr_set.is_empty()
+            && !self.global_ctx.enable_exit_node()
+            && !self.global_ctx.no_tun()
+        {
             return None;
         }
 
@@ -375,11 +547,15 @@ impl TcpProxy {
             return None;
         }
 
-        if !self.cidr_set.contains_v4(ipv4.get_destination()) && !is_exit_node {
+        if !self.cidr_set.contains_v4(ipv4.get_destination())
+            && !is_exit_node
+            && !(self.global_ctx.no_tun()
+                && Some(ipv4.get_destination()) == self.global_ctx.get_ipv4())
+        {
             return None;
         }
 
-        tracing::info!(ipv4 = ?ipv4, cidr_set = ?self.cidr_set, "proxy tcp packet received");
+        tracing::trace!(ipv4 = ?ipv4, cidr_set = ?self.cidr_set, "proxy tcp packet received");
 
         let ip_packet = Ipv4Packet::new(payload_bytes).unwrap();
         let tcp_packet = TcpPacket::new(ip_packet.payload()).unwrap();
@@ -397,7 +573,7 @@ impl TcpProxy {
             let old_val = self
                 .syn_map
                 .insert(src, Arc::new(NatDstEntry::new(src, dst)));
-            tracing::trace!(src = ?src, dst = ?dst, old_entry = ?old_val, "tcp syn received");
+            tracing::info!(src = ?src, dst = ?dst, old_entry = ?old_val, "tcp syn received");
         }
 
         let mut ip_packet = MutableIpv4Packet::new(payload_bytes).unwrap();
@@ -411,7 +587,18 @@ impl TcpProxy {
         drop(tcp_packet);
         Self::update_ip_packet_checksum(&mut ip_packet);
 
-        tracing::info!(?source, ?ipv4_addr, ?packet, "tcp packet after modified");
+        tracing::trace!(?source, ?ipv4_addr, ?packet, "tcp packet after modified");
+
+        if self
+            .enable_smoltcp
+            .load(std::sync::atomic::Ordering::Relaxed)
+        {
+            let smoltcp_stack_sender = self.smoltcp_stack_sender.as_ref().unwrap();
+            if let Err(e) = smoltcp_stack_sender.try_send(packet.clone()) {
+                tracing::error!("send to smoltcp stack failed: {:?}", e);
+            }
+            return None;
+        }
 
         Some(())
     }