quickwit_common/
net.rs

1// Copyright (C) 2021 Quickwit, Inc.
2//
3// Quickwit is offered under the AGPL v3.0 and as commercial software.
4// For commercial licensing, contact us at hello@quickwit.io.
5//
6// AGPL:
7// This program is free software: you can redistribute it and/or modify
8// it under the terms of the GNU Affero General Public License as
9// published by the Free Software Foundation, either version 3 of the
10// License, or (at your option) any later version.
11//
12// This program is distributed in the hope that it will be useful,
13// but WITHOUT ANY WARRANTY; without even the implied warranty of
14// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15// GNU Affero General Public License for more details.
16//
17// You should have received a copy of the GNU Affero General Public License
18// along with this program. If not, see <http://www.gnu.org/licenses/>.
19
20use std::net::{SocketAddr, TcpListener, ToSocketAddrs};
21
22/// Finds a random available TCP port.
23pub fn find_available_tcp_port() -> anyhow::Result<u16> {
24    let socket: SocketAddr = ([127, 0, 0, 1], 0u16).into();
25    let listener = TcpListener::bind(socket)?;
26    let port = listener.local_addr()?.port();
27    Ok(port)
28}
29
30/// Converts an object into a resolved `SocketAddr`.
31pub fn get_socket_addr<T: ToSocketAddrs + std::fmt::Debug>(addr: &T) -> anyhow::Result<SocketAddr> {
32    addr.to_socket_addrs()?
33        .next()
34        .ok_or_else(|| anyhow::anyhow!("Failed to resolve address `{:?}`.", addr))
35}
36
37/// Returns true if the socket addr is a valid socket address containing a port.
38///
39/// If the socket adddress looks invalid to begin with, we may return false or true.
40fn contains_port(addr: &str) -> bool {
41    // [IPv6]:port
42    if let Some((_, colon_port)) = addr[1..].rsplit_once(']') {
43        return colon_port.starts_with(':');
44    }
45    if let Some((host, _port)) = addr[1..].rsplit_once(':') {
46        // if host contains a ":" then is thi is probably a IPv6 address.
47        return !host.contains(':');
48    }
49    false
50}
51
52/// Attempts to parse a `socket_addr`.
53/// If no port is defined, it just accepts the address and uses the given default port.
54///
55/// This function supports
56/// - IPv4
57/// - IPv4:port
58/// - IPv6
59/// - \[IPv6\]:port -- IpV6 contains colon. It is customary to require bracket for this reason.
60/// - hostname
61/// - hostname:port
62/// with or without a port.
63///
64/// Note that this function returns a SocketAddr, so that if a hostname
65/// is given, DNS resolution will happen once and for all.
66pub fn parse_socket_addr_with_default_port(
67    addr: &str,
68    default_port: u16,
69) -> anyhow::Result<SocketAddr> {
70    if contains_port(addr) {
71        get_socket_addr(&addr)
72    } else {
73        get_socket_addr(&(addr, default_port))
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    fn test_parse_socket_addr_helper(addr: &str, expected_opt: Option<&str>) {
82        let socket_addr_res = parse_socket_addr_with_default_port(addr, 1337);
83        if let Some(expected) = expected_opt {
84            assert!(
85                socket_addr_res.is_ok(),
86                "Parsing `{}` was expected to succeed",
87                addr
88            );
89            let socket_addr = socket_addr_res.unwrap();
90            let expected_socket_addr: SocketAddr = expected.parse().unwrap();
91            assert_eq!(socket_addr, expected_socket_addr);
92        } else {
93            assert!(
94                socket_addr_res.is_err(),
95                "Parsing `{}` was expected to fail",
96                addr
97            );
98        }
99    }
100
101    #[test]
102    fn test_parse_socket_addr_with_ips() {
103        test_parse_socket_addr_helper("127.0.0.1", Some("127.0.0.1:1337"));
104        test_parse_socket_addr_helper("127.0.0.1:100", Some("127.0.0.1:100"));
105        test_parse_socket_addr_helper("127.0..1:100", None);
106        test_parse_socket_addr_helper(
107            "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
108            Some("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1337"),
109        );
110        test_parse_socket_addr_helper("2001:0db8:85a3:0000:0000:8a2e:0370:7334:1000", None);
111        test_parse_socket_addr_helper(
112            "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1000",
113            Some("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1000"),
114        );
115        test_parse_socket_addr_helper("[2001:0db8:1000", None);
116        test_parse_socket_addr_helper("2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1000", None);
117    }
118
119    // This test require DNS.
120    #[test]
121    fn test_parse_socket_addr_with_resolution() {
122        let socket_addr = parse_socket_addr_with_default_port("google.com:1000", 1337).unwrap();
123        assert_eq!(socket_addr.port(), 1000);
124        let socket_addr = parse_socket_addr_with_default_port("google.com", 1337).unwrap();
125        assert_eq!(socket_addr.port(), 1337);
126    }
127}
quickwit_common/net.rs

quickwit_common/
net.rs
