feat(net): Remove NetworkAddress (#55)

### Description Removes the `NetworkAddress` which is effectively used as a conversion utility between the [`SocketAddr`](https://doc.rust-lang.org/stable/std/net/enum.SocketAddr.html) and [`MulitAddr`](https://docs.rs/libp2p/latest/libp2p/struct.Multiaddr.html). Closes #52 --------- Co-authored-by: nicolas <[email protected]>
ithacaxyz · Aug 28, 2024 · 9c08589 · 9c08589
1 parent 1482685
commit 9c08589
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 111 deletions.
diff --git a/crates/net/src/builder.rs b/crates/net/src/builder.rs
@@ -2,20 +2,19 @@
 
 use alloy::primitives::Address;
 use eyre::Result;
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use tokio::sync::watch::channel;
 
 use libp2p::{
-    gossipsub::Config as GossipConfig, noise::Config as NoiseConfig, tcp::Config as TcpConfig,
-    yamux::Config as YamuxConfig, Multiaddr, SwarmBuilder,
+    gossipsub::Config as GossipConfig, multiaddr::Protocol, noise::Config as NoiseConfig,
+    tcp::Config as TcpConfig, yamux::Config as YamuxConfig, Multiaddr, SwarmBuilder,
 };
 use libp2p_identity::Keypair;
 
 use crate::{
     discovery::builder::DiscoveryBuilder,
     driver::NetworkDriver,
     gossip::{behaviour::Behaviour, config, driver::GossipDriver, handler::BlockHandler},
-    types::address::NetworkAddress,
 };
 
 /// Constructs a [NetworkDriver] for Optimism's consensus-layer.
@@ -180,14 +179,18 @@ impl NetworkDriverBuilder {
             )?
             .with_behaviour(|_| behaviour)?
             .build();
-        let addr = self.socket.take().ok_or_else(|| eyre::eyre!("socket address not set"))?;
-        let addr = NetworkAddress::try_from(addr)?;
-        let swarm_addr = Multiaddr::from(addr);
-        let gossip = GossipDriver::new(swarm, swarm_addr, handler);
+        let socket = self.socket.take().ok_or(eyre::eyre!("socket address not set"))?;
+        let mut multiaddr = Multiaddr::empty();
+        match socket.ip() {
+            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
+            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
+        }
+        multiaddr.push(Protocol::Tcp(socket.port()));
+        let gossip = GossipDriver::new(swarm, multiaddr, handler);
 
         // Build the discovery service
         let discovery =
-            DiscoveryBuilder::new().with_address(addr).with_chain_id(chain_id).build()?;
+            DiscoveryBuilder::new().with_address(socket).with_chain_id(chain_id).build()?;
 
         Ok(NetworkDriver { unsafe_block_recv, unsafe_block_signer_sender, gossip, discovery })
     }
@@ -240,11 +243,15 @@ mod tests {
             .with_gossip_config(cfg)
             .build()
             .unwrap();
-        let signer_net_addr = NetworkAddress::try_from(socket).expect("network address");
-        let signer_multiaddr = Multiaddr::from(signer_net_addr);
+        let mut multiaddr = Multiaddr::empty();
+        match socket.ip() {
+            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
+            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
+        }
+        multiaddr.push(Protocol::Tcp(socket.port()));
 
         // Driver Assertions
-        assert_eq!(driver.gossip.addr, signer_multiaddr);
+        assert_eq!(driver.gossip.addr, multiaddr);
         assert_eq!(driver.discovery.chain_id, id);
 
         // Block Handler Assertions
@@ -268,11 +275,15 @@ mod tests {
             .with_socket(socket)
             .build()
             .unwrap();
-        let signer_net_addr = NetworkAddress::try_from(socket).expect("network address");
-        let signer_multiaddr = Multiaddr::from(signer_net_addr);
+        let mut multiaddr = Multiaddr::empty();
+        match socket.ip() {
+            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
+            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
+        }
+        multiaddr.push(Protocol::Tcp(socket.port()));
 
         // Driver Assertions
-        assert_eq!(driver.gossip.addr, signer_multiaddr);
+        assert_eq!(driver.gossip.addr, multiaddr);
         assert_eq!(driver.discovery.chain_id, id);
 
         // Block Handler Assertions

diff --git a/crates/net/src/discovery/builder.rs b/crates/net/src/discovery/builder.rs
@@ -1,22 +1,20 @@
 //! Contains a builder for the discovery service.
 
-use crate::{
-    discovery::driver::DiscoveryDriver,
-    types::{address::NetworkAddress, enr::OpStackEnr},
-};
+use crate::{discovery::driver::DiscoveryDriver, types::enr::OpStackEnr};
 use discv5::{
     enr::{CombinedKey, Enr},
     ConfigBuilder, Discv5, ListenConfig,
 };
 use eyre::Result;
+use std::net::SocketAddr;
 
 use crate::types::enr::OP_CL_KEY;
 
 /// Discovery service builder.
 #[derive(Debug, Default, Clone)]
 pub struct DiscoveryBuilder {
     /// The discovery service address.
-    address: Option<NetworkAddress>,
+    address: Option<SocketAddr>,
     /// The chain ID of the network.
     chain_id: Option<u64>,
 }
@@ -28,7 +26,7 @@ impl DiscoveryBuilder {
     }
 
     /// Sets the discovery service address.
-    pub fn with_address(mut self, address: NetworkAddress) -> Self {
+    pub fn with_address(mut self, address: SocketAddr) -> Self {
         self.address = Some(address);
         self
     }
@@ -48,7 +46,7 @@ impl DiscoveryBuilder {
 
         let key = CombinedKey::generate_secp256k1();
         let enr = Enr::builder().add_value_rlp(OP_CL_KEY, opstack_data.into()).build(&key)?;
-        let listen_config = ListenConfig::from_ip(addr.ip.into(), addr.port);
+        let listen_config = ListenConfig::from_ip(addr.ip(), addr.port());
         let config = ConfigBuilder::new(listen_config).build();
 
         let disc = Discv5::new(enr, key, config)

diff --git a/crates/net/src/discovery/driver.rs b/crates/net/src/discovery/driver.rs
@@ -12,7 +12,7 @@ use discv5::{enr::NodeId, Discv5};
 
 use crate::{
     discovery::{bootnodes::BOOTNODES, builder::DiscoveryBuilder},
-    types::{address::Peer, enr::OpStackEnr},
+    types::{enr::OpStackEnr, peer::Peer},
 };
 
 /// The number of peers to buffer in the channel.
@@ -49,14 +49,14 @@ impl DiscoveryDriver {
     /// ## Example
     ///
     /// ```no_run
-    /// use op_net::{discovery::builder::DiscoveryBuilder, types::address::NetworkAddress};
-    /// use std::net::Ipv4Addr;
+    /// use op_net::discovery::builder::DiscoveryBuilder;
+    /// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
     ///
     /// #[tokio::main]
     /// async fn main() {
-    ///     let network_addr = NetworkAddress { ip: Ipv4Addr::new(127, 0, 0, 1), port: 9000 };
+    ///     let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9099);
     ///     let mut discovery = DiscoveryBuilder::new()
-    ///         .with_address(network_addr)
+    ///         .with_address(socket)
     ///         .with_chain_id(10) // OP Mainnet chain id
     ///         .build()
     ///         .expect("Failed to build discovery service");

diff --git a/crates/net/src/types/address.rs b/crates/net/src/types/address.rs
diff --git a/crates/net/src/types/mod.rs b/crates/net/src/types/mod.rs
@@ -1,6 +1,6 @@
 //! Common types for the Networking Crate.
 
-pub mod address;
 pub mod enr;
 pub mod envelope;
 pub mod payload;
+pub mod peer;
diff --git a/crates/net/src/types/peer.rs b/crates/net/src/types/peer.rs
@@ -0,0 +1,38 @@
+//! Peer Types
+
+use discv5::enr::{CombinedKey, Enr};
+use eyre::Result;
+use libp2p::{multiaddr::Protocol, Multiaddr};
+use std::net::{IpAddr, SocketAddr};
+
+/// A wrapper around a peer's [SocketAddr].
+#[derive(Debug)]
+pub struct Peer {
+    /// The peer's [SocketAddr].
+    pub socket: SocketAddr,
+}
+
+impl TryFrom<&Enr<CombinedKey>> for Peer {
+    type Error = eyre::Report;
+
+    /// Converts an [Enr] to a Peer
+    fn try_from(value: &Enr<CombinedKey>) -> Result<Self> {
+        let ip = value.ip4().ok_or(eyre::eyre!("missing ip"))?;
+        let port = value.tcp4().ok_or(eyre::eyre!("missing port"))?;
+        let socket = SocketAddr::new(IpAddr::V4(ip), port);
+        Ok(Peer { socket })
+    }
+}
+
+impl From<Peer> for Multiaddr {
+    /// Converts a Peer to a [Multiaddr]
+    fn from(value: Peer) -> Self {
+        let mut multiaddr = Multiaddr::empty();
+        match value.socket.ip() {
+            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
+            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
+        }
+        multiaddr.push(Protocol::Tcp(value.socket.port()));
+        multiaddr
+    }
+}