From 889fa89d5e73b6145d49d0ed7962b9de4ce390ae Mon Sep 17 00:00:00 2001
From: Aditya Pandey <prblmslvr.aditya@gmail.com>
Date: Mon, 27 Mar 2023 17:31:45 +0530
Subject: [PATCH] Fixing parse socket address function (#1990)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 bin/reth/src/node/mod.rs            | 13 +++++++
 crates/staged-sync/src/utils/mod.rs | 59 ++++++++++++++++++++++-------
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/bin/reth/src/node/mod.rs b/bin/reth/src/node/mod.rs
index 90e61f5f6..e02d1a383 100644
--- a/bin/reth/src/node/mod.rs
+++ b/bin/reth/src/node/mod.rs
@@ -623,6 +623,7 @@ async fn run_network_until_shutdown<C>(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::net::IpAddr;
 
     #[test]
     fn parse_help_node_command() {
@@ -651,4 +652,16 @@ mod tests {
         assert_eq!(cmd.network.discovery.port, Some(300));
         assert_eq!(cmd.network.port, Some(99));
     }
+
+    #[test]
+    fn parse_metrics_port() {
+        let cmd = Command::try_parse_from(["reth", "--metrics", "9000"]).unwrap();
+        assert_eq!(cmd.metrics, Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9000)));
+
+        let cmd = Command::try_parse_from(["reth", "--metrics", ":9000"]).unwrap();
+        assert_eq!(cmd.metrics, Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9000)));
+
+        let cmd = Command::try_parse_from(["reth", "--metrics", "localhost:9000"]).unwrap();
+        assert_eq!(cmd.metrics, Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9000)));
+    }
 }
diff --git a/crates/staged-sync/src/utils/mod.rs b/crates/staged-sync/src/utils/mod.rs
index 1c0016bea..c189bb2b5 100644
--- a/crates/staged-sync/src/utils/mod.rs
+++ b/crates/staged-sync/src/utils/mod.rs
@@ -2,7 +2,7 @@
 use reth_primitives::{BlockHashOrNumber, H256};
 use std::{
     env::VarError,
-    net::{SocketAddr, ToSocketAddrs},
+    net::{IpAddr, Ipv4Addr, SocketAddr, ToSocketAddrs},
     path::{Path, PathBuf},
     str::FromStr,
 };
@@ -38,6 +38,23 @@ pub fn hash_or_num_value_parser(value: &str) -> Result<BlockHashOrNumber, eyre::
     }
 }
 
+/// Error thrown while parsing a socket address.
+#[derive(thiserror::Error, Debug)]
+pub enum SocketAddressParsingError {
+    /// Failed to convert the string into a socket addr
+    #[error("Cannot parse socket address: {0}")]
+    Io(#[from] std::io::Error),
+    /// Input must not be empty
+    #[error("Cannot parse socket address from empty string")]
+    Empty,
+    /// Failed to parse the address
+    #[error("Could not parse socket address from {0}")]
+    Parse(String),
+    /// Failed to parse port
+    #[error("Could not parse port: {0}")]
+    Port(#[from] std::num::ParseIntError),
+}
+
 /// Parse a [SocketAddr] from a `str`.
 ///
 /// The following formats are checked:
@@ -48,29 +65,32 @@ pub fn hash_or_num_value_parser(value: &str) -> Result<BlockHashOrNumber, eyre::
 /// - Otherwise it is assumed to be a hostname
 ///
 /// An error is returned if the value is empty.
-pub fn parse_socket_address(value: &str) -> Result<SocketAddr, eyre::Error> {
+pub fn parse_socket_address(value: &str) -> Result<SocketAddr, SocketAddressParsingError> {
     if value.is_empty() {
-        eyre::bail!("Cannot parse socket address from an empty string");
+        return Err(SocketAddressParsingError::Empty)
     }
 
-    if value.starts_with(':') || value.parse::<u16>().is_ok() {
-        ("localhost", 9000).to_socket_addrs()
-    } else if value.contains(':') {
-        value.to_socket_addrs()
-    } else {
-        (value, 9000).to_socket_addrs()
-    }?
-    .next()
-    .ok_or_else(|| eyre::eyre!("Could not parse socket address from {}", value))
+    if let Some(port) = value.strip_prefix(':').or_else(|| value.strip_prefix("localhost:")) {
+        let port: u16 = port.parse()?;
+        return Ok(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port))
+    }
+    if let Ok(port) = value.parse::<u16>() {
+        return Ok(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port))
+    }
+    value
+        .to_socket_addrs()?
+        .next()
+        .ok_or_else(|| SocketAddressParsingError::Parse(value.to_string()))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rand::{thread_rng, Rng};
 
     #[test]
     fn parse_socket_addresses() {
-        for value in ["localhost:9000", ":9000", "9000", "localhost"] {
+        for value in ["localhost:9000", ":9000", "9000"] {
             let socket_addr = parse_socket_address(value)
                 .unwrap_or_else(|_| panic!("could not parse socket address: {value}"));
 
@@ -78,4 +98,17 @@ mod tests {
             assert_eq!(socket_addr.port(), 9000);
         }
     }
+
+    #[test]
+    fn parse_socket_address_random() {
+        let port: u16 = thread_rng().gen();
+
+        for value in [format!("localhost:{port}"), format!(":{port}"), port.to_string()] {
+            let socket_addr = parse_socket_address(&value)
+                .unwrap_or_else(|_| panic!("could not parse socket address: {value}"));
+
+            assert!(socket_addr.ip().is_loopback());
+            assert_eq!(socket_addr.port(), port);
+        }
+    }
 }