diff --git a/bin/hera/src/network.rs b/bin/hera/src/network.rs index d31e9d4..5a58452 100644 --- a/bin/hera/src/network.rs +++ b/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"); + } + } + } + } }