portpicker/
lib.rs

1// Modified by `Vector Contributors <vector@datadoghq.com>`.
2// Based on `https://github.com/Dentosal/portpicker-rs` by `Hannes Karppila <hannes.karppila@gmail.com>`.
3// `portpicker-rs` LICENSE:
4// This is free and unencumbered software released into the public domain.
5
6// Anyone is free to copy, modify, publish, use, compile, sell, or
7// distribute this software, either in source code form or as a compiled
8// binary, for any purpose, commercial or non-commercial, and by any
9// means.
10
11// In jurisdictions that recognize copyright laws, the author or authors
12// of this software dedicate any and all copyright interest in the
13// software to the public domain. We make this dedication for the benefit
14// of the public at large and to the detriment of our heirs and
15// successors. We intend this dedication to be an overt act of
16// relinquishment in perpetuity of all present and future rights to this
17// software under copyright law.
18
19// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
20// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
21// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
22// IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
23// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
24// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
25// OTHER DEALINGS IN THE SOFTWARE.
26
27// For more information, please refer to <http://unlicense.org>
28
29#![deny(warnings)]
30
31use std::net::{IpAddr, SocketAddr, TcpListener, ToSocketAddrs, UdpSocket};
32
33use rand::{rng, Rng};
34
35pub type Port = u16;
36
37// Try to bind to a socket using UDP
38fn test_bind_udp<A: ToSocketAddrs>(addr: A) -> Option<Port> {
39    Some(UdpSocket::bind(addr).ok()?.local_addr().ok()?.port())
40}
41
42// Try to bind to a socket using TCP
43fn test_bind_tcp<A: ToSocketAddrs>(addr: A) -> Option<Port> {
44    Some(TcpListener::bind(addr).ok()?.local_addr().ok()?.port())
45}
46
47/// Check if a port is free on UDP
48pub fn is_free_udp(ip: IpAddr, port: Port) -> bool {
49    test_bind_udp(SocketAddr::new(ip, port)).is_some()
50}
51
52/// Check if a port is free on TCP
53pub fn is_free_tcp(ip: IpAddr, port: Port) -> bool {
54    test_bind_tcp(SocketAddr::new(ip, port)).is_some()
55}
56
57/// Check if a port is free on both TCP and UDP
58pub fn is_free(ip: IpAddr, port: Port) -> bool {
59    is_free_tcp(ip, port) && is_free_udp(ip, port)
60}
61
62/// Asks the OS for a free port
63fn ask_free_tcp_port(ip: IpAddr) -> Option<Port> {
64    test_bind_tcp(SocketAddr::new(ip, 0))
65}
66
67/// Picks an available port that is available on both TCP and UDP
68/// ```rust
69/// use portpicker::pick_unused_port;
70/// use std::net::{IpAddr, Ipv4Addr};
71/// let port: u16 = pick_unused_port(IpAddr::V4(Ipv4Addr::LOCALHOST));
72/// ```
73pub fn pick_unused_port(ip: IpAddr) -> Port {
74    let mut rng = rng();
75
76    loop {
77        // Try random port first
78        for _ in 0..10 {
79            let port = rng.random_range(15000..25000);
80            if is_free(ip, port) {
81                return port;
82            }
83        }
84
85        // Ask the OS for a port
86        for _ in 0..10 {
87            if let Some(port) = ask_free_tcp_port(ip) {
88                // Test that the udp port is free as well
89                if is_free_udp(ip, port) {
90                    return port;
91                }
92            }
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
100
101    use super::pick_unused_port;
102
103    #[test]
104    fn ipv4_localhost() {
105        pick_unused_port(IpAddr::V4(Ipv4Addr::LOCALHOST));
106    }
107
108    #[test]
109    fn ipv6_localhost() {
110        pick_unused_port(IpAddr::V6(Ipv6Addr::LOCALHOST));
111    }
112}
portpicker/lib.rs

portpicker/
lib.rs
