Skip to main content

stratum_apps/network_helpers/
resolve_hostname.rs

1use std::{
2    net::{IpAddr, SocketAddr},
3    time::Duration,
4};
5
6use tracing::{debug, info};
7
8/// Maximum time to wait for a DNS lookup before giving up.
9/// DNS resolution should complete in milliseconds on a healthy network;
10/// 5 seconds is generous enough for slow links while still failing fast
11/// when DNS is broken.
12const DNS_TIMEOUT: Duration = Duration::from_secs(5);
13
14/// Errors that can occur during address resolution.
15#[derive(Debug)]
16pub enum ResolveError {
17    /// DNS lookup returned no results for the given hostname.
18    NoResults(String),
19    /// DNS lookup failed with an IO error.
20    LookupFailed(std::io::Error),
21    /// DNS lookup did not complete within the timeout.
22    Timeout(String),
23}
24
25impl std::fmt::Display for ResolveError {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            ResolveError::NoResults(host) => {
29                write!(f, "DNS resolution returned no results for '{host}'")
30            }
31            ResolveError::LookupFailed(e) => write!(f, "DNS resolution failed: {e}"),
32            ResolveError::Timeout(host) => {
33                write!(
34                    f,
35                    "DNS resolution for '{host}' timed out after {}s",
36                    DNS_TIMEOUT.as_secs()
37                )
38            }
39        }
40    }
41}
42
43impl std::error::Error for ResolveError {}
44
45/// Resolves a host string and port to every [`SocketAddr`] it maps to.
46///
47/// This function first attempts to parse the host as an IP address (fast path, no DNS).
48/// If that fails, it performs an async DNS lookup via [`tokio::net::lookup_host`].
49///
50/// Every resolved address is returned, in resolver order. Pass the slice straight to
51/// [`tokio::net::TcpStream::connect`], which attempts each address until one succeeds, so a
52/// dual-stack host still connects over IPv4 when its IPv6 address is unreachable.
53///
54/// This should be called at connection time (not config parse time) so that DNS changes
55/// are picked up on reconnection attempts.
56///
57/// # Examples
58///
59/// ```ignore
60/// // IP address (fast path)
61/// let addrs = resolve_host("127.0.0.1", 3333).await?;
62///
63/// // Hostname (DNS lookup)
64/// let addrs = resolve_host("pool.example.com", 3333).await?;
65/// let stream = TcpStream::connect(&addrs[..]).await?;
66/// ```
67pub async fn resolve_host(host: &str, port: u16) -> Result<Vec<SocketAddr>, ResolveError> {
68    // Fast path: try parsing as an IP address directly (no DNS needed)
69    if let Ok(ip) = host.parse::<IpAddr>() {
70        return Ok(vec![SocketAddr::new(ip, port)]);
71    }
72
73    // Slow path: perform async DNS resolution
74    info!("Resolving hostname '{host}' via DNS...");
75    let lookup = format!("{host}:{port}");
76    let addrs: Vec<SocketAddr> =
77        tokio::time::timeout(DNS_TIMEOUT, tokio::net::lookup_host(&lookup))
78            .await
79            .map_err(|_| ResolveError::Timeout(host.to_string()))?
80            .map_err(ResolveError::LookupFailed)?
81            .collect();
82
83    if addrs.is_empty() {
84        return Err(ResolveError::NoResults(host.to_string()));
85    }
86
87    debug!("Resolved '{host}' -> {addrs:?}");
88    Ok(addrs)
89}
90
91/// Resolves a `"host:port"` string to a single [`SocketAddr`].
92///
93/// Accepts both IP addresses and hostnames in the `"host:port"` format.
94/// For hostnames, performs async DNS resolution via [`tokio::net::lookup_host`].
95///
96/// Only the first resolved address is returned, which suits callers that need one concrete
97/// address to describe an endpoint. To open a connection, use [`resolve_host`] or hand the
98/// `"host:port"` string to [`tokio::net::TcpStream::connect`] instead, so that every address
99/// behind the name gets an attempt.
100///
101/// # Examples
102///
103/// ```ignore
104/// // IP address (fast path)
105/// let addr = resolve_host_port("127.0.0.1:3333").await?;
106///
107/// // Hostname (DNS lookup)
108/// let addr = resolve_host_port("pool.example.com:3333").await?;
109/// ```
110pub async fn resolve_host_port(addr: &str) -> Result<SocketAddr, ResolveError> {
111    // Fast path: try parsing as a SocketAddr directly (no DNS needed)
112    if let Ok(socket) = addr.parse::<SocketAddr>() {
113        return Ok(socket);
114    }
115
116    // Slow path: perform async DNS resolution
117    info!("Resolving address '{addr}' via DNS...");
118    let resolved = tokio::time::timeout(DNS_TIMEOUT, tokio::net::lookup_host(addr))
119        .await
120        .map_err(|_| ResolveError::Timeout(addr.to_string()))?
121        .map_err(ResolveError::LookupFailed)?
122        // DNS can return multiple addresses; take the first one
123        .next()
124        .ok_or_else(|| ResolveError::NoResults(addr.to_string()))?;
125
126    debug!("Resolved '{addr}' -> {resolved}");
127    Ok(resolved)
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[tokio::test]
135    async fn resolve_ipv4_address() {
136        let addrs = resolve_host("127.0.0.1", 3333).await.unwrap();
137        assert_eq!(
138            addrs,
139            vec![SocketAddr::new("127.0.0.1".parse().unwrap(), 3333)]
140        );
141    }
142
143    #[tokio::test]
144    async fn resolve_ipv6_address() {
145        let addrs = resolve_host("::1", 3333).await.unwrap();
146        assert_eq!(addrs, vec![SocketAddr::new("::1".parse().unwrap(), 3333)]);
147    }
148
149    #[tokio::test]
150    async fn resolve_localhost_hostname() {
151        let addrs = resolve_host("localhost", 3333).await.unwrap();
152        assert!(!addrs.is_empty());
153        // localhost can resolve to 127.0.0.1, ::1, or both depending on the system
154        assert!(
155            addrs
156                .iter()
157                .all(|a| a.port() == 3333 && a.ip().is_loopback())
158        );
159    }
160
161    /// Every address behind a name must survive resolution: dropping any of them leaves a
162    /// dual-stack pool unreachable whenever the resolver happens to list a dead address first.
163    #[tokio::test]
164    async fn resolve_keeps_every_address_dns_returns() {
165        let expected: Vec<SocketAddr> = tokio::net::lookup_host("localhost:3333")
166            .await
167            .unwrap()
168            .collect();
169        assert_eq!(resolve_host("localhost", 3333).await.unwrap(), expected);
170    }
171
172    #[tokio::test]
173    async fn resolve_invalid_hostname_fails() {
174        let result = resolve_host("this.hostname.definitely.does.not.exist.invalid", 3333).await;
175        assert!(result.is_err());
176    }
177
178    #[tokio::test]
179    async fn resolve_host_port_ipv4() {
180        let addr = resolve_host_port("127.0.0.1:3333").await.unwrap();
181        assert_eq!(addr, SocketAddr::new("127.0.0.1".parse().unwrap(), 3333));
182    }
183
184    #[tokio::test]
185    async fn resolve_host_port_localhost() {
186        let addr = resolve_host_port("localhost:3333").await.unwrap();
187        assert_eq!(addr.port(), 3333);
188        assert!(addr.ip().is_loopback());
189    }
190
191    #[tokio::test]
192    async fn resolve_host_port_invalid_fails() {
193        let result =
194            resolve_host_port("this.hostname.definitely.does.not.exist.invalid:3333").await;
195        assert!(result.is_err());
196    }
197}