Skip to main content

agent_procs/daemon/
port_allocator.rs

1use crate::error::ProxyError;
2use std::collections::HashSet;
3
4const AUTO_PORT_MIN: u16 = 4000;
5const AUTO_PORT_MAX: u16 = 4999;
6
7impl Default for PortAllocator {
8    fn default() -> Self {
9        Self::new()
10    }
11}
12
13pub struct PortAllocator {
14    proxy_enabled: bool,
15    next_auto_port: u16,
16}
17
18impl PortAllocator {
19    pub fn new() -> Self {
20        Self {
21            proxy_enabled: false,
22            next_auto_port: AUTO_PORT_MIN,
23        }
24    }
25
26    pub fn enable_proxy(&mut self) {
27        self.proxy_enabled = true;
28    }
29
30    pub fn is_proxy_enabled(&self) -> bool {
31        self.proxy_enabled
32    }
33
34    pub fn auto_assign_port(&mut self, assigned: &HashSet<u16>) -> Result<u16, ProxyError> {
35        let start = self.next_auto_port;
36        let range_size = (AUTO_PORT_MAX - AUTO_PORT_MIN + 1) as usize;
37
38        for i in 0..range_size {
39            let candidate = AUTO_PORT_MIN
40                + (((self.next_auto_port - AUTO_PORT_MIN) as usize + i) % range_size) as u16;
41            if assigned.contains(&candidate) {
42                continue;
43            }
44            // Bind-test: if we can bind, the port is free (listener drops immediately)
45            if std::net::TcpListener::bind(("127.0.0.1", candidate)).is_ok() {
46                self.next_auto_port = if candidate >= AUTO_PORT_MAX {
47                    AUTO_PORT_MIN
48                } else {
49                    candidate + 1
50                };
51                return Ok(candidate);
52            }
53        }
54        Err(ProxyError::NoFreeAutoPort {
55            min: AUTO_PORT_MIN,
56            max: AUTO_PORT_MAX,
57            start,
58        })
59    }
60}