kona_net/
builder.rs

1//! Network Builder Module.
2
3use alloy_primitives::Address;
4use discv5::{Config, ListenConfig};
5use std::{
6    net::{IpAddr, SocketAddr},
7    time::Duration,
8};
9use tokio::sync::watch::channel;
10
11use libp2p::{
12    gossipsub::Config as GossipConfig, multiaddr::Protocol, noise::Config as NoiseConfig,
13    tcp::Config as TcpConfig, yamux::Config as YamuxConfig, Multiaddr, SwarmBuilder,
14};
15use libp2p_identity::Keypair;
16
17use crate::{
18    discovery::builder::{DiscoveryBuilder, DiscoveryBuilderError},
19    driver::NetworkDriver,
20    gossip::{
21        behaviour::{Behaviour, BehaviourError},
22        config,
23        driver::GossipDriver,
24        handler::BlockHandler,
25    },
26};
27
28/// An error from the [NetworkDriverBuilder].
29#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
30pub enum NetworkDriverBuilderError {
31    /// The unsafe block signer is not set.
32    #[error("unsafe block signer not set")]
33    UnsafeBlockSignerNotSet,
34    /// The chain ID is not set.
35    #[error("chain ID not set")]
36    ChainIdNotSet,
37    /// The gossip address is not set.
38    #[error("gossip_addr address not set")]
39    GossipAddrNotSet,
40    /// A behaviour error.
41    #[error("behaviour error: {0}")]
42    BehaviourError(#[from] BehaviourError),
43    /// A config builder error.
44    #[error("config builder error")]
45    ConfigBuilderError,
46    /// A TCP error.
47    #[error("TCP error")]
48    TcpError,
49    /// An error when setting the behaviour on the swarm builder.
50    #[error("error setting behaviour on swarm builder")]
51    WithBehaviourError,
52    /// A discovery builder error.
53    #[error("discovery builder error: {0}")]
54    DiscoveryBuilderError(#[from] DiscoveryBuilderError),
55}
56
57/// Constructs a [NetworkDriver] for Optimism's consensus-layer.
58#[derive(Default)]
59pub struct NetworkDriverBuilder {
60    /// The chain ID of the network.
61    pub chain_id: Option<u64>,
62    /// The unsafe block signer.
63    pub unsafe_block_signer: Option<Address>,
64    /// The socket address that the gossip service is listening on.
65    pub gossip_addr: Option<SocketAddr>,
66    /// The listen config that the discovery service is listening on.
67    pub discovery_addr: Option<ListenConfig>,
68    /// The [GossipConfig] constructs the config for `gossipsub`.
69    pub gossip_config: Option<GossipConfig>,
70    /// The interval to discovery random nodes.
71    pub interval: Option<Duration>,
72    /// The [Config] constructs the config for `discv5`.
73    pub discovery_config: Option<Config>,
74    /// The [Keypair] for the node.
75    pub keypair: Option<Keypair>,
76    /// The [TcpConfig] for the swarm.
77    pub tcp_config: Option<TcpConfig>,
78    /// The [NoiseConfig] for the swarm.
79    pub noise_config: Option<NoiseConfig>,
80    /// The [YamuxConfig] for the swarm.
81    pub yamux_config: Option<YamuxConfig>,
82    /// The idle connection timeout.
83    pub timeout: Option<Duration>,
84}
85
86impl NetworkDriverBuilder {
87    /// Creates a new [NetworkDriverBuilder].
88    pub fn new() -> Self {
89        Self::default()
90    }
91
92    /// Specifies the chain ID of the network.
93    pub fn with_chain_id(&mut self, chain_id: u64) -> &mut Self {
94        self.chain_id = Some(chain_id);
95        self
96    }
97
98    /// Specifies the unsafe block signer.
99    pub fn with_unsafe_block_signer(&mut self, unsafe_block_signer: Address) -> &mut Self {
100        self.unsafe_block_signer = Some(unsafe_block_signer);
101        self
102    }
103
104    /// Specifies the interval to discovery random nodes.
105    pub fn with_interval(&mut self, interval: Duration) -> &mut Self {
106        self.interval = Some(interval);
107        self
108    }
109
110    /// Specifies the socket address that the gossip service is listening on.
111    pub fn with_gossip_addr(&mut self, socket: SocketAddr) -> &mut Self {
112        self.gossip_addr = Some(socket);
113        self
114    }
115
116    /// Specifies the listen config that the discovery service is listening on.
117    pub fn with_discovery_addr(&mut self, listen_config: ListenConfig) -> &mut Self {
118        self.discovery_addr = Some(listen_config);
119        self
120    }
121
122    /// Specifies the keypair for the node.
123    pub fn with_keypair(&mut self, keypair: Keypair) -> &mut Self {
124        self.keypair = Some(keypair);
125        self
126    }
127
128    /// Specifies the [TcpConfig] for the swarm.
129    pub fn with_tcp_config(&mut self, tcp_config: TcpConfig) -> &mut Self {
130        self.tcp_config = Some(tcp_config);
131        self
132    }
133
134    /// Specifies the [NoiseConfig] for the swarm.
135    pub fn with_noise_config(&mut self, noise_config: NoiseConfig) -> &mut Self {
136        self.noise_config = Some(noise_config);
137        self
138    }
139
140    /// Specifies the [YamuxConfig] for the swarm.
141    pub fn with_yamux_config(&mut self, yamux_config: YamuxConfig) -> &mut Self {
142        self.yamux_config = Some(yamux_config);
143        self
144    }
145
146    /// Set the swarm's idle connection timeout.
147    pub fn with_idle_connection_timeout(&mut self, timeout: Duration) -> &mut Self {
148        self.timeout = Some(timeout);
149        self
150    }
151
152    /// Specifies the [GossipConfig] for the `gossipsub` configuration.
153    ///
154    /// If not set, the [NetworkDriverBuilder] will use the default gossipsub
155    /// configuration defined in [config::default_config]. These defaults can
156    /// be extended by using the [config::default_config_builder] method to
157    /// build a custom [GossipConfig].
158    ///
159    /// ## Example
160    ///
161    /// ```rust,ignore
162    /// use kona_net::gossip::config;
163    /// use kona_net::NetworkDriverBuilder;
164    /// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
165    ///
166    /// let chain_id = 10;
167    /// let signer = Address::random();
168    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
169    ///
170    /// // Let's say we want to enable flood publishing and use all other default settings.
171    /// let cfg = config::default_config_builder().flood_publish(true).build().unwrap();
172    /// let mut builder = NetworkDriverBuilder::new()
173    ///    .with_unsafe_block_signer(signer)
174    ///    .with_chain_id(chain_id)
175    ///    .with_gossip_addr(socket)
176    ///    .with_gossip_config(cfg);
177    ///    .build()
178    ///    .unwrap();
179    /// ```
180    pub fn with_gossip_config(&mut self, cfg: GossipConfig) -> &mut Self {
181        self.gossip_config = Some(cfg);
182        self
183    }
184
185    /// Specifies the [Config] for the `discv5` configuration.
186    ///
187    /// If not set, the [NetworkDriverBuilder] will fall back to use the [ListenConfig]
188    /// to construct [Config]. These defaults can be extended by using the
189    /// [discv5::ConfigBuilder::new] method to build a custom [Config].
190    ///
191    /// ## Example
192    ///
193    /// ```rust
194    /// use alloy_primitives::{address, Address};
195    /// use discv5::{ConfigBuilder, ListenConfig};
196    /// use kona_net::builder::NetworkDriverBuilder;
197    /// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
198    ///
199    /// let id = 10;
200    /// let signer = Address::random();
201    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
202    /// let discovery_config =
203    ///     ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
204    ///         .build();
205    /// let driver = NetworkDriverBuilder::new()
206    ///     .with_unsafe_block_signer(signer)
207    ///     .with_chain_id(id)
208    ///     .with_gossip_addr(socket)
209    ///     .with_discovery_config(discovery_config)
210    ///     .build()
211    ///     .unwrap();
212    /// ```
213    pub fn with_discovery_config(&mut self, cfg: Config) -> &mut Self {
214        self.discovery_config = Some(cfg);
215        self
216    }
217
218    /// Builds the [NetworkDriver].
219    ///
220    /// ## Errors
221    ///
222    /// Returns an error if any of the following required fields are not set:
223    /// - [NetworkDriverBuilder::unsafe_block_signer]
224    /// - [NetworkDriverBuilder::chain_id]
225    /// - [NetworkDriverBuilder::gossip_addr]
226    ///
227    /// If explicitly set, the following fields are used for discovery address, otherwise the gossip
228    /// address is used:
229    /// - [NetworkDriverBuilder::discovery_addr]
230    ///
231    /// Set these fields using the respective methods on the [NetworkDriverBuilder]
232    /// before calling this method.
233    ///
234    /// ## Example
235    ///
236    /// ```rust,ignore
237    /// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
238    /// use kona_net::NetworkDriverBuilder;
239    ///
240    /// let chain_id = 10;
241    /// let signer = Address::random();
242    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
243    /// let driver = NetworkDriverBuilder::new()
244    ///    .with_unsafe_block_signer(signer)
245    ///    .with_chain_id(chain_id)
246    ///    .with_gossip_addr(socket)
247    ///    .build()
248    ///    .unwrap();
249    ///
250    /// Or if you want to use a different discovery address:
251    ///
252    /// let chain_id = 10;
253    /// let signer = Address::random();
254    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
255    /// let listen_config = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9999);
256    /// let driver = NetworkDriverBuilder::new()
257    ///    .with_unsafe_block_signer(signer)
258    ///    .with_chain_id(chain_id)
259    ///    .with_gossip_addr(socket)
260    ///    .with_discovery_addr(listen_config)
261    ///    .build()
262    ///    .unwrap();
263    /// ```
264    pub fn build(&mut self) -> Result<NetworkDriver, NetworkDriverBuilderError> {
265        // Build the config for gossipsub.
266        let config = match self.gossip_config.take() {
267            Some(cfg) => cfg,
268            None => config::default_config()
269                .map_err(|_| NetworkDriverBuilderError::ConfigBuilderError)?,
270        };
271        let unsafe_block_signer =
272            self.unsafe_block_signer.ok_or(NetworkDriverBuilderError::UnsafeBlockSignerNotSet)?;
273        let chain_id = self.chain_id.ok_or(NetworkDriverBuilderError::ChainIdNotSet)?;
274
275        // Create the block handler.
276        let (unsafe_block_signer_sender, unsafe_block_signer_recv) = channel(unsafe_block_signer);
277        let (handler, unsafe_block_recv) = BlockHandler::new(chain_id, unsafe_block_signer_recv);
278
279        // Construct the gossipsub behaviour.
280        let behaviour = Behaviour::new(config, &[Box::new(handler.clone())])?;
281
282        // Build the swarm.
283        let timeout = self.timeout.take().unwrap_or(Duration::from_secs(60));
284        let noise_config = self.noise_config.take();
285        let keypair = self.keypair.take().unwrap_or(Keypair::generate_secp256k1());
286        let swarm = SwarmBuilder::with_existing_identity(keypair)
287            .with_tokio()
288            .with_tcp(
289                self.tcp_config.take().unwrap_or_default(),
290                |i: &Keypair| match noise_config {
291                    Some(cfg) => Ok(cfg),
292                    None => NoiseConfig::new(i),
293                },
294                || self.yamux_config.take().unwrap_or_default(),
295            )
296            .map_err(|_| NetworkDriverBuilderError::TcpError)?
297            .with_behaviour(|_| behaviour)
298            .map_err(|_| NetworkDriverBuilderError::WithBehaviourError)?
299            .with_swarm_config(|c| c.with_idle_connection_timeout(timeout))
300            .build();
301
302        let gossip_addr =
303            self.gossip_addr.take().ok_or(NetworkDriverBuilderError::GossipAddrNotSet)?;
304        let mut multiaddr = Multiaddr::empty();
305        match gossip_addr.ip() {
306            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
307            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
308        }
309        multiaddr.push(Protocol::Tcp(gossip_addr.port()));
310        let gossip = GossipDriver::new(swarm, multiaddr, handler.clone());
311
312        // Build the discovery service
313        let mut discovery_builder =
314            DiscoveryBuilder::new().with_address(gossip_addr).with_chain_id(chain_id);
315
316        if let Some(discovery_addr) = self.discovery_addr.take() {
317            discovery_builder = discovery_builder.with_listen_config(discovery_addr);
318        }
319
320        if let Some(discovery_config) = self.discovery_config.take() {
321            discovery_builder = discovery_builder.with_discovery_config(discovery_config);
322        }
323
324        let mut discovery = discovery_builder.build()?;
325        discovery.interval = self.interval.unwrap_or(Duration::from_secs(10));
326
327        Ok(NetworkDriver {
328            discovery,
329            gossip,
330            unsafe_block_recv: Some(unsafe_block_recv),
331            unsafe_block_signer_sender: Some(unsafe_block_signer_sender),
332        })
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use discv5::ConfigBuilder;
340    use libp2p::gossipsub::IdentTopic;
341    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
342
343    #[test]
344    fn test_build_missing_unsafe_block_signer() {
345        let mut builder = NetworkDriverBuilder::new();
346        let Err(err) = builder.build() else {
347            panic!("expected error when building NetworkDriver without unsafe block signer");
348        };
349        assert_eq!(err, NetworkDriverBuilderError::UnsafeBlockSignerNotSet);
350    }
351
352    #[test]
353    fn test_build_missing_chain_id() {
354        let mut builder = NetworkDriverBuilder::new();
355        let Err(err) = builder.with_unsafe_block_signer(Address::random()).build() else {
356            panic!("expected error when building NetworkDriver without chain id");
357        };
358        assert_eq!(err, NetworkDriverBuilderError::ChainIdNotSet);
359    }
360
361    #[test]
362    fn test_build_missing_socket() {
363        let mut builder = NetworkDriverBuilder::new();
364        let Err(err) = builder.with_unsafe_block_signer(Address::random()).with_chain_id(1).build()
365        else {
366            panic!("expected error when building NetworkDriver without socket");
367        };
368        assert_eq!(err, NetworkDriverBuilderError::GossipAddrNotSet);
369    }
370
371    #[test]
372    fn test_build_custom_gossip_config() {
373        let id = 10;
374        let signer = Address::random();
375        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
376        let cfg = config::default_config_builder().flood_publish(true).build().unwrap();
377        let driver = NetworkDriverBuilder::new()
378            .with_unsafe_block_signer(signer)
379            .with_chain_id(id)
380            .with_gossip_addr(socket)
381            .with_gossip_config(cfg)
382            .build()
383            .unwrap();
384        let mut multiaddr = Multiaddr::empty();
385        match socket.ip() {
386            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
387            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
388        }
389        multiaddr.push(Protocol::Tcp(socket.port()));
390
391        // Driver Assertions
392        assert_eq!(driver.gossip.addr, multiaddr);
393        assert_eq!(driver.discovery.chain_id, id);
394
395        // Block Handler Assertions
396        assert_eq!(driver.gossip.handler.chain_id, id);
397        let v1 = IdentTopic::new(format!("/optimism/{}/0/blocks", id));
398        assert_eq!(driver.gossip.handler.blocks_v1_topic.hash(), v1.hash());
399        let v2 = IdentTopic::new(format!("/optimism/{}/1/blocks", id));
400        assert_eq!(driver.gossip.handler.blocks_v2_topic.hash(), v2.hash());
401        let v3 = IdentTopic::new(format!("/optimism/{}/2/blocks", id));
402        assert_eq!(driver.gossip.handler.blocks_v3_topic.hash(), v3.hash());
403    }
404
405    #[test]
406    fn test_build_default_network_driver() {
407        let id = 10;
408        let signer = Address::random();
409        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
410        let driver = NetworkDriverBuilder::new()
411            .with_unsafe_block_signer(signer)
412            .with_chain_id(id)
413            .with_gossip_addr(socket)
414            .build()
415            .unwrap();
416        let mut multiaddr = Multiaddr::empty();
417        match socket.ip() {
418            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
419            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
420        }
421        multiaddr.push(Protocol::Tcp(socket.port()));
422
423        // Driver Assertions
424        assert_eq!(driver.gossip.addr, multiaddr);
425        assert_eq!(driver.discovery.chain_id, id);
426        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9099);
427
428        // Block Handler Assertions
429        assert_eq!(driver.gossip.handler.chain_id, id);
430        let v1 = IdentTopic::new(format!("/optimism/{}/0/blocks", id));
431        assert_eq!(driver.gossip.handler.blocks_v1_topic.hash(), v1.hash());
432        let v2 = IdentTopic::new(format!("/optimism/{}/1/blocks", id));
433        assert_eq!(driver.gossip.handler.blocks_v2_topic.hash(), v2.hash());
434        let v3 = IdentTopic::new(format!("/optimism/{}/2/blocks", id));
435        assert_eq!(driver.gossip.handler.blocks_v3_topic.hash(), v3.hash());
436    }
437
438    #[test]
439    fn test_build_network_driver_with_discovery_addr() {
440        let id = 10;
441        let signer = Address::random();
442        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
443        let discovery_addr = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098);
444        let driver = NetworkDriverBuilder::new()
445            .with_unsafe_block_signer(signer)
446            .with_chain_id(id)
447            .with_gossip_addr(socket)
448            .with_discovery_addr(discovery_addr)
449            .build()
450            .unwrap();
451
452        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
453    }
454
455    #[test]
456    fn test_build_network_driver_with_discovery_config() {
457        let id = 10;
458        let signer = Address::random();
459        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
460        let discovery_config =
461            ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
462                .build();
463        let driver = NetworkDriverBuilder::new()
464            .with_unsafe_block_signer(signer)
465            .with_chain_id(id)
466            .with_gossip_addr(socket)
467            .with_discovery_config(discovery_config)
468            .build()
469            .unwrap();
470
471        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
472    }
473
474    #[test]
475    fn test_build_network_driver_with_discovery_config_and_listen_config() {
476        let id = 10;
477        let signer = Address::random();
478        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
479        let discovery_config =
480            ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
481                .build();
482        let discovery_addr = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9097);
483        let driver = NetworkDriverBuilder::new()
484            .with_unsafe_block_signer(signer)
485            .with_chain_id(id)
486            .with_gossip_addr(socket)
487            .with_discovery_addr(discovery_addr)
488            .with_discovery_config(discovery_config)
489            .build()
490            .unwrap();
491
492        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9097);
493    }
494}