nuclei/net/
addr.rs

1//! The addresses an operation takes: numbers, used at once, or names,
2//! looked up on the blocking pool, so no shard ever waits on a resolver.
3//!
4//! Numbers make no allocation: one address, or a slice of them, is walked
5//! where it is. A name is copied once, to cross to the pool's thread, and
6//! the resolver's own list of results is handed back as it is.
7
8use std::io;
9use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
10use std::{iter, option, slice, vec};
11
12/// What an operation that takes an address accepts: a `SocketAddr`, an
13/// `(IpAddr, u16)` pair, a slice or array of addresses, or a string,
14/// `"host:port"` or `("host", port)`. Numbers are used as they are; a name
15/// is looked up on the blocking pool (`getaddrinfo`), never on the calling
16/// thread.
17///
18/// The trait is sealed: its implementations are the ones listed here.
19pub trait ToSocketAddrs: sealed::Sealed {}
20
21pub(crate) mod sealed {
22    /// Each implementation's addresses, or the name to look up.
23    pub trait Sealed {
24        fn resolve(&self) -> super::Resolve<'_>;
25    }
26}
27
28/// What an address names before any lookup, borrowed from it.
29pub enum Resolve<'a> {
30    One(SocketAddr),
31    Slice(&'a [SocketAddr]),
32    Name(&'a str),
33    HostPort(&'a str, u16),
34}
35
36/// The addresses to try, in order.
37pub(crate) enum Addrs<'a> {
38    One(option::IntoIter<SocketAddr>),
39    Slice(iter::Copied<slice::Iter<'a, SocketAddr>>),
40    /// The resolver's own list.
41    Looked(vec::IntoIter<SocketAddr>),
42}
43
44impl Iterator for Addrs<'_> {
45    type Item = SocketAddr;
46
47    fn next(&mut self) -> Option<SocketAddr> {
48        match self {
49            Self::One(one) => one.next(),
50            Self::Slice(slice) => slice.next(),
51            Self::Looked(looked) => looked.next(),
52        }
53    }
54
55    fn size_hint(&self) -> (usize, Option<usize>) {
56        match self {
57            Self::One(one) => one.size_hint(),
58            Self::Slice(slice) => slice.size_hint(),
59            Self::Looked(looked) => looked.size_hint(),
60        }
61    }
62}
63
64impl<'a> Resolve<'a> {
65    /// The addresses, looked up on the blocking pool if they are a name.
66    async fn addrs(self) -> io::Result<Addrs<'a>> {
67        use std::net::ToSocketAddrs as _;
68        let looked = match self {
69            Self::One(addr) => return Ok(Addrs::One(Some(addr).into_iter())),
70            Self::Slice(addrs) => return Ok(Addrs::Slice(addrs.iter().copied())),
71            // Owned to cross threads: the pool's job may outlive a caller
72            // that was cancelled.
73            Self::Name(name) => {
74                let name = name.to_owned();
75                crate::rt::blocking::call(move || name.as_str().to_socket_addrs()).await?
76            }
77            Self::HostPort(host, port) => {
78                let host = host.to_owned();
79                crate::rt::blocking::call(move || (host.as_str(), port).to_socket_addrs()).await?
80            }
81        };
82        Ok(Addrs::Looked(looked))
83    }
84}
85
86macro_rules! at_once {
87    ($($ty:ty, $this:ident => $addr:expr;)*) => {$(
88        impl ToSocketAddrs for $ty {}
89        impl sealed::Sealed for $ty {
90            fn resolve(&self) -> Resolve<'_> {
91                let $this = self;
92                Resolve::One($addr)
93            }
94        }
95    )*};
96}
97
98at_once! {
99    SocketAddr, a => *a;
100    SocketAddrV4, a => SocketAddr::V4(*a);
101    SocketAddrV6, a => SocketAddr::V6(*a);
102    (IpAddr, u16), a => SocketAddr::new(a.0, a.1);
103    (Ipv4Addr, u16), a => SocketAddr::new(a.0.into(), a.1);
104    (Ipv6Addr, u16), a => SocketAddr::new(a.0.into(), a.1);
105}
106
107impl ToSocketAddrs for [SocketAddr] {}
108impl sealed::Sealed for [SocketAddr] {
109    fn resolve(&self) -> Resolve<'_> {
110        Resolve::Slice(self)
111    }
112}
113
114impl<const N: usize> ToSocketAddrs for [SocketAddr; N] {}
115impl<const N: usize> sealed::Sealed for [SocketAddr; N] {
116    fn resolve(&self) -> Resolve<'_> {
117        Resolve::Slice(self)
118    }
119}
120
121impl ToSocketAddrs for str {}
122impl sealed::Sealed for str {
123    fn resolve(&self) -> Resolve<'_> {
124        match self.parse() {
125            Ok(addr) => Resolve::One(addr),
126            Err(_) => Resolve::Name(self),
127        }
128    }
129}
130
131impl ToSocketAddrs for String {}
132impl sealed::Sealed for String {
133    fn resolve(&self) -> Resolve<'_> {
134        self.as_str().resolve()
135    }
136}
137
138/// A host, by number or by name, and a port.
139fn host_port(host: &str, port: u16) -> Resolve<'_> {
140    match host.parse::<IpAddr>() {
141        Ok(ip) => Resolve::One(SocketAddr::new(ip, port)),
142        Err(_) => Resolve::HostPort(host, port),
143    }
144}
145
146impl ToSocketAddrs for (&str, u16) {}
147impl sealed::Sealed for (&str, u16) {
148    fn resolve(&self) -> Resolve<'_> {
149        host_port(self.0, self.1)
150    }
151}
152
153impl ToSocketAddrs for (String, u16) {}
154impl sealed::Sealed for (String, u16) {
155    fn resolve(&self) -> Resolve<'_> {
156        host_port(&self.0, self.1)
157    }
158}
159
160impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {}
161impl<T: ToSocketAddrs + ?Sized> sealed::Sealed for &T {
162    fn resolve(&self) -> Resolve<'_> {
163        (**self).resolve()
164    }
165}
166
167/// Every address `host` names: numbers at once, a name looked up on the
168/// blocking pool. The iterator owns what it walks, so a slice given here is
169/// copied into it; the operations that take an address walk it in place.
170///
171/// ```
172/// # nuclei::block_on(async {
173/// let addrs: Vec<_> = nuclei::net::lookup_host("127.0.0.1:80").await?.collect();
174/// assert_eq!(addrs, ["127.0.0.1:80".parse().unwrap()]);
175/// # std::io::Result::Ok(())
176/// # }).unwrap();
177/// ```
178///
179/// # Errors
180///
181/// When the name cannot be looked up.
182pub async fn lookup_host<T: ToSocketAddrs>(
183    host: T,
184) -> io::Result<impl Iterator<Item = SocketAddr> + Send + 'static> {
185    Ok(match resolve(&host).await? {
186        Addrs::One(one) => Addrs::One(one),
187        Addrs::Slice(slice) => Addrs::Looked(slice.collect::<Vec<_>>().into_iter()),
188        Addrs::Looked(looked) => Addrs::Looked(looked),
189    })
190}
191
192/// The addresses `addr` names, walked in place for numbers.
193pub(crate) async fn resolve<A: ToSocketAddrs + ?Sized>(addr: &A) -> io::Result<Addrs<'_>> {
194    sealed::Sealed::resolve(addr).addrs().await
195}
196
197/// The first address `addr` names, or `None` if it names none: a number at
198/// once, with no lookup's future polled for it.
199pub(crate) async fn first<A: ToSocketAddrs + ?Sized>(addr: &A) -> io::Result<Option<SocketAddr>> {
200    match sealed::Sealed::resolve(addr) {
201        Resolve::One(addr) => Ok(Some(addr)),
202        Resolve::Slice(addrs) => Ok(addrs.first().copied()),
203        name => Ok(name.addrs().await?.next()),
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    fn here<T: ToSocketAddrs + ?Sized>(addr: &T) -> Vec<SocketAddr> {
212        crate::block_on(resolve(addr)).unwrap().collect()
213    }
214
215    #[test]
216    fn numbers_are_used_as_they_are() {
217        let a: SocketAddr = "127.0.0.1:80".parse().unwrap();
218        let b: SocketAddr = "[::1]:443".parse().unwrap();
219        assert_eq!(here(&a), [a]);
220        assert_eq!(here("127.0.0.1:80"), [a]);
221        assert_eq!(here(&("127.0.0.1", 80)), [a]);
222        assert_eq!(here(&("::1", 443)), [b]);
223        assert_eq!(here(&(IpAddr::from([127, 0, 0, 1]), 80)), [a]);
224        assert_eq!(here(&[a, b]), [a, b]);
225        assert_eq!(here(&[a, b][..]), [a, b]);
226        assert!(matches!(
227            sealed::Sealed::resolve("127.0.0.1:80"),
228            Resolve::One(_)
229        ));
230        assert!(matches!(
231            sealed::Sealed::resolve("localhost:80"),
232            Resolve::Name(_)
233        ));
234    }
235
236    #[test]
237    fn a_name_is_looked_up_on_the_pool() {
238        let addrs = here("localhost:80");
239        assert!(!addrs.is_empty());
240        assert!(addrs.iter().all(|a| a.ip().is_loopback() && a.port() == 80));
241        assert!(crate::block_on(resolve("no.such.host.invalid:80")).is_err());
242    }
243
244    #[test]
245    fn the_first_address_is_the_lookups_first() {
246        let a: SocketAddr = "127.0.0.1:80".parse().unwrap();
247        let b: SocketAddr = "[::1]:443".parse().unwrap();
248        let none: [SocketAddr; 0] = [];
249        assert_eq!(crate::block_on(first(&a)).unwrap(), Some(a));
250        assert_eq!(crate::block_on(first(&[b, a][..])).unwrap(), Some(b));
251        assert_eq!(crate::block_on(first(&none[..])).unwrap(), None);
252        assert_eq!(crate::block_on(first("[::1]:443")).unwrap(), Some(b));
253        let looked = crate::block_on(first("localhost:80")).unwrap();
254        assert!(looked.is_some_and(|l| l.ip().is_loopback() && l.port() == 80));
255        assert!(crate::block_on(first("no.such.host.invalid:80")).is_err());
256    }
257}
nuclei/net/addr.rs

nuclei/net/
addr.rs
