fix(hera): Networking

bin/hera/src/network.rs

@@ -3,7 +3,7 @@
 use crate::globals::GlobalArgs;
 use clap::Args;
 use eyre::Result;
-use op_net::driver::NetworkDriver;
+use op_net::{discovery::builder::DiscoveryBuilder, driver::NetworkDriver};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use superchain_registry::ROLLUP_CONFIGS;
 
@@ -12,11 +12,8 @@ use superchain_registry::ROLLUP_CONFIGS;
 #[non_exhaustive]
 pub struct NetworkCommand {
     /// Run the peer discovery service.
-    #[clap(long, short = 'p', help = "Runs peer discovery")]
-    pub disc: bool,
-    /// Run the gossip driver.
-    #[clap(long, short = 'g', help = "Runs the unsafe block gossipping service")]
-    pub gossip: bool,
+    #[clap(long, short = 'p', help = "Runs only peer discovery")]
+    pub only_disc: bool,
     /// Port to listen for gossip on.
     #[clap(long, short = 'l', default_value = "9099", help = "Port to listen for gossip on")]
     pub gossip_port: u16,
@@ -25,12 +22,16 @@ pub struct NetworkCommand {
 impl NetworkCommand {
     /// Run the network subcommand.
     pub async fn run(&self, args: &GlobalArgs) -> Result<()> {
-        if self.disc {
-            println!("Running peer discovery");
-        }
-        if self.gossip {
-            println!("Running gossip driver");
+        if self.only_disc {
+            self.run_discovery(args).await?;
+        } else {
+            self.run_network(args)?;
         }
+        Ok(())
+    }
+
+    /// Runs the full network.
+    pub fn run_network(&self, args: &GlobalArgs) -> Result<()> {
         let signer = ROLLUP_CONFIGS
             .get(&args.l2_chain_id)
             .ok_or(eyre::eyre!("No rollup config found for chain ID"))?
@@ -44,16 +45,10 @@ impl NetworkCommand {
             .with_chain_id(args.l2_chain_id)
             .with_unsafe_block_signer(signer)
             .with_gossip_addr(socket)
-            .build()
-            .expect("Failed to builder network driver");
-
-        // Call `.start()` on the driver.
+            .build()?;
         let recv =
             driver.take_unsafe_block_recv().ok_or(eyre::eyre!("No unsafe block receiver"))?;
-        driver.start().expect("Failed to start network driver");
-
-        tracing::info!("NetworkDriver started successfully.");
-
+        driver.start()?;
         loop {
             match recv.recv() {
                 Ok(block) => {
@@ -65,4 +60,23 @@ impl NetworkCommand {
             }
         }
     }
+
+    /// Runs only the discovery service.
+    pub async fn run_discovery(&self, args: &GlobalArgs) -> Result<()> {
+        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), self.gossip_port);
+        let mut discovery_builder =
+            DiscoveryBuilder::new().with_address(socket).with_chain_id(args.l2_chain_id);
+        let discovery = discovery_builder.build()?;
+        let mut peer_recv = discovery.start()?;
+        loop {
+            match peer_recv.recv().await {
+                Some(peer) => {
+                    tracing::info!("Received peer: {:?}", peer);
+                }
+                None => {
+                    tracing::warn!("Failed to receive peer");
+                }
+            }
+        }
+    }
 }