port_selector/
lib.rs

1pub mod take_up;
2
3use rand::prelude::*;
4use std::net::{
5    Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6, TcpListener, ToSocketAddrs, UdpSocket,
6};
7
8pub type Port = u16;
9
10#[derive(Copy, Clone, Debug)]
11pub struct Selector {
12    pub check_tcp: bool,
13    pub check_udp: bool,
14    pub port_range: (u16, u16),
15    pub max_random_times: u16,
16}
17
18impl Default for Selector {
19    fn default() -> Self {
20        Selector {
21            check_tcp: true,
22            check_udp: true,
23            port_range: (0, 65535),
24            max_random_times: 100,
25        }
26    }
27}
28
29// Try to bind to a socket using TCP
30fn test_bind_tcp<A: ToSocketAddrs>(addr: A) -> Option<Port> {
31    Some(TcpListener::bind(addr).ok()?.local_addr().ok()?.port())
32}
33
34// Try to bind to a socket using UDP
35fn test_bind_udp<A: ToSocketAddrs>(addr: A) -> Option<Port> {
36    Some(UdpSocket::bind(addr).ok()?.local_addr().ok()?.port())
37}
38
39/// Check whether the port is not used on TCP
40pub fn is_free_tcp(port: Port) -> bool {
41    let ipv4 = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port);
42    let ipv6 = SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, port, 0, 0);
43
44    test_bind_tcp(ipv6).is_some() && test_bind_tcp(ipv4).is_some()
45}
46
47/// Check whether the port is not used on UDP
48pub fn is_free_udp(port: Port) -> bool {
49    let ipv4 = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port);
50    let ipv6 = SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, port, 0, 0);
51
52    test_bind_udp(ipv6).is_some() && test_bind_udp(ipv4).is_some()
53}
54
55/// Check whether the port is not used on TCP and UDP
56pub fn is_free(port: Port) -> bool {
57    is_free_tcp(port) && is_free_udp(port)
58}
59
60/// The system randomly assigns available TCP ports
61pub fn random_free_tcp_port() -> Option<Port> {
62    let ipv4 = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0);
63    let ipv6 = SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0);
64
65    test_bind_tcp(ipv6).or_else(|| test_bind_tcp(ipv4))
66}
67
68/// The system randomly assigns available UDP ports
69pub fn random_free_udp_port() -> Option<Port> {
70    let ipv4 = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0);
71    let ipv6 = SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0);
72
73    test_bind_udp(ipv6).or_else(|| test_bind_udp(ipv4))
74}
75
76/// The system randomly assigns available TCP and UDP ports
77pub fn random_free_port() -> Option<Port> {
78    loop {
79        let free_tcp_port = random_free_tcp_port();
80        if free_tcp_port.is_some() && is_free_udp(free_tcp_port?) {
81            break free_tcp_port;
82        }
83    }
84}
85
86/// Check from `given_port` and return the first available port
87/// Return if `given_port` is available; Otherwise `given_port += given_port` until the port is available
88pub fn select_from_given_port(given_port: Port) -> Option<Port> {
89    let mut port = given_port;
90    loop {
91        if is_free(port) {
92            break Some(port);
93        } else {
94            port += 1;
95        }
96    }
97}
98
99/// Gets a matching port based on the `Selector` parameter constraint
100pub fn select_free_port(selector: Selector) -> Option<Port> {
101    let mut rng = rand::thread_rng();
102    let (from, to) = selector.port_range;
103    if selector.check_tcp && selector.check_udp {
104        for _ in 0..selector.max_random_times {
105            let port = rng.gen_range(from..to);
106            if is_free(port) {
107                return Some(port);
108            }
109        }
110    } else if selector.check_tcp {
111        for _ in 0..selector.max_random_times {
112            let port = rng.gen_range(from..to);
113            if is_free_tcp(port) {
114                return Some(port);
115            }
116        }
117    } else if selector.check_udp {
118        for _ in 0..selector.max_random_times {
119            let port = rng.gen_range(from..to);
120            if is_free_udp(port) {
121                return Some(port);
122            }
123        }
124    }
125    // Give up
126    None
127}
128
129#[cfg(test)]
130mod tests {
131
132    use crate::{
133        is_free, is_free_tcp, is_free_udp, random_free_port, random_free_tcp_port,
134        random_free_udp_port, select_free_port, select_from_given_port,
135        take_up::{random_take_up_port, random_take_up_tcp_port, random_take_up_udp_port},
136        test_bind_tcp, test_bind_udp, Selector,
137    };
138    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
139
140    #[test]
141    fn test_is_free() {
142        let used_port = random_take_up_port();
143        assert!(!is_free(used_port));
144        let free_port = random_free_port();
145        assert!(is_free(free_port.unwrap()));
146    }
147
148    #[test]
149    fn test_is_free_tcp() {
150        let used_tcp_port = random_take_up_tcp_port();
151        assert!(!is_free_tcp(used_tcp_port));
152        let free_port = random_free_tcp_port();
153        assert!(is_free_tcp(free_port.unwrap()));
154    }
155
156    #[test]
157    fn test_is_free_udp() {
158        let used_tcp_port = random_take_up_udp_port();
159        assert!(!is_free_udp(used_tcp_port));
160        let free_port = random_free_udp_port();
161        assert!(is_free_udp(free_port.unwrap()));
162    }
163
164    #[test]
165    fn test_free_tcp_port() {
166        let free_tcp_port = random_free_tcp_port();
167        assert!(free_tcp_port.is_some());
168        assert!(is_free_tcp(free_tcp_port.unwrap()));
169    }
170
171    #[test]
172    fn test_free_udp_port() {
173        let free_udp_port = random_free_udp_port();
174        assert!(free_udp_port.is_some());
175        assert!(is_free_udp(free_udp_port.unwrap()));
176    }
177
178    #[test]
179    fn test_pick_unused_port() {
180        let used_port = random_take_up_port();
181        let selector_fail: Selector = Selector {
182            port_range: (used_port, used_port + 1),
183            ..Default::default()
184        };
185        let port = select_free_port(selector_fail);
186        println!("selector_fail, port: {:#?}", &port);
187        assert!(!port.is_some());
188
189        let selector: Selector = Selector {
190            port_range: (50000, 60000),
191            ..Default::default()
192        };
193        for i in 0..100 {
194            let port = select_free_port(selector);
195            println!("index: {}, port: {:#?}", i, &port.unwrap());
196            assert!(&port.unwrap() >= &50000 && &port.unwrap() <= &60000);
197            assert!(port.is_some());
198        }
199    }
200
201    #[test]
202    fn test_test_bind_tcp() {
203        assert!(test_bind_tcp(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)).is_some());
204        assert!(test_bind_tcp(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0)).is_some());
205    }
206
207    #[test]
208    fn test_test_bind_udp() {
209        assert!(test_bind_udp(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)).is_some());
210        assert!(test_bind_udp(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0)).is_some());
211    }
212
213    #[test]
214    fn test_random_free_port() {
215        let port = random_free_port().unwrap();
216        println!("port: {}", &port);
217        assert!(is_free_tcp(port) && is_free_udp(port));
218        assert!(random_free_port().is_some());
219    }
220
221    #[test]
222    fn test_select_from_given_port() {
223        let port = select_from_given_port(30000).unwrap();
224        println!("port: {}", &port);
225        assert!(is_free_tcp(port) && is_free_udp(port));
226        let mut used_port = random_take_up_port();
227        println!("used_port: {}", &used_port);
228        let new_port = select_from_given_port(used_port).unwrap();
229        println!("new_port: {}", &new_port);
230        let used_port = loop {
231            if is_free(used_port) {
232                break used_port;
233            }
234            used_port += 1;
235        };
236        assert_eq!(new_port, used_port);
237    }
238}
port_selector/lib.rs

port_selector/
lib.rs
