From f5ae970e63f718b5c57b213e1b74aae7fb26b477 Mon Sep 17 00:00:00 2001 From: GeemoCandama <104614073+GeemoCandama@users.noreply.github.com> Date: Fri, 30 Dec 2022 03:56:59 -0600 Subject: [PATCH] feat(cli): add `--disable-discovery` (#597) * redo add disable_discovery cli flag * incorporate option into discovery test * smol touch up * rustmft Co-authored-by: Matthias Seitz --- bin/reth/src/config.rs | 2 ++ bin/reth/src/node/mod.rs | 10 ++++-- crates/net/network/src/config.rs | 25 +++++++++---- crates/net/network/src/discovery.rs | 55 ++++++++++++++++------------- crates/net/network/src/manager.rs | 9 +++-- 5 files changed, 66 insertions(+), 35 deletions(-) diff --git a/bin/reth/src/config.rs b/bin/reth/src/config.rs index ffe5bdf42..67a75214e 100644 --- a/bin/reth/src/config.rs +++ b/bin/reth/src/config.rs @@ -27,6 +27,7 @@ impl Config { db: Arc, chain_id: u64, genesis_hash: H256, + disable_discovery: bool, ) -> NetworkConfig> { let peer_config = reth_network::PeersConfig::default() .with_trusted_nodes(self.peers.trusted_nodes.clone()) @@ -36,6 +37,7 @@ impl Config { .peer_config(peer_config) .genesis_hash(genesis_hash) .chain_id(chain_id) + .set_discovery(disable_discovery) .build() } } diff --git a/bin/reth/src/node/mod.rs b/bin/reth/src/node/mod.rs index be2dbbfd1..db7ebf52c 100644 --- a/bin/reth/src/node/mod.rs +++ b/bin/reth/src/node/mod.rs @@ -75,6 +75,10 @@ pub struct Command { /// NOTE: This is a temporary flag #[arg(long = "debug.tip")] tip: Option, + + /// Disable the discovery service. + #[arg(short, long)] + disable_discovery: bool, } impl Command { @@ -98,8 +102,10 @@ impl Command { let consensus = Arc::new(BeaconConsensus::new(self.chain.consensus.clone())); let genesis_hash = init_genesis(db.clone(), self.chain.genesis.clone())?; - let network = - config.network_config(db.clone(), chain_id, genesis_hash).start_network().await?; + let network = config + .network_config(db.clone(), chain_id, genesis_hash, self.disable_discovery) + .start_network() + .await?; info!(peer_id = ?network.peer_id(), local_addr = %network.local_addr(), "Started p2p networking"); diff --git a/crates/net/network/src/config.rs b/crates/net/network/src/config.rs index a629d563b..4bbb12042 100644 --- a/crates/net/network/src/config.rs +++ b/crates/net/network/src/config.rs @@ -42,7 +42,7 @@ pub struct NetworkConfig { /// All boot nodes to start network discovery with. pub boot_nodes: HashSet, /// How to set up discovery. - pub discovery_v4_config: Discv4Config, + pub discovery_v4_config: Option, /// Address to use for discovery pub discovery_addr: SocketAddr, /// Address to listen for incoming connections @@ -89,7 +89,7 @@ impl NetworkConfig { /// Sets the config to use for the discovery v4 protocol. pub fn set_discovery_v4(mut self, discovery_config: Discv4Config) -> Self { - self.discovery_v4_config = discovery_config; + self.discovery_v4_config = Some(discovery_config); self } @@ -125,7 +125,7 @@ pub struct NetworkConfigBuilder { /// The node's secret key, from which the node's identity is derived. secret_key: SecretKey, /// How to set up discovery. - discovery_v4_builder: Discv4ConfigBuilder, + discovery_v4_builder: Option, /// All boot nodes to start network discovery with. boot_nodes: HashSet, /// Address to use for discovery @@ -164,7 +164,7 @@ impl NetworkConfigBuilder { Self { client, secret_key, - discovery_v4_builder: Default::default(), + discovery_v4_builder: Some(Default::default()), boot_nodes: Default::default(), discovery_addr: None, listener_addr: None, @@ -270,7 +270,7 @@ impl NetworkConfigBuilder { /// Sets the discv4 config to use. pub fn discovery(mut self, builder: Discv4ConfigBuilder) -> Self { - self.discovery_v4_builder = builder; + self.discovery_v4_builder = Some(builder); self } @@ -280,6 +280,19 @@ impl NetworkConfigBuilder { self } + /// Sets the discovery service off on true. + pub fn set_discovery(mut self, disable_discovery: bool) -> Self { + if disable_discovery { + self.disable_discovery(); + } + self + } + + /// disables discovery. + pub fn disable_discovery(&mut self) { + self.discovery_v4_builder = None; + } + /// Consumes the type and creates the actual [`NetworkConfig`] pub fn build(self) -> NetworkConfig { let peer_id = self.get_peer_id(); @@ -322,7 +335,7 @@ impl NetworkConfigBuilder { client, secret_key, boot_nodes, - discovery_v4_config: discovery_v4_builder.build(), + discovery_v4_config: discovery_v4_builder.map(|builder| builder.build()), discovery_addr: discovery_addr.unwrap_or_else(|| { SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT)) }), diff --git a/crates/net/network/src/discovery.rs b/crates/net/network/src/discovery.rs index 820cf63f2..68c799909 100644 --- a/crates/net/network/src/discovery.rs +++ b/crates/net/network/src/discovery.rs @@ -24,15 +24,13 @@ pub struct Discovery { /// Local ENR of the discovery service. local_enr: NodeRecord, /// Handler to interact with the Discovery v4 service - discv4: Discv4, + discv4: Option, /// All KAD table updates from the discv4 service. - discv4_updates: ReceiverStream, - /// The initial config for the discv4 service - _dsicv4_config: Discv4Config, + discv4_updates: Option>, /// Events buffered until polled. queued_events: VecDeque, /// The handle to the spawned discv4 service - _discv4_service: JoinHandle<()>, + _discv4_service: Option>, } impl Discovery { @@ -43,23 +41,26 @@ impl Discovery { pub async fn new( discovery_addr: SocketAddr, sk: SecretKey, - dsicv4_config: Discv4Config, + discv4_config: Option, ) -> Result { let local_enr = NodeRecord::from_secret_key(discovery_addr, &sk); - let (discv4, mut discv4_service) = - Discv4::bind(discovery_addr, local_enr, sk, dsicv4_config.clone()) - .await - .map_err(NetworkError::Discovery)?; - let discv4_updates = discv4_service.update_stream(); - - // spawn the service - let _discv4_service = discv4_service.spawn(); + let (discv4, discv4_updates, _discv4_service) = if let Some(disc_config) = discv4_config { + let (discv4, mut discv4_service) = + Discv4::bind(discovery_addr, local_enr, sk, disc_config) + .await + .map_err(NetworkError::Discovery)?; + let discv4_updates = discv4_service.update_stream(); + // spawn the service + let _discv4_service = discv4_service.spawn(); + (Some(discv4), Some(discv4_updates), Some(_discv4_service)) + } else { + (None, None, None) + }; Ok(Self { local_enr, discv4, discv4_updates, - _dsicv4_config: dsicv4_config, _discv4_service, discovered_nodes: Default::default(), queued_events: Default::default(), @@ -69,17 +70,23 @@ impl Discovery { /// Updates the `eth:ForkId` field in discv4. #[allow(unused)] pub(crate) fn update_fork_id(&self, fork_id: ForkId) { - self.discv4.set_eip868_rlp("eth".as_bytes().to_vec(), fork_id) + if let Some(discv4) = &self.discv4 { + discv4.set_eip868_rlp("eth".as_bytes().to_vec(), fork_id) + } } /// Bans the [`IpAddr`] in the discovery service. pub(crate) fn ban_ip(&self, ip: IpAddr) { - self.discv4.ban_ip(ip) + if let Some(discv4) = &self.discv4 { + discv4.ban_ip(ip) + } } /// Bans the [`PeerId`] and [`IpAddr`] in the discovery service. pub(crate) fn ban(&self, peer_id: PeerId, ip: IpAddr) { - self.discv4.ban(peer_id, ip) + if let Some(discv4) = &self.discv4 { + discv4.ban(peer_id, ip) + } } /// Returns the id with which the local identifies itself in the network @@ -122,7 +129,9 @@ impl Discovery { } // drain the update stream - while let Poll::Ready(Some(update)) = self.discv4_updates.poll_next_unpin(cx) { + while let Some(Poll::Ready(Some(update))) = + self.discv4_updates.as_mut().map(|disc_updates| disc_updates.poll_next_unpin(cx)) + { self.on_discv4_update(update) } @@ -139,7 +148,6 @@ impl Discovery { /// /// NOTE: This instance does nothing pub(crate) fn noop() -> Self { - let (_tx, rx) = tokio::sync::mpsc::channel(1); Self { discovered_nodes: Default::default(), local_enr: NodeRecord { @@ -148,11 +156,10 @@ impl Discovery { udp_port: 0, id: PeerId::random(), }, - discv4: Discv4::noop(), - discv4_updates: ReceiverStream::new(rx), - _dsicv4_config: Default::default(), + discv4: Default::default(), + discv4_updates: Default::default(), queued_events: Default::default(), - _discv4_service: tokio::task::spawn(async move {}), + _discv4_service: Default::default(), } } } diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs index 3cea4b1a0..1074b9cbd 100644 --- a/crates/net/network/src/manager.rs +++ b/crates/net/network/src/manager.rs @@ -160,9 +160,12 @@ where let incoming = ConnectionListener::bind(listener_addr).await?; let listener_address = Arc::new(Mutex::new(incoming.local_address())); - // merge configured boot nodes - discovery_v4_config.bootstrap_nodes.extend(boot_nodes.clone()); - discovery_v4_config.add_eip868_pair("eth", status.forkid); + discovery_v4_config = discovery_v4_config.map(|mut disc_config| { + // merge configured boot nodes + disc_config.bootstrap_nodes.extend(boot_nodes.clone()); + disc_config.add_eip868_pair("eth", status.forkid); + disc_config + }); let discovery = Discovery::new(discovery_addr, secret_key, discovery_v4_config).await?; // need to retrieve the addr here since provided port could be `0`