refactor from_str fun

lightningdevkit · Jun 28, 2023 · d74c87f · d74c87f
1 parent 90ae34c
commit d74c87f
Showing 1 changed file with 34 additions and 50 deletions.
diff --git a/lightning/src/ln/msgs.rs b/lightning/src/ln/msgs.rs
@@ -904,6 +904,7 @@ impl Readable for NetAddress {
 	}
 }
 
+
 /// NetAddress error variants
 #[cfg(feature = "std")]
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -918,25 +919,22 @@ pub enum NetAddressError {
 	InvalidOnionV3,
 }
 
+impl From<std::net::SocketAddr> for NetAddress {
+    fn from(addr: std::net::SocketAddr) -> Self {
+        match addr {
+            std::net::SocketAddr::V4(addr) => NetAddress::IPv4 { addr: addr.ip().octets(), port: addr.port() },
+            std::net::SocketAddr::V6(addr) => NetAddress::IPv6 { addr: addr.ip().octets(), port: addr.port() },
+        }
+    }
+}
+
 #[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 });
-					},
-				}
-			},
+			Ok(addr) => Ok(addr.into()),
 			Err(e) => {
 				let trimmed_input = match s.rfind(":") {
 					Some(pos) => pos,
@@ -947,44 +945,30 @@ impl FromStr for NetAddress {
 					Ok(port) => port,
 					Err(_) => return Err(NetAddressError::InvalidPort),
 				};
-				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::SocketAdrrParseError(e))
+				if host.ends_with(".onion") {
+					let domain = host.trim_end_matches(".onion");
+					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);
+					}
+					let version = onion[0];
+					let first_checksum_flag = onion[1];
+					let second_checksum_flag = onion[2];
+					let mut ed25519_pubkey = [0;32];
+					ed25519_pubkey.copy_from_slice(&onion[3..35]);
+					let checksum = u16::from_be_bytes([first_checksum_flag, second_checksum_flag]);
+					return Ok(NetAddress::OnionV3 { ed25519_pubkey, checksum, version, port });
+
+				};
+				if let Ok(hostname) = Hostname::try_from(s[..trimmed_input].to_string()) {
+					return Ok(NetAddress::Hostname { hostname, port });
+				};
 				return Err(NetAddressError::SocketAdrrParseError(e))
 			},
 		}