add test to cover err scenario

Author: jbesraa

lightning/src/ln/msgs.rs

@@ -896,7 +896,7 @@ impl Readable for NetAddress {
 
 /// NetAddress error variants
 #[cfg(feature = "std")]
-#[derive(Debug)]
+#[derive(Debug, Eq, PartialEq, Clone)]
 pub enum NetAddressError {
 	/// Socket address(IPv4/IPv6) parsing error
 	SocketAdrrParseError(std::net::AddrParseError),
@@ -910,24 +910,24 @@ pub enum NetAddressError {
 
 #[cfg(feature = "std")]
 impl FromStr for NetAddress {
-    type Err = NetAddressError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match std::net::SocketAddr::from_str(s) {
-            Ok(addr) => {
-                let port: u16 = addr.port();
-                match addr {
-                    std::net::SocketAddr::V4(addr) => {
-                        let addr = addr.ip().octets();
-                        return Ok(NetAddress::IPv4 { addr, port });
-                    },
-                    std::net::SocketAddr::V6(addr) => {
-                        let addr = addr.ip().octets();
-                        return Ok(NetAddress::IPv6 { addr, port });
-                    },
-                }
-            },
-            Err(e) => {
+	type Err = NetAddressError;
+
+	fn from_str(s: &str) -> Result<Self, Self::Err> {
+		match std::net::SocketAddr::from_str(s) {
+			Ok(addr) => {
+				let port: u16 = addr.port();
+				match addr {
+					std::net::SocketAddr::V4(addr) => {
+						let addr = addr.ip().octets();
+						return Ok(NetAddress::IPv4 { addr, port });
+					},
+					std::net::SocketAddr::V6(addr) => {
+						let addr = addr.ip().octets();
+						return Ok(NetAddress::IPv6 { addr, port });
+					},
+				}
+			},
+			Err(e) => {
 				let trimmed_input = match s.rfind(":") {
 					Some(pos) => pos,
 					None => return Err(NetAddressError::InvalidInput("Invalid input. Expected format: \"<host>:<port>\"" .to_string())),
@@ -937,36 +937,48 @@ impl FromStr for NetAddress {
 					Ok(port) => port,
 					Err(_) => return Err(NetAddressError::InvalidPort),
 				};
-				if host.ends_with(".onion") {
-					let onion = match host.split(".onion").nth(0) {
-						Some(onion) => onion,
-						None => return Err(NetAddressError::InvalidOnionV3),
-					};
-					let onion = match (base32::Alphabet::RFC4648 { padding: false }.decode(&onion)) {
-						Ok(onion) => onion,
-						Err(_) => return Err(NetAddressError::InvalidOnionV3),
-					};
-					match (onion.get(0), onion.get(1), onion.get(2)) {
-						(Some(version), Some(first_checksum_flag), Some(second_checksum_flag)) => {
-							let checksum = u16::from_be_bytes([*first_checksum_flag, *second_checksum_flag]);
-							let ed25519_pubkey = match onion[3..35].try_into() {
-								Ok(ed25519_pubkey) => ed25519_pubkey,
-								Err(_) => return Err(NetAddressError::InvalidOnionV3),
-							};
-							return Ok(NetAddress::OnionV3 { ed25519_pubkey, checksum, version: *version, port });
-
+				let host: Vec<&str> = host.split(".").collect();
+				if host.len() == 2 {
+					match (host.get(0), host.get(1)) {
+						(Some(domain), Some(suffix)) => {
+							if suffix.to_string() == "onion".to_owned() {
+								if domain.len() != 56 {
+									return Err(NetAddressError::InvalidOnionV3);
+								}
+								let onion = match (base32::Alphabet::RFC4648 { padding: false }.decode(&domain)) {
+									Ok(onion) => onion,
+									Err(_) => return Err(NetAddressError::InvalidOnionV3),
+								};
+								if onion.len() < 35 {
+									return Err(NetAddressError::InvalidOnionV3);
+								}
+								match (onion.get(0), onion.get(1), onion.get(2)) {
+									(Some(version), Some(first_checksum_flag), Some(second_checksum_flag)) => {
+										let checksum = u16::from_be_bytes([*first_checksum_flag, *second_checksum_flag]);
+										let ed25519_pubkey = match onion[3..35].try_into() {
+											Ok(ed25519_pubkey) => ed25519_pubkey,
+											Err(_) => return Err(NetAddressError::InvalidOnionV3),
+										};
+										return Ok(NetAddress::OnionV3 { ed25519_pubkey, checksum, version: *version, port });
+
+									},
+									_ => return Err(NetAddressError::InvalidOnionV3),
+								}
+							} else {
+								if let Ok(hostname) = Hostname::try_from(s[..trimmed_input].to_string()) {
+									return Ok(NetAddress::Hostname { hostname, port });
+								} else {
+									return Err(NetAddressError::InvalidInput("Invalid input. Expected format: \"<host>:<port>\"" .to_string()));
+								}
+							}
 						},
-						_ => return Err(NetAddressError::InvalidOnionV3),
-					}
-                }
-
-                if let Ok(hostname) = Hostname::try_from(host.to_string()) {
-                    return Ok(NetAddress::Hostname { hostname, port });
-                }
-                return Err(NetAddressError::SocketAdrrParseError(e))
-            },
-        }
-    }
+						_ => return Err(NetAddressError::SocketAdrrParseError(e))
+					};
+				}
+				return Err(NetAddressError::SocketAdrrParseError(e))
+			},
+		}
+	}
 }
 
 /// Represents the set of gossip messages that require a signature from a node's identity key.
@@ -3816,21 +3828,35 @@ mod tests {
 			addr: Ipv4Addr::new(127, 0, 0, 1).octets(),
 			port: 1234,
 		}, NetAddress::from_str("127.0.0.1:1234").unwrap());
+
 		assert_eq!(NetAddress::IPv6 {
 			addr : Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1).octets(),
 			port: 1234,
 		}, NetAddress::from_str("[0:0:0:0:0:0:0:1]:1234").unwrap());
+
 		assert_eq!(
 			NetAddress::Hostname {
 				hostname : Hostname::try_from("example.com".to_string()).unwrap(),
 				port: 1234,
 			} , NetAddress::from_str("example.com:1234").unwrap());
+
 		assert_eq!(NetAddress::OnionV3 {
 			ed25519_pubkey: [37, 24, 75, 5, 25, 73, 117, 194, 139, 102, 182, 107, 4, 105, 247, 246, 85,
 			111, 177, 172, 49, 137, 167, 155, 64, 221, 163, 47, 31, 33, 71, 3],
 			checksum: 48326,
 			version: 121,
 			port: 1234
 		}, NetAddress::from_str( "pg6mmjiyjmcrsslvykfwnntlaru7p5svn6y2ymmju6nubxndf4pscryd.onion:1234").unwrap());
+
+		let invalid_onion_address = "pg6mmjiyjmcrsslvykfwnntlaru7p5svn6.onion:1234";
+		assert_eq!(super::NetAddressError::InvalidOnionV3, NetAddress::from_str(invalid_onion_address).unwrap_err().into());
+
+		assert_eq!(super::NetAddressError::InvalidInput("Invalid input. Expected format: \"<host>:<port>\"" .to_string()), NetAddress::from_str("127.0.0.1@1234").unwrap_err().into());
+
+		assert!(NetAddress::from_str("pg6mmjiyjmcrsslvykfwnntlaru7p5svn6y2ymmju6nubxndf4pscryd.onion.onion:9735:94").is_err());
+
+		assert_eq!(super::NetAddressError::InvalidInput("Invalid input. Expected format: \"<host>:<port>\"" .to_string()), NetAddress::from_str("wrong$%#.com:1234").unwrap_err().into());
+
+		assert_eq!(super::NetAddressError::InvalidPort, NetAddress::from_str("example.com:wrong").unwrap_err().into());
 	}
 }