refactor(server): port range and exclusive port

Author: sword-jin

Cargo.lock

@@ -38,6 +38,8 @@ fastrand = "2.1.0"
 prost = "0.12.6"
 async-shutdown = "0.2.2"
 httparse = "1.9.4"
+rand = "0.8.5"
+rand_chacha = "0.3.1"
 
 [build-dependencies]
 tonic-build = "0.11.0"

Cargo.toml

@@ -40,6 +40,9 @@ struct Args {
     /// Maximum accepted port number.
     #[clap(long, default_value_t = 65535)]
     random_max_port: u16,
+
+    #[clap(long, default_value = "[]")]
+    exclude_ports: Vec<u16>,
 }
 
 #[tokio::main]
@@ -71,6 +74,7 @@ async fn main() {
             ip: args.ip,
             vhttp_behind_proxy_tls: args.vhttp_behind_proxy_tls,
             port_range: args.random_min_port..=args.random_max_port,
+            exclude_ports: args.exclude_ports,
         },
     });
     if let Err(err) = server.run(cancel.cancelled()).await {

src/bin/castled.rs

@@ -1,4 +1,7 @@
-use crate::event::{self, ClientEventResponse, Payload};
+use crate::{
+    event::{self, ClientEventResponse, Payload},
+    server::port::{Available, PortManager},
+};
 
 use super::{
     tunnel::{
@@ -29,14 +32,20 @@ pub(crate) struct DataServer {
     http_tunnel: Http,
     http_registry: DynamicRegistry,
     entrypoint_config: EntrypointConfig,
+    port_manager: PortManager,
 }
 
 impl DataServer {
     pub(crate) fn new(vhttp_port: u16, entrypoint_config: EntrypointConfig) -> Self {
         let http_registry = DynamicRegistry::new();
+        let port_manager = PortManager::new(
+            entrypoint_config.port_range.clone(),
+            entrypoint_config.exclude_ports.clone(),
+        );
         Self {
             http_registry: http_registry.clone(),
             http_tunnel: Http::new(vhttp_port, Arc::new(Box::new(http_registry))),
+            port_manager,
             entrypoint_config,
         }
     }
@@ -62,8 +71,8 @@ impl DataServer {
         while let Some(event) = receiver.recv().await {
             match event.payload {
                 event::Payload::RegisterTcp { port } => {
-                    let result: Result<(u16, TcpListener), tonic::Status> =
-                        create_socket::<Tcp>(port, this.entrypoint_config.port_range.clone()).await;
+                    let result: Result<(Available, TcpListener), tonic::Status> =
+                        create_socket::<Tcp>(port, &mut self.port_manager).await;
                     match result {
                         Ok((port, listener)) => {
                             let cancel = event.close_listener;
@@ -72,12 +81,13 @@ impl DataServer {
                                 Tcp::new(listener, conn_event_chan.clone())
                                     .serve(cancel)
                                     .await;
-                                info!(port, "tcp server closed");
+                                info!(port = *port, "tcp server closed");
                             });
                             event
                                 .resp
                                 .send(ClientEventResponse::registered(
-                                    this.entrypoint_config.make_entrypoint(&event.payload, port),
+                                    this.entrypoint_config
+                                        .make_entrypoint(&event.payload, *port),
                                 ))
                                 .unwrap(); // success
                         }
@@ -90,8 +100,8 @@ impl DataServer {
                     }
                 }
                 event::Payload::RegisterUdp { port } => {
-                    let result: Result<(u16, UdpSocket), tonic::Status> =
-                        create_socket::<Udp>(port, this.entrypoint_config.port_range.clone()).await;
+                    let result: Result<(Available, UdpSocket), tonic::Status> =
+                        create_socket::<Udp>(port, &mut self.port_manager).await;
 
                     match result {
                         Ok((port, socket)) => {
@@ -102,12 +112,13 @@ impl DataServer {
                                 Udp::new(socket, conn_event_chan.clone())
                                     .serve(cancel)
                                     .await;
-                                info!(port, "udp server closed");
+                                info!(port = *port, "udp server closed");
                             });
                             event
                                 .resp
                                 .send(ClientEventResponse::registered(
-                                    this.entrypoint_config.make_entrypoint(&event.payload, port),
+                                    this.entrypoint_config
+                                        .make_entrypoint(&event.payload, *port),
                                 ))
                                 .unwrap(); // success
                         }
@@ -179,7 +190,7 @@ impl DataServer {
     }
 
     async fn register_http(
-        &self,
+        mut self,
         shutdown: CancellationToken,
         domain: Bytes,
         subdomain: &mut Bytes,
@@ -234,15 +245,14 @@ impl DataServer {
                 None
             }
         } else {
-            let result =
-                create_socket::<Tcp>(*port, self.entrypoint_config.port_range.clone()).await;
+            let result = create_socket::<Tcp>(*port, &mut self.port_manager).await;
             match result {
-                Ok((random_port, listener)) => {
-                    *port = random_port;
+                Ok((available_port, listener)) => {
+                    *port = *available_port;
                     let conn_event_chan = conn_event_chan.clone();
                     spawn(async move {
                         Http::new(
-                            random_port,
+                            *available_port,
                             Arc::new(Box::new(FixedRegistry::new(conn_event_chan))),
                         )
                         .serve_with_listener(listener, shutdown);

src/server/data_server.rs

@@ -1,5 +1,6 @@
 mod control_server;
 mod data_server;
+mod port;
 mod tunnel;
 pub use control_server::Server;
 
@@ -24,6 +25,7 @@ pub struct EntrypointConfig {
     pub ip: Vec<IpAddr>,
     pub vhttp_behind_proxy_tls: bool,
     pub port_range: RangeInclusive<u16>,
+    pub exclude_ports: Vec<u16>,
 }
 
 impl Default for Config {
@@ -43,6 +45,7 @@ impl Default for EntrypointConfig {
             ip: Vec::new(),
             vhttp_behind_proxy_tls: false,
             port_range: 1024..=65535,
+            exclude_ports: Vec::new(),
         }
     }
 }

src/server/mod.rs

@@ -0,0 +1,123 @@
+use std::sync::Arc;
+
+use dashmap::DashSet;
+use rand::prelude::*;
+use rand::seq::IteratorRandom;
+use rand_chacha::ChaCha20Rng;
+
+pub struct PortManager {
+    rng: ChaCha20Rng,
+    pool: Arc<DashSet<u16>>,
+}
+
+impl PortManager {
+    pub fn new(port_range: std::ops::RangeInclusive<u16>, exclude_ports: Vec<u16>) -> Self {
+        let ports: Vec<_> = port_range
+            .filter(|port| !exclude_ports.contains(port))
+            .collect();
+
+        let pool = DashSet::from_iter(ports);
+        let rng = ChaCha20Rng::from_entropy();
+        Self {
+            rng,
+            pool: Arc::new(pool),
+        }
+    }
+
+    pub fn get(&mut self) -> Option<Available> {
+        let port = *self.pool.iter().choose(&mut self.rng)?;
+        self.pool.remove(&port);
+        Some(Available {
+            port,
+            pool: Some(Arc::clone(&self.pool)),
+        })
+    }
+
+    pub fn remove(&mut self, port: u16) {
+        self.pool.remove(&port);
+    }
+}
+
+#[derive(Debug)]
+pub struct Available {
+    port: u16,
+    /// the pool may not exist if creates `Available` directly.
+    pool: Option<Arc<DashSet<u16>>>,
+}
+
+impl Drop for Available {
+    fn drop(&mut self) {
+        if let Some(pool) = &self.pool {
+            pool.insert(self.port);
+        }
+    }
+}
+
+impl std::ops::Deref for Available {
+    type Target = u16;
+
+    fn deref(&self) -> &Self::Target {
+        &self.port
+    }
+}
+
+impl From<u16> for Available {
+    fn from(port: u16) -> Self {
+        Self { port, pool: None }
+    }
+}
+
+impl From<Available> for u16 {
+    fn from(available: Available) -> Self {
+        available.port
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::ops::RangeInclusive;
+
+    use super::*;
+
+    #[test]
+    fn test_port_manager() {
+        let port_range: RangeInclusive<u16> = 2000..=2100;
+        let exclude_ports = vec![2010, 2020];
+        let len = port_range.len() - exclude_ports.len();
+        let mut port_manager = PortManager::new(port_range, exclude_ports.clone());
+
+        for _ in 0..10000 {
+            let port = port_manager.get();
+            assert!(port.is_some());
+            // auto drop
+        }
+
+        let ports: DashSet<u16> = Default::default();
+        for _ in 0..1000 {
+            let port = port_manager.get();
+            if port.is_none() {
+                break;
+            }
+            ports.insert(*port.unwrap());
+        }
+        assert_eq!(ports.len(), len);
+        assert!(!ports.contains(&2010));
+        assert!(!ports.contains(&2020));
+
+        drop(ports);
+
+        port_manager.remove(2000);
+        let ports: DashSet<u16> = Default::default();
+        for _ in 0..1000 {
+            let port = port_manager.get();
+            if port.is_none() {
+                break;
+            }
+            ports.insert(*port.unwrap());
+        }
+        assert_eq!(ports.len(), len - 1);
+        assert!(!ports.contains(&2010));
+        assert!(!ports.contains(&2020));
+        assert!(!ports.contains(&2000));
+    }
+}

src/server/port.rs

@@ -1,5 +1,3 @@
-use std::ops::RangeInclusive;
-
 use crate::{
     bridge::{self, DataSenderBridge, IdDataSenderBridge},
     event,
@@ -11,6 +9,8 @@ use tokio_util::sync::CancellationToken;
 use tonic::Status;
 use uuid::Uuid;
 
+use super::port::{Available, PortManager};
+
 pub(crate) mod http;
 pub(crate) mod tcp;
 pub(crate) mod udp;
@@ -87,21 +87,26 @@ pub(crate) trait SocketCreator {
 
 pub(crate) async fn create_socket<T: SocketCreator>(
     port: u16,
-    free_port_range: RangeInclusive<u16>,
-) -> anyhow::Result<(u16, T::Output), Status> {
+    port_manager: &mut PortManager,
+) -> anyhow::Result<(Available, T::Output), Status> {
     if port > 0 {
         let socket = T::create_socket(port).await?;
-        Ok((port, socket))
+        Ok((port.into(), socket))
     } else {
         // refer: https://github.com/ekzhang/bore/blob/v0.5.1/src/server.rs#L88
         // todo: a better way to find a free port
         for _ in 0..150 {
-            let freeport = fastrand::u16(free_port_range.clone());
-            let result = T::create_socket(freeport).await;
+            let port: Available = match port_manager.get().await {
+                None => {
+                    return Err(Status::resource_exhausted("no available port"));
+                }
+                Some(port) => port,
+            };
+            let result = T::create_socket(*port).await;
             if result.is_err() {
                 continue;
             }
-            return Ok((freeport, result.unwrap()));
+            return Ok((port, result.unwrap()));
         }
         Err(Status::internal("failed to find a free port"))
     }