support socket addresses in metatls (#1397)

Summary:
Pull Request resolved: #1397

Fix up address handling in general.

This will also let us support wildcard addresses.
ghstack-source-id: 313884234
exported-using-ghexport

Reviewed By: zdevito

Differential Revision: D83689820

fbshipit-source-id: bef02792f3207acb2311a5c795512e5687bd9cd2

Author: mariusae

hyper/src/utils/system_address.rs

@@ -10,6 +10,7 @@ use std::str::FromStr;
 
 use anyhow;
 use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::MetaTlsAddr;
 
 /// Extended type to represent a system address which can be a ChannelAdd or a MAST job name.
 #[derive(Clone, Debug)]
@@ -62,7 +63,10 @@ async fn parse_system_address_or_mast_job(address: &str) -> Result<ChannelAddr,
             let (host, port) = SMCClient::new(fbinit::expect_init(), smc_tier)?
                 .get_system_address()
                 .await?;
-            let channel_address = ChannelAddr::MetaTls(canonicalize_hostname(&host), port);
+            let channel_address = ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: canonicalize_hostname(&host),
+                port,
+            });
             Ok(channel_address)
         }
     }

hyperactor/src/channel.rs

@@ -19,6 +19,7 @@ use std::os::linux::net::SocketAddrExt;
 use std::str::FromStr;
 
 use async_trait::async_trait;
+use enum_as_inner::EnumAsInner;
 use lazy_static::lazy_static;
 use local_ip_address::local_ipv6;
 use serde::Deserialize;
@@ -256,6 +257,60 @@ pub enum TlsMode {
     // TODO: consider adding IpV4 support.
 }
 
+/// Address format for MetaTls channels. Supports both hostname/port pairs
+/// (required for clients for host identity) and direct socket addresses
+/// (allowed for servers).
+#[derive(
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+    Ord,
+    PartialOrd,
+    EnumAsInner
+)]
+pub enum MetaTlsAddr {
+    /// Hostname and port pair. Required for clients to establish host identity.
+    Host {
+        /// The hostname to connect to.
+        hostname: Hostname,
+        /// The port to connect to.
+        port: Port,
+    },
+    /// Direct socket address. Allowed for servers.
+    Socket(SocketAddr),
+}
+
+impl MetaTlsAddr {
+    /// Returns the port number for this address.
+    pub fn port(&self) -> Port {
+        match self {
+            Self::Host { port, .. } => *port,
+            Self::Socket(addr) => addr.port(),
+        }
+    }
+
+    /// Returns the hostname if this is a Host variant, None otherwise.
+    pub fn hostname(&self) -> Option<&str> {
+        match self {
+            Self::Host { hostname, .. } => Some(hostname),
+            Self::Socket(_) => None,
+        }
+    }
+}
+
+impl fmt::Display for MetaTlsAddr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Host { hostname, port } => write!(f, "{}:{}", hostname, port),
+            Self::Socket(addr) => write!(f, "{}", addr),
+        }
+    }
+}
+
 /// Types of channel transports.
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Named)]
 pub enum ChannelTransport {
@@ -402,8 +457,9 @@ pub enum ChannelAddr {
     Tcp(SocketAddr),
 
     /// An address to establish TCP channels with TLS support within Meta.
-    /// Composed of hostname and port.
-    MetaTls(Hostname, Port),
+    /// Supports both hostname/port pairs (required for clients) and
+    /// socket addresses (allowed for servers).
+    MetaTls(MetaTlsAddr),
 
     /// Local addresses are registered in-process and given an integral
     /// index.
@@ -471,7 +527,10 @@ impl ChannelAddr {
                         .and_then(|addr| addr.to_string().parse().ok())
                         .expect("failed to retrieve ipv6 address"),
                 };
-                Self::MetaTls(host_address, 0)
+                Self::MetaTls(MetaTlsAddr::Host {
+                    hostname: host_address,
+                    port: 0,
+                })
             }
             ChannelTransport::Local => Self::Local(0),
             ChannelTransport::Sim(transport) => sim::any(*transport),
@@ -484,12 +543,16 @@ impl ChannelAddr {
     pub fn transport(&self) -> ChannelTransport {
         match self {
             Self::Tcp(_) => ChannelTransport::Tcp,
-            Self::MetaTls(address, _) => match address.parse::<IpAddr>() {
-                Ok(ip) => match ip {
+            Self::MetaTls(addr) => match addr {
+                MetaTlsAddr::Host { hostname, .. } => match hostname.parse::<IpAddr>() {
+                    Ok(IpAddr::V6(_)) => ChannelTransport::MetaTls(TlsMode::IpV6),
+                    Ok(IpAddr::V4(_)) => ChannelTransport::MetaTls(TlsMode::Hostname),
+                    Err(_) => ChannelTransport::MetaTls(TlsMode::Hostname),
+                },
+                MetaTlsAddr::Socket(socket_addr) => match socket_addr.ip() {
                     IpAddr::V6(_) => ChannelTransport::MetaTls(TlsMode::IpV6),
                     IpAddr::V4(_) => ChannelTransport::MetaTls(TlsMode::Hostname),
                 },
-                Err(_) => ChannelTransport::MetaTls(TlsMode::Hostname),
             },
             Self::Local(_) => ChannelTransport::Local,
             Self::Sim(addr) => ChannelTransport::Sim(Box::new(addr.transport())),
@@ -502,7 +565,7 @@ impl fmt::Display for ChannelAddr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::Tcp(addr) => write!(f, "tcp:{}", addr),
-            Self::MetaTls(hostname, port) => write!(f, "metatls:{}:{}", hostname, port),
+            Self::MetaTls(addr) => write!(f, "metatls:{}", addr),
             Self::Local(index) => write!(f, "local:{}", index),
             Self::Sim(sim_addr) => write!(f, "sim:{}", sim_addr),
             Self::Unix(addr) => write!(f, "unix:{}", addr),
@@ -630,7 +693,7 @@ pub fn dial<M: RemoteMessage>(addr: ChannelAddr) -> Result<ChannelTx<M>, Channel
     let inner = match addr {
         ChannelAddr::Local(port) => ChannelTxKind::Local(local::dial(port)?),
         ChannelAddr::Tcp(addr) => ChannelTxKind::Tcp(net::tcp::dial(addr)),
-        ChannelAddr::MetaTls(host, port) => ChannelTxKind::MetaTls(net::meta::dial(host, port)),
+        ChannelAddr::MetaTls(meta_addr) => ChannelTxKind::MetaTls(net::meta::dial(meta_addr)?),
         ChannelAddr::Sim(sim_addr) => ChannelTxKind::Sim(sim::dial::<M>(sim_addr)?),
         ChannelAddr::Unix(path) => ChannelTxKind::Unix(net::unix::dial(path)),
     };
@@ -649,8 +712,8 @@ pub fn serve<M: RemoteMessage>(
             let (addr, rx) = net::tcp::serve::<M>(addr)?;
             Ok((addr, ChannelRxKind::Tcp(rx)))
         }
-        ChannelAddr::MetaTls(hostname, port) => {
-            let (addr, rx) = net::meta::serve::<M>(hostname, port)?;
+        ChannelAddr::MetaTls(meta_addr) => {
+            let (addr, rx) = net::meta::serve::<M>(meta_addr)?;
             Ok((addr, ChannelRxKind::MetaTls(rx)))
         }
         ChannelAddr::Unix(path) => {

hyperactor/src/channel/net.rs

@@ -1175,6 +1175,8 @@ pub enum ClientError {
     Io(ChannelAddr, std::io::Error),
     #[error("send {0}: serialize: {1}")]
     Serialize(ChannelAddr, bincode::ErrorKind),
+    #[error("invalid address: {0}")]
+    InvalidAddress(String),
 }
 
 #[derive(EnumAsInner)]
@@ -2175,14 +2177,23 @@ pub(crate) mod meta {
 
     #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ChannelError`.
     pub(crate) fn parse(addr_string: &str) -> Result<ChannelAddr, ChannelError> {
+        // Try to parse as a socket address first
+        if let Ok(socket_addr) = addr_string.parse::<SocketAddr>() {
+            return Ok(ChannelAddr::MetaTls(MetaTlsAddr::Socket(socket_addr)));
+        }
+
+        // Otherwise, parse as hostname:port
         // use right split to allow for ipv6 addresses where ":" is expected.
         let parts = addr_string.rsplit_once(":");
         match parts {
             Some((hostname, port_str)) => {
                 let Ok(port) = port_str.parse() else {
                     return Err(ChannelError::InvalidAddress(addr_string.to_string()));
                 };
-                Ok(ChannelAddr::MetaTls(hostname.to_string(), port))
+                Ok(ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                    hostname: hostname.to_string(),
+                    port,
+                }))
             }
             _ => Err(ChannelError::InvalidAddress(addr_string.to_string())),
         }
@@ -2314,7 +2325,10 @@ pub(crate) mod meta {
         type Stream = TlsStream<TcpStream>;
 
         fn dest(&self) -> ChannelAddr {
-            ChannelAddr::MetaTls(self.hostname.clone(), self.port)
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: self.hostname.clone(),
+                port: self.port,
+            })
         }
 
         async fn connect(&self) -> Result<Self::Stream, ClientError> {
@@ -2353,42 +2367,93 @@ pub(crate) mod meta {
         }
     }
 
-    pub fn dial<M: RemoteMessage>(hostname: Hostname, port: Port) -> NetTx<M> {
-        NetTx::new(MetaLink { hostname, port })
+    pub fn dial<M: RemoteMessage>(addr: MetaTlsAddr) -> Result<NetTx<M>, ClientError> {
+        match addr {
+            MetaTlsAddr::Host { hostname, port } => Ok(NetTx::new(MetaLink { hostname, port })),
+            MetaTlsAddr::Socket(_) => Err(ClientError::InvalidAddress(
+                "MetaTls clients require hostname/port for host identity, not socket addresses"
+                    .to_string(),
+            )),
+        }
     }
 
-    /// Serve the given address with hostname and port. If port 0 is provided,
-    /// dynamic port will be resolved and is available on the returned ServerHandle.
+    /// Serve the given address. If port 0 is provided in a Host address,
+    /// a dynamic port will be resolved and is available in the returned ChannelAddr.
+    /// For Host addresses, binds to all resolved socket addresses.
     pub fn serve<M: RemoteMessage>(
-        hostname: Hostname,
-        port: Port,
+        addr: MetaTlsAddr,
     ) -> Result<(ChannelAddr, NetRx<M>), ServerError> {
-        let mut addrs = (hostname.as_ref(), port).to_socket_addrs().map_err(|err| {
-            ServerError::Resolve(ChannelAddr::MetaTls(hostname.clone(), port), err)
-        })?;
-        let addr = addrs.next().ok_or(ServerError::Resolve(
-            ChannelAddr::MetaTls(hostname.clone(), port),
-            io::Error::other("no available socket addr"),
-        ))?;
-        let channel_addr = ChannelAddr::MetaTls(hostname.clone(), port);
+        match addr {
+            MetaTlsAddr::Host { hostname, port } => {
+                // Resolve all addresses for the hostname
+                let addrs: Vec<SocketAddr> = (hostname.as_ref(), port)
+                    .to_socket_addrs()
+                    .map_err(|err| {
+                        ServerError::Resolve(
+                            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                                hostname: hostname.clone(),
+                                port,
+                            }),
+                            err,
+                        )
+                    })?
+                    .collect();
+
+                if addrs.is_empty() {
+                    return Err(ServerError::Resolve(
+                        ChannelAddr::MetaTls(MetaTlsAddr::Host { hostname, port }),
+                        io::Error::other("no available socket addr"),
+                    ));
+                }
 
-        // Go by way of a std listener to avoid making this function async.
-        let std_listener = std::net::TcpListener::bind(addr)
-            .map_err(|err| ServerError::Listen(channel_addr.clone(), err))?;
-        std_listener
-            .set_nonblocking(true)
-            .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
-        let listener = TcpListener::from_std(std_listener)
-            .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
+                let channel_addr = ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                    hostname: hostname.clone(),
+                    port,
+                });
 
-        let local_addr = listener
-            .local_addr()
-            .map_err(|err| ServerError::Resolve(channel_addr, err))?;
-        super::serve(
-            listener,
-            ChannelAddr::MetaTls(hostname, local_addr.port()),
-            true,
-        )
+                // Bind to all resolved addresses
+                let std_listener = std::net::TcpListener::bind(&addrs[..])
+                    .map_err(|err| ServerError::Listen(channel_addr.clone(), err))?;
+                std_listener
+                    .set_nonblocking(true)
+                    .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
+                let listener = TcpListener::from_std(std_listener)
+                    .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
+
+                let local_addr = listener
+                    .local_addr()
+                    .map_err(|err| ServerError::Resolve(channel_addr, err))?;
+                super::serve(
+                    listener,
+                    ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                        hostname,
+                        port: local_addr.port(),
+                    }),
+                    true,
+                )
+            }
+            MetaTlsAddr::Socket(socket_addr) => {
+                let channel_addr = ChannelAddr::MetaTls(MetaTlsAddr::Socket(socket_addr));
+
+                // Bind directly to the socket address
+                let std_listener = std::net::TcpListener::bind(socket_addr)
+                    .map_err(|err| ServerError::Listen(channel_addr.clone(), err))?;
+                std_listener
+                    .set_nonblocking(true)
+                    .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
+                let listener = TcpListener::from_std(std_listener)
+                    .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
+
+                let local_addr = listener
+                    .local_addr()
+                    .map_err(|err| ServerError::Resolve(channel_addr, err))?;
+                super::serve(
+                    listener,
+                    ChannelAddr::MetaTls(MetaTlsAddr::Socket(local_addr)),
+                    true,
+                )
+            }
+        }
     }
 }
 
@@ -2569,11 +2634,11 @@ mod tests {
     #[tokio::test]
     async fn test_meta_tls_basic() {
         let addr = ChannelAddr::any(ChannelTransport::MetaTls(TlsMode::IpV6));
-        let (hostname, port) = match addr {
-            ChannelAddr::MetaTls(hostname, port) => (hostname, port),
-            _ => ("".to_string(), 0),
+        let meta_addr = match addr {
+            ChannelAddr::MetaTls(meta_addr) => meta_addr,
+            _ => panic!("expected MetaTls address"),
         };
-        let (local_addr, mut rx) = net::meta::serve::<u64>(hostname, port).unwrap();
+        let (local_addr, mut rx) = net::meta::serve::<u64>(meta_addr).unwrap();
         {
             let tx = dial::<u64>(local_addr.clone()).unwrap();
             tx.try_post(123, unused_return_channel()).unwrap();
@@ -3736,17 +3801,35 @@ mod tests {
     fn test_metatls_parsing() {
         // host:port
         let channel: ChannelAddr = "metatls!localhost:1234".parse().unwrap();
-        assert_eq!(channel, ChannelAddr::MetaTls("localhost".to_string(), 1234));
-        // ipv4:port
+        assert_eq!(
+            channel,
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: "localhost".to_string(),
+                port: 1234
+            })
+        );
+        // ipv4:port - can be parsed as hostname or socket address
         let channel: ChannelAddr = "metatls!1.2.3.4:1234".parse().unwrap();
-        assert_eq!(channel, ChannelAddr::MetaTls("1.2.3.4".to_string(), 1234));
+        assert_eq!(
+            channel,
+            ChannelAddr::MetaTls(MetaTlsAddr::Socket("1.2.3.4:1234".parse().unwrap()))
+        );
         // ipv6:port
         let channel: ChannelAddr = "metatls!2401:db00:33c:6902:face:0:2a2:0:1234"
             .parse()
             .unwrap();
         assert_eq!(
             channel,
-            ChannelAddr::MetaTls("2401:db00:33c:6902:face:0:2a2:0".to_string(), 1234)
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: "2401:db00:33c:6902:face:0:2a2:0".to_string(),
+                port: 1234
+            })
+        );
+
+        let channel: ChannelAddr = "metatls![::]:1234".parse().unwrap();
+        assert_eq!(
+            channel,
+            ChannelAddr::MetaTls(MetaTlsAddr::Socket("[::]:1234".parse().unwrap()))
         );
     }