move bind_tcp_socket and bind_udp_socket to common

easytier/src/gateway/socks5.rs

@@ -18,7 +18,7 @@ use crate::gateway::kcp_proxy::NatDstKcpConnector;
 use crate::{
     common::{
         config::PortForwardConfig, global_ctx::GlobalCtxEvent, join_joinset_background,
-        netns::NetNS, scoped_task::ScopedTask,
+        scoped_task::ScopedTask,
     },
     gateway::{
         fast_socks5::{
@@ -30,10 +30,7 @@ use crate::{
         ip_reassembler::IpReassembler,
         tokio_smoltcp::{channel_device, BufferSize, Net, NetConfig},
     },
-    tunnel::{
-        common::setup_sokcet2,
-        packet_def::{PacketType, ZCPacket},
-    },
+    tunnel::packet_def::{PacketType, ZCPacket},
 };
 use anyhow::Context;
 use dashmap::DashMap;
@@ -42,21 +39,21 @@ use pnet::packet::{
 };
 use tokio::{
     io::{AsyncRead, AsyncWrite},
-    net::{TcpListener, TcpSocket, UdpSocket},
+    net::{TcpListener, UdpSocket},
     select,
     sync::{mpsc, Mutex, Notify},
     task::JoinSet,
     time::timeout,
 };
 
+#[cfg(feature = "kcp")]
+use super::tcp_proxy::NatDstConnector as _;
+use crate::tunnel::common::{bind_tcp_socket, bind_udp_socket};
 use crate::{
     common::{error::Error, global_ctx::GlobalCtx},
     peers::{peer_manager::PeerManager, PeerPacketFilter},
 };
 
-#[cfg(feature = "kcp")]
-use super::tcp_proxy::NatDstConnector as _;
-
 enum SocksUdpSocket {
     UdpSocket(Arc<tokio::net::UdpSocket>),
     SmolUdpSocket(super::tokio_smoltcp::UdpSocket),
@@ -332,38 +329,6 @@ impl AsyncTcpConnector for Socks5AutoConnector {
     }
 }
 
-fn bind_tcp_socket(addr: SocketAddr, net_ns: NetNS) -> Result<TcpListener, Error> {
-    let _g = net_ns.guard();
-    let socket2_socket = socket2::Socket::new(
-        socket2::Domain::for_address(addr),
-        socket2::Type::STREAM,
-        Some(socket2::Protocol::TCP),
-    )?;
-
-    setup_sokcet2(&socket2_socket, &addr)?;
-
-    let socket = TcpSocket::from_std_stream(socket2_socket.into());
-
-    if let Err(e) = socket.set_nodelay(true) {
-        tracing::warn!(?e, "set_nodelay fail in listen");
-    }
-
-    Ok(socket.listen(1024)?)
-}
-
-fn bind_udp_socket(addr: SocketAddr, net_ns: NetNS) -> Result<UdpSocket, Error> {
-    let _g = net_ns.guard();
-    let socket2_socket = socket2::Socket::new(
-        socket2::Domain::for_address(addr),
-        socket2::Type::DGRAM,
-        Some(socket2::Protocol::UDP),
-    )?;
-
-    setup_sokcet2(&socket2_socket, &addr)?;
-
-    Ok(UdpSocket::from_std(socket2_socket.into())?)
-}
-
 struct Socks5ServerNet {
     ipv4_addr: cidr::Ipv4Inet,
     auth: Option<SimpleUserPassword>,

easytier/src/tunnel/common.rs

@@ -11,13 +11,15 @@ use network_interface::NetworkInterfaceConfig as _;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
+use super::TunnelInfo;
+use crate::common::error::Error;
+use crate::common::netns::NetNS;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
+use tokio::net::{TcpListener, TcpSocket, UdpSocket};
 use tokio_stream::StreamExt;
 use tokio_util::io::poll_write_buf;
 use zerocopy::FromBytes as _;
 
-use super::TunnelInfo;
-
 use crate::tunnel::packet_def::{ZCPacket, PEER_MANAGER_HEADER_SIZE};
 
 use super::{
@@ -439,6 +441,56 @@ pub(crate) fn setup_sokcet2(
     )
 }
 
+pub trait Bindable: Sized {
+    const TY: socket2::Type;
+    const PROTOCOL: Option<socket2::Protocol>;
+
+    fn finalize(socket: socket2::Socket) -> Result<Self, Error>;
+}
+
+impl Bindable for TcpListener {
+    const TY: socket2::Type = socket2::Type::STREAM;
+    const PROTOCOL: Option<socket2::Protocol> = Some(socket2::Protocol::TCP);
+
+    fn finalize(socket: socket2::Socket) -> Result<Self, Error> {
+        let tcp_socket = TcpSocket::from_std_stream(socket.into());
+
+        if let Err(e) = tcp_socket.set_nodelay(true) {
+            tracing::warn!(?e, "set_nodelay fail in listen");
+        }
+
+        Ok(tcp_socket.listen(1024)?)
+    }
+}
+
+impl Bindable for UdpSocket {
+    const TY: socket2::Type = socket2::Type::DGRAM;
+    const PROTOCOL: Option<socket2::Protocol> = Some(socket2::Protocol::UDP);
+
+    fn finalize(socket: socket2::Socket) -> Result<Self, Error> {
+        Ok(UdpSocket::from_std(socket.into())?)
+    }
+}
+
+pub fn bind_socket<B: Bindable>(addr: SocketAddr, net_ns: Option<NetNS>) -> Result<B, Error> {
+    let _g = net_ns.map(|n| n.guard());
+
+    let socket2_socket =
+        socket2::Socket::new(socket2::Domain::for_address(addr), B::TY, B::PROTOCOL)?;
+
+    setup_sokcet2(&socket2_socket, &addr)?;
+
+    B::finalize(socket2_socket)
+}
+
+pub fn bind_tcp_socket(addr: SocketAddr, net_ns: NetNS) -> Result<TcpListener, Error> {
+    bind_socket(addr, Some(net_ns))
+}
+
+pub fn bind_udp_socket(addr: SocketAddr, net_ns: NetNS) -> Result<UdpSocket, Error> {
+    bind_socket(addr, Some(net_ns))
+}
+
 pub fn reserve_buf(buf: &mut BytesMut, min_size: usize, max_size: usize) {
     if buf.capacity() < min_size {
         buf.reserve(max_size);