From 3a965efab2bdc70c17f3a03744b4e0412eb1432e Mon Sep 17 00:00:00 2001
From: "sijie.sun" <sijie.sun@smartx.com>
Date: Sat, 27 Apr 2024 21:20:23 +0800
Subject: [PATCH] allow tunnel listener alloc port after listen

---
 easytier/src/tunnel/quic.rs      | 20 ++++++++++++++++++++
 easytier/src/tunnel/tcp.rs       | 19 +++++++++++++++++++
 easytier/src/tunnel/udp.rs       | 19 +++++++++++++++++++
 easytier/src/tunnel/wireguard.rs | 21 +++++++++++++++++++++
 4 files changed, 79 insertions(+)

diff --git a/easytier/src/tunnel/quic.rs b/easytier/src/tunnel/quic.rs
index 49dd455..9ae7cd5 100644
--- a/easytier/src/tunnel/quic.rs
+++ b/easytier/src/tunnel/quic.rs
@@ -120,6 +120,11 @@ impl TunnelListener for QUICTunnelListener {
         let (endpoint, server_cert) = make_server_endpoint(addr).unwrap();
         self.endpoint = Some(endpoint);
         self.server_cert = Some(server_cert);
+
+        self.addr
+            .set_port(Some(self.endpoint.as_ref().unwrap().local_addr()?.port()))
+            .unwrap();
+
         Ok(())
     }
 
@@ -267,4 +272,19 @@ mod tests {
         connector.set_ip_version(IpVersion::V4);
         _tunnel_pingpong(listener, connector).await;
     }
+
+    #[tokio::test]
+    async fn test_alloc_port() {
+        // v4
+        let mut listener = QUICTunnelListener::new("quic://0.0.0.0:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+
+        // v6
+        let mut listener = QUICTunnelListener::new("quic://[::]:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+    }
 }
diff --git a/easytier/src/tunnel/tcp.rs b/easytier/src/tunnel/tcp.rs
index 8dddbaa..0b334d6 100644
--- a/easytier/src/tunnel/tcp.rs
+++ b/easytier/src/tunnel/tcp.rs
@@ -45,6 +45,10 @@ impl TunnelListener for TcpTunnelListener {
         // socket.set_reuseport(true)?;
         socket.bind(addr)?;
 
+        self.addr
+            .set_port(Some(socket.local_addr()?.port()))
+            .unwrap();
+
         self.listener = Some(socket.listen(1024)?);
         Ok(())
     }
@@ -232,4 +236,19 @@ mod tests {
         connector.set_ip_version(IpVersion::V4);
         _tunnel_pingpong(listener, connector).await;
     }
+
+    #[tokio::test]
+    async fn test_alloc_port() {
+        // v4
+        let mut listener = TcpTunnelListener::new("tcp://0.0.0.0:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+
+        // v6
+        let mut listener = TcpTunnelListener::new("tcp://[::]:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+    }
 }
diff --git a/easytier/src/tunnel/udp.rs b/easytier/src/tunnel/udp.rs
index b12ab41..d4d7c21 100644
--- a/easytier/src/tunnel/udp.rs
+++ b/easytier/src/tunnel/udp.rs
@@ -388,6 +388,10 @@ impl TunnelListener for UdpTunnelListener {
         self.socket = Some(Arc::new(UdpSocket::from_std(socket2_socket.into())?));
         self.data.socket = self.socket.clone();
 
+        self.addr
+            .set_port(Some(self.socket.as_ref().unwrap().local_addr()?.port()))
+            .unwrap();
+
         self.forward_tasks
             .lock()
             .unwrap()
@@ -884,4 +888,19 @@ mod tests {
         connector.set_ip_version(IpVersion::V4);
         _tunnel_pingpong(listener, connector).await;
     }
+
+    #[tokio::test]
+    async fn test_alloc_port() {
+        // v4
+        let mut listener = UdpTunnelListener::new("udp://0.0.0.0:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+
+        // v6
+        let mut listener = UdpTunnelListener::new("udp://[::]:0".parse().unwrap());
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+    }
 }
diff --git a/easytier/src/tunnel/wireguard.rs b/easytier/src/tunnel/wireguard.rs
index ffc0475..b1cd804 100644
--- a/easytier/src/tunnel/wireguard.rs
+++ b/easytier/src/tunnel/wireguard.rs
@@ -537,6 +537,10 @@ impl TunnelListener for WgTunnelListener {
         }
 
         self.udp = Some(Arc::new(UdpSocket::from_std(socket2_socket.into())?));
+        self.addr
+            .set_port(Some(self.udp.as_ref().unwrap().local_addr()?.port()))
+            .unwrap();
+
         self.tasks.spawn(Self::handle_udp_incoming(
             self.get_udp_socket(),
             self.config.clone(),
@@ -868,4 +872,21 @@ pub mod tests {
         connector.set_ip_version(IpVersion::V4);
         _tunnel_pingpong(listener, connector).await;
     }
+
+    #[tokio::test]
+    async fn test_alloc_port() {
+        // v4
+        let (server_cfg, _client_cfg) = create_wg_config();
+        let mut listener = WgTunnelListener::new("wg://0.0.0.0:0".parse().unwrap(), server_cfg);
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+
+        // v6
+        let (server_cfg, _client_cfg) = create_wg_config();
+        let mut listener = WgTunnelListener::new("wg://[::]:0".parse().unwrap(), server_cfg);
+        listener.listen().await.unwrap();
+        let port = listener.local_url().port().unwrap();
+        assert!(port > 0);
+    }
 }