Skip to content

Commit

Permalink
chore: integrate discv5 config builder in networkconfig builder (#7856)
Browse files Browse the repository at this point in the history 
Co-authored-by: Emilia Hane <elsaemiliaevahane@gmail.com>
  • Loading branch information
@mattsse @emhane
mattsse and emhane committed May 20, 2024
1 parent db9c559 commit de79f26
Show file tree
Hide file tree
Showing 7 changed files with 155 additions and 157 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions bin/reth/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ reth-payload-builder.workspace = true
reth-payload-validator.workspace = true
reth-basic-payload-builder.workspace = true
reth-discv4.workspace = true
reth-discv5.workspace = true
reth-static-file = { workspace = true }
reth-trie = { workspace = true, features = ["metrics"] }
reth-nippy-jar.workspace = true
Expand Down
86 changes: 39 additions & 47 deletions bin/reth/src/commands/p2p/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use crate::{
args::{
get_secret_key,
utils::{chain_help, chain_spec_value_parser, hash_or_num_value_parser, SUPPORTED_CHAINS},
DatabaseArgs, DiscoveryArgs,
DatabaseArgs, DiscoveryArgs, NetworkArgs,
},
dirs::{DataDirPath, MaybePlatformPath},
utils::get_single_header,
Expand All @@ -14,12 +14,11 @@ use clap::{Parser, Subcommand};
use discv5::ListenConfig;
use reth_config::Config;
use reth_db::create_db;
use reth_discv4::NatResolver;
use reth_interfaces::p2p::bodies::client::BodiesClient;
use reth_primitives::{BlockHashOrNumber, ChainSpec, NodeRecord};
use reth_primitives::{BlockHashOrNumber, ChainSpec};
use reth_provider::ProviderFactory;
use std::{
net::{SocketAddrV4, SocketAddrV6},
net::{IpAddr, SocketAddrV4, SocketAddrV6},
path::PathBuf,
sync::Arc,
};
Expand Down Expand Up @@ -53,31 +52,14 @@ pub struct Command {
#[arg(long, value_name = "DATA_DIR", verbatim_doc_comment, default_value_t)]
datadir: MaybePlatformPath<DataDirPath>,

/// Secret key to use for this node.
///
/// This also will deterministically set the peer ID.
#[arg(long, value_name = "PATH")]
p2p_secret_key: Option<PathBuf>,

/// Disable the discovery service.
#[command(flatten)]
pub discovery: DiscoveryArgs,

/// Target trusted peer
#[arg(long)]
trusted_peer: Option<NodeRecord>,

/// Connect only to trusted peers
#[arg(long)]
trusted_only: bool,
pub network: NetworkArgs,

/// The number of retries per request
#[arg(long, default_value = "5")]
retries: usize,

#[arg(long, default_value = "any")]
nat: NatResolver,

#[command(flatten)]
db: DatabaseArgs,

Expand Down Expand Up @@ -113,65 +95,75 @@ impl Command {

let mut config: Config = confy::load_path(&config_path).unwrap_or_default();

if let Some(peer) = self.trusted_peer {
for &peer in &self.network.trusted_peers {
config.peers.trusted_nodes.insert(peer);
}

if config.peers.trusted_nodes.is_empty() && self.trusted_only {
if config.peers.trusted_nodes.is_empty() && self.network.trusted_only {
eyre::bail!("No trusted nodes. Set trusted peer with `--trusted-peer <enode record>` or set `--trusted-only` to `false`")
}

config.peers.trusted_nodes_only = self.trusted_only;
config.peers.trusted_nodes_only = self.network.trusted_only;

let default_secret_key_path = data_dir.p2p_secret();
let secret_key_path = self.p2p_secret_key.clone().unwrap_or(default_secret_key_path);
let secret_key_path =
self.network.p2p_secret_key.clone().unwrap_or(default_secret_key_path);
let p2p_secret_key = get_secret_key(&secret_key_path)?;
let rlpx_socket = (self.network.addr, self.network.port).into();
let boot_nodes = self.chain.bootnodes().unwrap_or_default();

let mut network_config_builder = config
.network_config(self.nat, None, p2p_secret_key)
.network_config(self.network.nat, None, p2p_secret_key)
.chain_spec(self.chain.clone())
.disable_discv4_discovery_if(self.chain.chain.is_optimism())
.boot_nodes(self.chain.bootnodes().unwrap_or_default());
.boot_nodes(boot_nodes.clone());

network_config_builder = self.discovery.apply_to_builder(network_config_builder);

let mut network_config = network_config_builder.build(Arc::new(ProviderFactory::new(
noop_db,
self.chain.clone(),
data_dir.static_files(),
)?));

if !self.discovery.disable_discovery &&
(self.discovery.enable_discv5_discovery ||
network_config.chain_spec.chain.is_optimism())
{
network_config = network_config.discovery_v5_with_config_builder(|builder| {
network_config_builder = self
.network
.discovery
.apply_to_builder(network_config_builder, rlpx_socket)
.map_discv5_config_builder(|builder| {
let DiscoveryArgs {
discv5_addr: discv5_addr_ipv4,
discv5_addr,
discv5_addr_ipv6,
discv5_port: discv5_port_ipv4,
discv5_port,
discv5_port_ipv6,
discv5_lookup_interval,
discv5_bootstrap_lookup_interval,
discv5_bootstrap_lookup_countdown,
..
} = self.discovery;
} = self.network.discovery;

// Use rlpx address if none given
let discv5_addr_ipv4 = discv5_addr.or(match self.network.addr {
IpAddr::V4(ip) => Some(ip),
IpAddr::V6(_) => None,
});
let discv5_addr_ipv6 = discv5_addr_ipv6.or(match self.network.addr {
IpAddr::V4(_) => None,
IpAddr::V6(ip) => Some(ip),
});

builder
.discv5_config(
discv5::ConfigBuilder::new(ListenConfig::from_two_sockets(
discv5_addr_ipv4.map(|addr| SocketAddrV4::new(addr, discv5_port_ipv4)),
discv5_addr_ipv4.map(|addr| SocketAddrV4::new(addr, discv5_port)),
discv5_addr_ipv6
.map(|addr| SocketAddrV6::new(addr, discv5_port_ipv6, 0, 0)),
))
.build(),
)
.add_unsigned_boot_nodes(boot_nodes.into_iter())
.lookup_interval(discv5_lookup_interval)
.bootstrap_lookup_interval(discv5_bootstrap_lookup_interval)
.bootstrap_lookup_countdown(discv5_bootstrap_lookup_countdown)
.build()
});
}

let network_config = network_config_builder.build(Arc::new(ProviderFactory::new(
noop_db,
self.chain.clone(),
data_dir.static_files(),
)?));

let network = network_config.start_network().await?;
let fetch_client = network.fetch_client().await?;
Expand Down
6 changes: 2 additions & 4 deletions crates/net/discv5/src/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ impl ConfigBuilder {
}

/// Config used to bootstrap [`discv5::Discv5`].
#[derive(Debug)]
#[derive(Clone, Debug)]
pub struct Config {
/// Config used by [`discv5::Discv5`]. Contains the [`ListenConfig`], with the discovery listen
/// socket.
Expand Down Expand Up @@ -296,9 +296,7 @@ impl Config {
discovered_peer_filter: None,
}
}
}

impl Config {
/// Returns the discovery (UDP) socket contained in the [`discv5::Config`]. Returns the IPv6
/// socket, if both IPv4 and v6 are configured. This socket will be advertised to peers in the
/// local [`Enr`](discv5::enr::Enr).
Expand Down Expand Up @@ -416,7 +414,7 @@ pub fn discv5_sockets_wrt_rlpx_addr(

/// A boot node can be added either as a string in either 'enode' URL scheme or serialized from
/// [`Enr`](discv5::Enr) type.
#[derive(Debug, PartialEq, Eq, Hash, Display)]
#[derive(Clone, Debug, PartialEq, Eq, Hash, Display)]
pub enum BootNode {
/// An unsigned node record.
#[display(fmt = "{_0}")]
Expand Down
81 changes: 42 additions & 39 deletions crates/net/network/src/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,43 +104,6 @@ impl<C> NetworkConfig<C> {
self
}

/// Sets the config to use for the discovery v5 protocol, with help of the
/// [`reth_discv5::ConfigBuilder`].
/// ```
/// use reth_network::NetworkConfigBuilder;
/// use secp256k1::{rand::thread_rng, SecretKey};
///
/// let sk = SecretKey::new(&mut thread_rng());
/// let network_config = NetworkConfigBuilder::new(sk).build(());
/// let fork_id = network_config.status.forkid;
/// let network_config = network_config
/// .discovery_v5_with_config_builder(|builder| builder.fork(b"eth", fork_id).build());
/// ```

pub fn discovery_v5_with_config_builder(
self,
f: impl FnOnce(reth_discv5::ConfigBuilder) -> reth_discv5::Config,
) -> Self {
let network_stack_id = NetworkStackId::id(&self.chain_spec);
let fork_id = self.chain_spec.latest_fork_id();
let boot_nodes = self.boot_nodes.clone();

let mut builder = reth_discv5::Config::builder(self.listener_addr)
.add_unsigned_boot_nodes(boot_nodes.into_iter());

if let Some(id) = network_stack_id {
builder = builder.fork(id, fork_id)
}

self.set_discovery_v5(f(builder))
}

/// Sets the config to use for the discovery v5 protocol.
pub fn set_discovery_v5(mut self, discv5_config: reth_discv5::Config) -> Self {
self.discovery_v5_config = Some(discv5_config);
self
}

/// Sets the address for the incoming RLPx connection listener.
pub fn set_listener_addr(mut self, listener_addr: SocketAddr) -> Self {
self.listener_addr = listener_addr;
Expand Down Expand Up @@ -180,6 +143,9 @@ pub struct NetworkConfigBuilder {
dns_discovery_config: Option<DnsDiscoveryConfig>,
/// How to set up discovery version 4.
discovery_v4_builder: Option<Discv4ConfigBuilder>,
/// How to set up discovery version 5.
#[serde(skip)]
discovery_v5_builder: Option<reth_discv5::ConfigBuilder>,
/// All boot nodes to start network discovery with.
boot_nodes: HashSet<NodeRecord>,
/// Address to use for discovery
Expand Down Expand Up @@ -222,6 +188,7 @@ impl NetworkConfigBuilder {
secret_key,
dns_discovery_config: Some(Default::default()),
discovery_v4_builder: Some(Default::default()),
discovery_v5_builder: None,
boot_nodes: Default::default(),
discovery_addr: None,
listener_addr: None,
Expand Down Expand Up @@ -348,12 +315,17 @@ impl NetworkConfigBuilder {
}

/// Sets the discv4 config to use.
//
pub fn discovery(mut self, builder: Discv4ConfigBuilder) -> Self {
self.discovery_v4_builder = Some(builder);
self
}

/// Sets the discv5 config to use.
pub fn discovery_v5(mut self, builder: reth_discv5::ConfigBuilder) -> Self {
self.discovery_v5_builder = Some(builder);
self
}

/// Sets the dns discovery config to use.
pub fn dns_discovery(mut self, config: DnsDiscoveryConfig) -> Self {
self.dns_discovery_config = Some(config);
Expand Down Expand Up @@ -420,6 +392,36 @@ impl NetworkConfigBuilder {
}
}

/// Calls a closure on [`reth_discv5::ConfigBuilder`], if discv5 discovery is enabled and the
/// builder has been set.
/// ```
/// use reth_network::NetworkConfigBuilder;
/// use reth_primitives::MAINNET;
/// use reth_provider::test_utils::NoopProvider;
/// use secp256k1::{rand::thread_rng, SecretKey};
///
/// let sk = SecretKey::new(&mut thread_rng());
/// let fork_id = MAINNET.latest_fork_id();
/// let network_config = NetworkConfigBuilder::new(sk)
/// .map_discv5_config_builder(|builder| builder.fork(b"eth", fork_id))
/// .build(NoopProvider::default());
/// ```
pub fn map_discv5_config_builder(
mut self,
f: impl FnOnce(reth_discv5::ConfigBuilder) -> reth_discv5::ConfigBuilder,
) -> Self {
if let Some(mut builder) = self.discovery_v5_builder {
if let Some(network_stack_id) = NetworkStackId::id(&self.chain_spec) {
let fork_id = self.chain_spec.latest_fork_id();
builder = builder.fork(network_stack_id, fork_id);
}

self.discovery_v5_builder = Some(f(builder));
}

self
}

/// Adds a new additional protocol to the RLPx sub-protocol list.
pub fn add_rlpx_sub_protocol(mut self, protocol: impl IntoRlpxSubProtocol) -> Self {
self.extra_protocols.push(protocol);
Expand Down Expand Up @@ -458,6 +460,7 @@ impl NetworkConfigBuilder {
secret_key,
mut dns_discovery_config,
discovery_v4_builder,
discovery_v5_builder,
boot_nodes,
discovery_addr,
listener_addr,
Expand Down Expand Up @@ -511,7 +514,7 @@ impl NetworkConfigBuilder {
boot_nodes,
dns_discovery_config,
discovery_v4_config: discovery_v4_builder.map(|builder| builder.build()),
discovery_v5_config: None,
discovery_v5_config: discovery_v5_builder.map(|builder| builder.build()),
discovery_v4_addr: discovery_addr.unwrap_or(DEFAULT_DISCOVERY_ADDRESS),
listener_addr,
peers_config: peers_config.unwrap_or_default(),
Expand Down

0 comments on commit de79f26

Please sign in to comment.