Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat:disc config param #81

Merged
merged 1 commit into from
Sep 5, 2024
Merged

feat:disc config param #81

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 87 additions & 7 deletions crates/net/src/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
//! Network Builder Module.

use alloy::primitives::Address;
use discv5::ListenConfig;
use discv5::{Config, ListenConfig};
use eyre::Result;
use std::{
net::{IpAddr, SocketAddr},
Expand Down Expand Up @@ -34,6 +34,9 @@ pub struct NetworkDriverBuilder {
pub discovery_addr: Option<ListenConfig>,
/// The [GossipConfig] constructs the config for `gossipsub`.
pub gossip_config: Option<GossipConfig>,

/// The [Config] constructs the config for `discv5`.
pub discovery_config: Option<Config>,
/// The [Keypair] for the node.
pub keypair: Option<Keypair>,
/// The [TcpConfig] for the swarm.
Expand Down Expand Up @@ -139,6 +142,39 @@ impl NetworkDriverBuilder {
self
}

/// Specifies the [Config] for the `discv5` configuration.
///
/// If not set, the [NetworkDriverBuilder] will fall back to use the [ListenConfig]
/// to construct [Config]. These defaults can be extended by using the
/// [discv5::ConfigBuilder::new] method to build a custom [Config].
///
/// ## Example
///
/// ```rust
/// use discv5::{ConfigBuilder, ListenConfig};
/// use kona_primitives::alloy_primitives::{address, Address};
/// use op_net::builder::NetworkDriverBuilder;
/// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
///
/// let id = 10;
/// let signer = Address::random();
/// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
/// let discovery_config =
/// ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
/// .build();
/// let driver = NetworkDriverBuilder::new()
/// .with_unsafe_block_signer(signer)
/// .with_chain_id(id)
/// .with_gossip_addr(socket)
/// .with_discovery_config(discovery_config)
/// .build()
/// .unwrap();
/// ```
pub fn with_discovery_config(&mut self, cfg: Config) -> &mut Self {
self.discovery_config = Some(cfg);
self
}

/// Builds the [NetworkDriver].
///
/// ## Errors
Expand Down Expand Up @@ -231,15 +267,18 @@ impl NetworkDriverBuilder {
let gossip = GossipDriver::new(swarm, multiaddr, handler.clone());

// Build the discovery service
let discovery_builder =
let mut discovery_builder =
DiscoveryBuilder::new().with_address(gossip_addr).with_chain_id(chain_id);

let discovery = if let Some(discovery_addr) = self.discovery_addr.take() {
discovery_builder.with_listen_config(discovery_addr)
} else {
discovery_builder
if let Some(discovery_addr) = self.discovery_addr.take() {
discovery_builder = discovery_builder.with_listen_config(discovery_addr);
}

if let Some(discovery_config) = self.discovery_config.take() {
discovery_builder = discovery_builder.with_discovery_config(discovery_config);
}
.build()?;

let discovery = discovery_builder.build()?;

Ok(NetworkDriver {
discovery,
Expand All @@ -253,6 +292,7 @@ impl NetworkDriverBuilder {
#[cfg(test)]
mod tests {
use super::*;
use discv5::ConfigBuilder;
use libp2p::gossipsub::IdentTopic;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};

Expand Down Expand Up @@ -367,4 +407,44 @@ mod tests {

assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
}

#[test]
fn test_build_network_driver_with_discovery_config() {
let id = 10;
let signer = Address::random();
let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
let discovery_config =
ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
.build();
let driver = NetworkDriverBuilder::new()
.with_unsafe_block_signer(signer)
.with_chain_id(id)
.with_gossip_addr(socket)
.with_discovery_config(discovery_config)
.build()
.unwrap();

assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
}

#[test]
fn test_build_network_driver_with_discovery_config_and_listen_config() {
let id = 10;
let signer = Address::random();
let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
let discovery_config =
ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
.build();
let discovery_addr = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9097);
let driver = NetworkDriverBuilder::new()
.with_unsafe_block_signer(signer)
.with_chain_id(id)
.with_gossip_addr(socket)
.with_discovery_addr(discovery_addr)
.with_discovery_config(discovery_config)
.build()
.unwrap();

assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9097);
}
}
35 changes: 25 additions & 10 deletions crates/net/src/discovery/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
use crate::{discovery::driver::DiscoveryDriver, types::enr::OpStackEnr};
use discv5::{
enr::{CombinedKey, Enr},
ConfigBuilder, Discv5, ListenConfig,
Config, ConfigBuilder, Discv5, ListenConfig,
};
use eyre::{Report, Result};
use std::net::SocketAddr;
Expand All @@ -19,6 +19,9 @@ pub struct DiscoveryBuilder {
chain_id: Option<u64>,
/// The listen config for the discovery service.
listen_config: Option<ListenConfig>,

/// The discovery config for the discovery service.
discovery_config: Option<Config>,
}

impl DiscoveryBuilder {
Expand All @@ -45,6 +48,12 @@ impl DiscoveryBuilder {
self
}

/// Sets the discovery config for the discovery service.
pub fn with_discovery_config(mut self, config: Config) -> Self {
self.discovery_config = Some(config);
self
}

/// Builds a [DiscoveryDriver].
pub fn build(&mut self) -> Result<DiscoveryDriver> {
let chain_id = self.chain_id.ok_or_else(|| eyre::eyre!("chain ID not set"))?;
Expand All @@ -53,17 +62,24 @@ impl DiscoveryBuilder {
use alloy_rlp::Encodable;
opstack.encode(&mut opstack_data);

let config = if let Some(mut discovery_config) = self.discovery_config.take() {
if let Some(listen_config) = self.listen_config.take() {
discovery_config.listen_config = listen_config;
}
Ok::<Config, Report>(discovery_config)
} else {
let listen_config = self
.listen_config
.take()
.or_else(|| self.address.map(ListenConfig::from))
.ok_or_else(|| eyre::eyre!("listen config not set"))?;
Ok(ConfigBuilder::new(listen_config).build())
}?;

let key = CombinedKey::generate_secp256k1();
let mut enr_builder = Enr::builder();
enr_builder.add_value_rlp(OP_CL_KEY, opstack_data.into());
let listen_config = self.listen_config.take().map_or_else(
|| {
let addr = self.address.ok_or(eyre::eyre!("address not set"))?;
Ok::<ListenConfig, Report>(ListenConfig::from(addr))
},
Ok,
)?;
match listen_config {
match config.listen_config {
ListenConfig::Ipv4 { ip, port } => {
enr_builder.ip4(ip).tcp4(port);
}
Expand All @@ -76,7 +92,6 @@ impl DiscoveryBuilder {
}
}
let enr = enr_builder.build(&key)?;
let config = ConfigBuilder::new(listen_config).build();

let disc = Discv5::new(enr, key, config)
.map_err(|_| eyre::eyre!("could not create disc service"))?;
Expand Down