Workspacify

host/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "host"
+edition = "2024"
+
+[dependencies]
+error-stack = { workspace = true }
+futures = { workspace = true }
+resolver = { workspace = true }
+thiserror = { workspace = true }
+tokio = { workspace = true }
+tracing = { workspace = true }

host/src/lib.rs

@@ -0,0 +1,208 @@
+use error_stack::{report, ResultExt};
+use futures::stream::{FuturesUnordered, StreamExt};
+use resolver::name::Name;
+use resolver::{ResolveError, Resolver};
+use std::collections::HashSet;
+use std::fmt;
+use std::net::{AddrParseError, IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::str::FromStr;
+use thiserror::Error;
+use tokio::net::TcpStream;
+
+#[derive(Debug)]
+pub enum Host {
+    IPv4(Ipv4Addr),
+    IPv6(Ipv6Addr),
+    Hostname(Name),
+}
+
+#[derive(Debug, Error)]
+pub enum HostParseError {
+    #[error("Could not parse IPv6 address {address:?}: {error}")]
+    CouldNotParseIPv6 {
+        address: String,
+        error: AddrParseError,
+    },
+    #[error("Invalid hostname {hostname:?}")]
+    InvalidHostname { hostname: String },
+}
+
+impl fmt::Display for Host {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Host::IPv4(x) => x.fmt(f),
+            Host::IPv6(x) => x.fmt(f),
+            Host::Hostname(x) => x.fmt(f),
+        }
+    }
+}
+
+impl FromStr for Host {
+    type Err = HostParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Ok(addr) = Ipv4Addr::from_str(s) {
+            return Ok(Host::IPv4(addr));
+        }
+
+        if let Ok(addr) = Ipv6Addr::from_str(s) {
+            return Ok(Host::IPv6(addr));
+        }
+
+        if let Some(prefix_removed) = s.strip_prefix('[') {
+            if let Some(cleaned) = prefix_removed.strip_suffix(']') {
+                match Ipv6Addr::from_str(cleaned) {
+                    Ok(addr) => return Ok(Host::IPv6(addr)),
+                    Err(error) => {
+                        return Err(HostParseError::CouldNotParseIPv6 {
+                            address: s.to_string(),
+                            error,
+                        })
+                    }
+                }
+            }
+        }
+
+        if let Ok(name) = Name::from_str(s) {
+            return Ok(Host::Hostname(name));
+        }
+
+        println!("   ... not a hostname");
+        Err(HostParseError::InvalidHostname {
+            hostname: s.to_string(),
+        })
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum ConnectionError {
+    #[error("Connection error: failed to resolve host")]
+    ResolveError,
+    #[error("No valid IP addresses found")]
+    NoAddresses,
+    #[error("Error connecting to host: {error}")]
+    ConnectionError {
+        #[from]
+        error: std::io::Error,
+    },
+}
+
+impl Host {
+    /// Resolve this host address to a set of underlying IP addresses.
+    ///
+    /// It is possible that the set of addresses provided may be empty, if the
+    /// address properly resolves (as in, we get a good DNS response) but there
+    /// are no relevant records for us to use for IPv4 or IPv6 connections. There
+    /// is also no guarantee that the host will have both IPv4 and IPv6 addresses,
+    /// so you may only see one or the other.
+    pub async fn resolve(
+        &self,
+        resolver: &mut Resolver,
+    ) -> error_stack::Result<HashSet<IpAddr>, ResolveError> {
+        match self {
+            Host::IPv4(addr) => Ok(HashSet::from([IpAddr::V4(*addr)])),
+            Host::IPv6(addr) => Ok(HashSet::from([IpAddr::V6(*addr)])),
+            Host::Hostname(name) => resolver.lookup(name).await,
+        }
+    }
+
+    /// Connect to this host and port.
+    ///
+    /// This routine will attempt to connect to every address provided by the
+    /// resolver, and return the first successful connection. If all of the
+    /// connections fail, it will return the first error it receives. This routine
+    /// will also return an error if there are no addresses to connect to (which
+    /// can happen in cases in which [`Host::resolve`] would return an empty set.
+    pub async fn connect(
+        &self,
+        resolver: &mut Resolver,
+        port: u16,
+    ) -> error_stack::Result<TcpStream, ConnectionError> {
+        let addresses = self
+            .resolve(resolver)
+            .await
+            .change_context(ConnectionError::ResolveError)
+            .attach_printable_lazy(|| format!("target address {}", self))?;
+
+        let mut connectors = FuturesUnordered::new();
+
+        for address in addresses.into_iter() {
+            tracing::trace!(?address, "adding possible target address");
+            let connect_future = TcpStream::connect(SocketAddr::new(address, port));
+            connectors.push(connect_future);
+        }
+
+        let mut error = None;
+
+        while let Some(result) = connectors.next().await {
+            match result {
+                Err(e) if error.is_none() => error = Some(e),
+                Err(_) => {}
+                Ok(v) => return Ok(v),
+            }
+        }
+
+        let final_error = if let Some(e) = error {
+            ConnectionError::ConnectionError { error: e }
+        } else {
+            ConnectionError::NoAddresses
+        };
+
+        Err(report!(final_error)).attach_printable_lazy(|| format!("target address {}", self))
+    }
+}
+
+#[test]
+fn ip4_hosts_work() {
+    assert!(
+        matches!(Host::from_str("127.0.0.1"), Ok(Host::IPv4(addr)) if addr == Ipv4Addr::new(127, 0, 0, 1))
+    );
+}
+
+#[test]
+fn bare_ip6_hosts_work() {
+    assert!(matches!(
+        Host::from_str("2001:0db8:85a3:0000:0000:8a2e:0370:7334"),
+        Ok(Host::IPv6(_))
+    ));
+    assert!(matches!(Host::from_str("2001:db8::1"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("2001:DB8::1"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("::1"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("::"), Ok(Host::IPv6(_))));
+}
+
+#[test]
+fn wrapped_ip6_hosts_work() {
+    assert!(matches!(
+        Host::from_str("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]"),
+        Ok(Host::IPv6(_))
+    ));
+    assert!(matches!(Host::from_str("[2001:db8::1]"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("[2001:DB8::1]"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("[::1]"), Ok(Host::IPv6(_))));
+    assert!(matches!(Host::from_str("[::]"), Ok(Host::IPv6(_))));
+}
+
+#[test]
+fn valid_domains_work() {
+    assert!(matches!(
+        Host::from_str("uhsure.com"),
+        Ok(Host::Hostname(_))
+    ));
+    assert!(matches!(
+        Host::from_str("www.cs.indiana.edu"),
+        Ok(Host::Hostname(_))
+    ));
+}
+
+#[test]
+fn invalid_inputs_fail() {
+    assert!(matches!(
+        Host::from_str("[uhsure.com]"),
+        Err(HostParseError::CouldNotParseIPv6 { .. })
+    ));
+    assert!(matches!(
+        Host::from_str("-uhsure.com"),
+        Err(HostParseError::InvalidHostname { .. })
+    ));
+}