port_selector/
take_up.rs

1use std::{
2    net::{Ipv4Addr, SocketAddrV4, TcpListener, UdpSocket},
3    thread,
4};
5
6use crate::{is_free, is_free_tcp, is_free_udp, random_free_port, Port};
7
8/// Run TCP Server to take up port on TCP
9fn set_up_tcp_listener(port: Option<Port>) -> Port {
10    let port = port.unwrap_or(0);
11    if port > 0 && !is_free_tcp(port) {
12        return port;
13    }
14    let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port)).unwrap();
15    let port = listener.local_addr().ok().unwrap().port();
16    thread::spawn(move || {
17        listener.accept().expect("Failed to accept TCP connection");
18    });
19    port
20}
21
22/// Take up given port on TCP
23pub fn take_up_tcp_port(port: Port) -> bool {
24    !is_free_tcp(set_up_tcp_listener(Some(port)))
25}
26
27/// Take up port randomly on TCP
28pub fn random_take_up_tcp_port() -> Port {
29    set_up_tcp_listener(None)
30}
31
32/// Run UDP Server to take up port on UDP
33fn set_up_udp_listener(port: Option<Port>) -> Port {
34    let port = port.unwrap_or(0);
35    if port > 0 && !is_free_udp(port) {
36        return port;
37    }
38    let listener = UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port)).unwrap();
39    let port = listener.local_addr().ok().unwrap().port();
40    thread::spawn(move || {
41        listener
42            .recv_from(&mut [0u8])
43            .expect("Failed to receive UDP packet");
44    });
45    port
46}
47
48/// Take up given port on UDP
49pub fn take_up_udp_port(port: Port) -> bool {
50    !is_free_udp(set_up_udp_listener(Some(port)))
51}
52
53/// Take up port randomly on UDP
54pub fn random_take_up_udp_port() -> Port {
55    set_up_udp_listener(None)
56}
57
58/// Take up port randomly on TCP && UDP
59pub fn take_up_port(port: Port) -> bool {
60    if is_free_tcp(port) {
61        take_up_tcp_port(port);
62    }
63    if is_free_udp(port) {
64        take_up_udp_port(port);
65    }
66    !is_free(port)
67}
68
69/// Take up port randomly on TCP && UDP
70pub fn random_take_up_port() -> Port {
71    loop {
72        let free_port = random_free_port().expect("Fail to get free port");
73        if take_up_port(free_port) {
74            break free_port;
75        }
76    }
77}
78
79#[cfg(test)]
80mod take_up_tests {
81    use crate::{
82        is_free, is_free_tcp, is_free_udp, random_free_port, random_free_tcp_port,
83        random_free_udp_port,
84        take_up::{
85            random_take_up_port, random_take_up_tcp_port, random_take_up_udp_port,
86            set_up_udp_listener, take_up_port, take_up_tcp_port, take_up_udp_port,
87        },
88    };
89
90    use super::set_up_tcp_listener;
91
92    #[test]
93    fn test_set_up_tcp_listener() {
94        let free_tcp_port = random_free_tcp_port();
95        assert_eq!(free_tcp_port.unwrap(), set_up_tcp_listener(free_tcp_port));
96        assert!(!is_free_tcp(free_tcp_port.unwrap()));
97    }
98
99    #[test]
100    fn test_set_up_udp_listener() {
101        let free_udp_port = random_free_udp_port();
102        assert_eq!(free_udp_port.unwrap(), set_up_udp_listener(free_udp_port));
103        assert!(!is_free_udp(free_udp_port.unwrap()));
104    }
105
106    #[test]
107    fn test_take_up_tcp_port() {
108        let free_tcp_port = random_free_tcp_port();
109        assert!(free_tcp_port.is_some());
110        let res = take_up_tcp_port(free_tcp_port.unwrap());
111        assert_eq!(res, !is_free_tcp(free_tcp_port.unwrap()));
112    }
113
114    #[test]
115    fn test_random_take_up_tcp_port() {
116        let used_tcp_port = random_take_up_tcp_port();
117        assert_eq!(is_free_tcp(used_tcp_port), false);
118    }
119
120    #[test]
121    fn test_take_up_udp_port() {
122        let free_udp_port = random_free_udp_port();
123        assert!(free_udp_port.is_some());
124        assert_eq!(
125            take_up_udp_port(free_udp_port.unwrap()),
126            !is_free_udp(free_udp_port.unwrap())
127        );
128    }
129
130    #[test]
131    fn test_random_take_up_udp_port() {
132        let used_udp_port = random_take_up_udp_port();
133        assert!(!is_free_udp(used_udp_port));
134    }
135
136    #[test]
137    fn test_take_up_port() {
138        let free_port = random_free_port();
139        assert!(free_port.is_some());
140        let is_used_port = take_up_port(free_port.unwrap());
141        assert!(is_used_port);
142        assert_eq!(is_free(free_port.unwrap()), false);
143    }
144
145    #[test]
146    fn test_random_take_up_port() {
147        let used_port = random_take_up_port();
148        assert!(!is_free(used_port));
149        assert!(!is_free_tcp(used_port));
150        assert!(!is_free_udp(used_port));
151    }
152}
port_selector/take_up.rs

port_selector/
take_up.rs
