Fix handling of IPv6 for Destination and DistantSingleKeyCredentials

CHANGELOG.md

@@ -6,6 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+## [0.17.4] - 2022-08-18
+### Fixed
+
+- Parsing of a host for `Destination` now correctly handles IPv6 addresses such
+  that `::1` and `[::1]:12345` are captured into host and port
+- Displaying of `Distant` and `DistantSingleKeyCredentials` now properly wrap
+  IPv6 addresses in square brackets when a port is available
 
 ## [0.17.3] - 2022-08-18
 ### Added

distant-core/src/credentials.rs

@@ -1,6 +1,6 @@
 use crate::{
     serde_str::{deserialize_from_str, serialize_to_str},
-    Destination,
+    Destination, Host,
 };
 use distant_net::SecretKey32;
 use serde::{de::Deserializer, ser::Serializer, Deserialize, Serialize};
@@ -13,7 +13,7 @@ const SCHEME_WITH_SEP: &str = "distant://";
 /// across all connections
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct DistantSingleKeyCredentials {
-    pub host: String,
+    pub host: Host,
     pub port: u16,
     pub key: SecretKey32,
     pub username: Option<String>,
@@ -23,10 +23,21 @@ impl fmt::Display for DistantSingleKeyCredentials {
     /// Converts credentials into string in the form of `distant://[username]:{key}@{host}:{port}`
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{SCHEME}://")?;
+
         if let Some(username) = self.username.as_ref() {
             write!(f, "{}", username)?;
         }
-        write!(f, ":{}@{}:{}", self.key, self.host, self.port)
+
+        write!(f, ":{}@", self.key)?;
+
+        // If we are IPv6, we need to include square brackets
+        if self.host.is_ipv6() {
+            write!(f, "[{}]", self.host)?;
+        } else {
+            write!(f, "{}", self.host)?;
+        }
+
+        write!(f, ":{}", self.port)
     }
 }
 
@@ -53,7 +64,7 @@ impl FromStr for DistantSingleKeyCredentials {
         }
 
         Ok(Self {
-            host: destination.host.to_string(),
+            host: destination.host,
             port: destination
                 .port
                 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Missing port"))?,
@@ -132,10 +143,7 @@ impl TryFrom<DistantSingleKeyCredentials> for Destination {
             scheme: Some("distant".to_string()),
             username: credentials.username,
             password: Some(credentials.key.to_string()),
-            host: credentials
-                .host
-                .parse()
-                .map_err(|x| io::Error::new(io::ErrorKind::InvalidData, x))?,
+            host: credentials.host,
             port: Some(credentials.port),
         })
     }
@@ -145,6 +153,7 @@ impl TryFrom<DistantSingleKeyCredentials> for Destination {
 mod tests {
     use super::*;
     use once_cell::sync::Lazy;
+    use std::net::{Ipv4Addr, Ipv6Addr};
 
     const HOST: &str = "testhost";
     const PORT: u16 = 12345;
@@ -239,4 +248,36 @@ mod tests {
         let credentials = DistantSingleKeyCredentials::find(&s);
         assert_eq!(credentials, None);
     }
+
+    #[test]
+    fn display_should_not_wrap_ipv4_address() {
+        let key = KEY.as_str();
+        let credentials = DistantSingleKeyCredentials {
+            host: Host::Ipv4(Ipv4Addr::LOCALHOST),
+            port: 12345,
+            username: None,
+            key: key.parse().unwrap(),
+        };
+
+        assert_eq!(
+            credentials.to_string(),
+            format!("{SCHEME}://:{key}@127.0.0.1:12345")
+        );
+    }
+
+    #[test]
+    fn display_should_wrap_ipv6_address_in_square_brackets() {
+        let key = KEY.as_str();
+        let credentials = DistantSingleKeyCredentials {
+            host: Host::Ipv6(Ipv6Addr::LOCALHOST),
+            port: 12345,
+            username: None,
+            key: key.parse().unwrap(),
+        };
+
+        assert_eq!(
+            credentials.to_string(),
+            format!("{SCHEME}://:{key}@[::1]:12345")
+        );
+    }
 }

distant-core/src/manager/data/destination.rs

@@ -86,7 +86,11 @@ impl fmt::Display for Destination {
             write!(f, "@")?;
         }
 
-        write!(f, "{}", self.host)?;
+        // For host, if we have a port and are IPv6, we need to wrap in [{}]
+        match &self.host {
+            Host::Ipv6(x) if self.port.is_some() => write!(f, "[{}]", x)?,
+            x => write!(f, "{}", x)?,
+        }
 
         if let Some(port) = self.port {
             write!(f, ":{port}")?;
@@ -154,4 +158,28 @@ mod tests {
         };
         assert_eq!(destination, "example.com");
     }
+
+    #[test]
+    fn display_should_not_wrap_ipv6_in_square_brackets_if_has_no_port() {
+        let destination = Destination {
+            scheme: None,
+            username: None,
+            password: None,
+            host: Host::Ipv6("::1".parse().unwrap()),
+            port: None,
+        };
+        assert_eq!(destination, "::1");
+    }
+
+    #[test]
+    fn display_should_wrap_ipv6_in_square_brackets_if_has_port() {
+        let destination = Destination {
+            scheme: None,
+            username: None,
+            password: None,
+            host: Host::Ipv6("::1".parse().unwrap()),
+            port: Some(12345),
+        };
+        assert_eq!(destination, "[::1]:12345");
+    }
 }

distant-core/src/manager/data/destination/host.rs

@@ -3,7 +3,7 @@ use derive_more::{Display, Error, From};
 use serde::{de::Deserializer, ser::Serializer, Deserialize, Serialize};
 use std::{
     fmt,
-    net::{Ipv4Addr, Ipv6Addr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
     str::FromStr,
 };
 
@@ -41,6 +41,30 @@ impl Host {
             Self::Name(_) => false,
         }
     }
+
+    /// Returns true if host is an IPv4 address
+    pub const fn is_ipv4(&self) -> bool {
+        matches!(self, Self::Ipv4(_))
+    }
+
+    /// Returns true if host is an IPv6 address
+    pub const fn is_ipv6(&self) -> bool {
+        matches!(self, Self::Ipv6(_))
+    }
+
+    /// Returns true if host is a name
+    pub const fn is_name(&self) -> bool {
+        matches!(self, Self::Name(_))
+    }
+}
+
+impl From<IpAddr> for Host {
+    fn from(addr: IpAddr) -> Self {
+        match addr {
+            IpAddr::V4(x) => Self::Ipv4(x),
+            IpAddr::V6(x) => Self::Ipv6(x),
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug, Error, Hash, PartialEq, Eq)]

distant-core/src/manager/data/destination/parser.rs

@@ -7,7 +7,32 @@ type PError = &'static str;
 pub fn parse(s: &str) -> Result<Destination, &'static str> {
     let (s, scheme) = maybe(parse_scheme)(s)?;
     let (s, username_password) = maybe(parse_username_password)(s)?;
-    let (s, host) = parse_and_then(parse_until(|c| c == ':'), parse_host)(s)?;
+
+    // NOTE: We can have a host or host/port in a couple of different ways
+    //
+    // 1. IPv4 - 127.0.0.1 or 127.0.0.1:1234
+    // 2. IPv6 - ::1 or [::1]:1234
+    // 3. Name - localhost or localhost:1234
+    //
+    // To determine path to take, we count the colons. If there is more than 1, we can assume IPv6
+    // is involved and can try to parse entirely as IPv6, or if that fails split off one colon and
+    // try a second time. Otherwise, we assume that it is not IPv6 and can parse with a single
+    // colon at most for a port.
+    let colon_cnt = s.chars().filter(|c| *c == ':').count();
+    let (s, host) = if colon_cnt > 1 {
+        // Either the host is [{}] with a port following, or the host is everything
+        either(
+            delimited(
+                parse_char('['),
+                parse_and_then(parse_until(|c| c == ']'), parse_host),
+                parse_char(']'),
+            ),
+            parse_host,
+        )(s)?
+    } else {
+        parse_and_then(parse_until(|c| c == ':'), parse_host)(s)?
+    };
+
     let (s, port) = maybe(prefixed(parse_char(':'), parse_port))(s)?;
 
     if !s.is_empty() {
@@ -70,6 +95,34 @@ fn parse_port(s: &str) -> PResult<u16> {
     Ok(("", port))
 }
 
+/// Execute parsers in order from left to right, returning the result of the first that succeeds
+fn either<'a, T>(
+    left: impl Fn(&'a str) -> PResult<'a, T>,
+    right: impl Fn(&'a str) -> PResult<'a, T>,
+) -> impl Fn(&'a str) -> PResult<'a, T> {
+    move |s: &str| {
+        if let Ok((s, value)) = left(s) {
+            Ok((s, value))
+        } else {
+            right(s)
+        }
+    }
+}
+
+/// Execute three parsers in a row, failing if any fails, and returns second parser's result
+fn delimited<'a, T1, T2, T3>(
+    p1: impl Fn(&'a str) -> PResult<'a, T1>,
+    p2: impl Fn(&'a str) -> PResult<'a, T2>,
+    p3: impl Fn(&'a str) -> PResult<'a, T3>,
+) -> impl Fn(&'a str) -> PResult<'a, T2> {
+    move |s: &str| {
+        let (s, _) = p1(s)?;
+        let (s, value) = p2(s)?;
+        let (s, _) = p3(s)?;
+        Ok((s, value))
+    }
+}
+
 /// Execute two parsers in a row, failing if either fails, and returns second parser's result
 fn prefixed<'a, T1, T2>(
     prefix_parser: impl Fn(&'a str) -> PResult<'a, T1>,
@@ -619,6 +672,46 @@ mod tests {
             assert_eq!(destination.port, Some(22));
         }
 
+        #[test]
+        fn parse_should_succeed_if_given_ipv4_host() {
+            let destination = parse("127.0.0.1").unwrap();
+            assert_eq!(destination.scheme, None);
+            assert_eq!(destination.username, None);
+            assert_eq!(destination.password, None);
+            assert_eq!(destination.host, "127.0.0.1");
+            assert_eq!(destination.port, None);
+        }
+
+        #[test]
+        fn parse_should_succeed_if_given_ipv4_host_and_port() {
+            let destination = parse("127.0.0.1:12345").unwrap();
+            assert_eq!(destination.scheme, None);
+            assert_eq!(destination.username, None);
+            assert_eq!(destination.password, None);
+            assert_eq!(destination.host, "127.0.0.1");
+            assert_eq!(destination.port, Some(12345));
+        }
+
+        #[test]
+        fn parse_should_succeed_if_given_ipv6_host() {
+            let destination = parse("::1").unwrap();
+            assert_eq!(destination.scheme, None);
+            assert_eq!(destination.username, None);
+            assert_eq!(destination.password, None);
+            assert_eq!(destination.host, "::1");
+            assert_eq!(destination.port, None);
+        }
+
+        #[test]
+        fn parse_should_succeed_if_given_ipv6_host_and_port() {
+            let destination = parse("[::1]:12345").unwrap();
+            assert_eq!(destination.scheme, None);
+            assert_eq!(destination.username, None);
+            assert_eq!(destination.password, None);
+            assert_eq!(destination.host, "::1");
+            assert_eq!(destination.port, Some(12345));
+        }
+
         #[test]
         fn parse_should_succeed_with_distant_server_output() {
             // This is an example of what a server might output that includes a 32-byte key

distant-ssh2/src/lib.rs

@@ -643,7 +643,9 @@ impl Ssh {
             trace!("Searching for credentials");
             match DistantSingleKeyCredentials::find(&String::from_utf8_lossy(&output.stdout)) {
                 Some(mut info) => {
-                    info.host = host;
+                    info.host = host
+                        .parse()
+                        .map_err(|x| io::Error::new(io::ErrorKind::InvalidData, x))?;
                     Ok(info)
                 }
                 None => Err(io::Error::new(

src/cli/commands/server.rs

@@ -6,7 +6,7 @@ use anyhow::Context;
 use clap::Subcommand;
 use distant_core::{
     net::{SecretKey32, ServerRef, TcpServerExt, XChaCha20Poly1305Codec},
-    DistantApiServer, DistantSingleKeyCredentials,
+    DistantApiServer, DistantSingleKeyCredentials, Host,
 };
 use log::*;
 use std::io::{self, Read, Write};
@@ -177,7 +177,7 @@ impl ServerSubcommand {
                     })?;
 
                 let credentials = DistantSingleKeyCredentials {
-                    host: addr.to_string(),
+                    host: Host::from(addr),
                     port: server.port(),
                     key,
                     username: None,