From 9c266e7fd2d9880a49f2a6d86a8bb86836d17049 Mon Sep 17 00:00:00 2001
From: Yilin Chen <sticnarf@gmail.com>
Date: Thu, 5 Sep 2019 11:14:42 +0800
Subject: [PATCH 1/2] Implement ToTargetAddr for more

Signed-off-by: Yilin Chen <sticnarf@gmail.com>
---
 src/lib.rs | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 4bc60ae..1a180f9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -154,6 +154,12 @@ trivial_impl_into_target_addr!((Ipv6Addr, u16));
 trivial_impl_into_target_addr!(SocketAddrV4);
 trivial_impl_into_target_addr!(SocketAddrV6);
 
+impl<'a> IntoTargetAddr<'a> for TargetAddr<'a> {
+    fn into_target_addr(self) -> Result<TargetAddr<'a>> {
+        Ok(self)
+    }
+}
+
 impl<'a> IntoTargetAddr<'a> for (&'a str, u16) {
     fn into_target_addr(self) -> Result<TargetAddr<'a>> {
         // Try IP address first
@@ -187,7 +193,28 @@ impl<'a> IntoTargetAddr<'a> for &'a str {
         let domain = parts_iter
             .next()
             .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
-        (domain, port).into_target_addr()
+        Ok(TargetAddr::Domain(domain.into(), port))
+    }
+}
+
+impl IntoTargetAddr<'static> for String {
+    fn into_target_addr(mut self) -> Result<TargetAddr<'static>> {
+        // Try IP address first
+        if let Ok(addr) = self.parse::<SocketAddr>() {
+            return addr.into_target_addr();
+        }
+
+        let mut parts_iter = self.rsplitn(2, ':');
+        let port: u16 = parts_iter
+            .next()
+            .and_then(|port_str| port_str.parse().ok())
+            .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
+        let domain_len = parts_iter
+            .next()
+            .ok_or(Error::InvalidTargetAddress("invalid address format"))?
+            .len();
+        self.truncate(domain_len);
+        Ok(TargetAddr::Domain(self.into(), port))
     }
 }
 
@@ -317,6 +344,12 @@ mod tests {
             TargetAddr::Domain(Cow::Borrowed("www.example.com"), 80),
             res
         );
+
+        let res = into_target_addr(domain.to_owned())?;
+        assert_eq!(
+            TargetAddr::Domain(Cow::Owned("www.example.com".to_owned()), 80),
+            res
+        );
         Ok(())
     }
 

From e435b5ea6fa84d8aff2cd40b358382daf6939f3c Mon Sep 17 00:00:00 2001
From: Yilin Chen <sticnarf@gmail.com>
Date: Thu, 5 Sep 2019 11:23:15 +0800
Subject: [PATCH 2/2] Check domain length

Signed-off-by: Yilin Chen <sticnarf@gmail.com>
---
 src/lib.rs | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 1a180f9..cfebd73 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -168,8 +168,7 @@ impl<'a> IntoTargetAddr<'a> for (&'a str, u16) {
         }
 
         // Treat as domain name
-        let len = self.0.as_bytes().len();
-        if len > 255 {
+        if self.0.len() > 255 {
             return Err(Error::InvalidTargetAddress("overlong domain"));
         }
         // TODO: Should we validate the domain format here?
@@ -193,6 +192,9 @@ impl<'a> IntoTargetAddr<'a> for &'a str {
         let domain = parts_iter
             .next()
             .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
+        if domain.len() > 255 {
+            return Err(Error::InvalidTargetAddress("overlong domain"));
+        }
         Ok(TargetAddr::Domain(domain.into(), port))
     }
 }
@@ -213,6 +215,9 @@ impl IntoTargetAddr<'static> for String {
             .next()
             .ok_or(Error::InvalidTargetAddress("invalid address format"))?
             .len();
+        if domain_len > 255 {
+            return Err(Error::InvalidTargetAddress("overlong domain"));
+        }
         self.truncate(domain_len);
         Ok(TargetAddr::Domain(self.into(), port))
     }