From 60867680fd097807eb718cbd041ef12b323ce00b Mon Sep 17 00:00:00 2001
From: Federico Gimenez <fgimenez@users.noreply.github.com>
Date: Wed, 2 Apr 2025 20:37:55 +0200
Subject: [PATCH] feat(examples): add custom NetworkPrimitives and
 Networkbuilder to custom-node example (#15468)

---
 Cargo.lock                          |  2 +
 examples/custom-node/Cargo.toml     |  3 +
 examples/custom-node/src/lib.rs     |  1 +
 examples/custom-node/src/network.rs | 96 +++++++++++++++++++++++++++++
 4 files changed, 102 insertions(+)
 create mode 100644 examples/custom-node/src/network.rs

diff --git a/Cargo.lock b/Cargo.lock
index 652e4b585c..3373ee441f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3477,6 +3477,7 @@ dependencies = [
  "reth-chain-state",
  "reth-chainspec",
  "reth-codecs",
+ "reth-network",
  "reth-network-peers",
  "reth-node-api",
  "reth-node-builder",
@@ -3488,6 +3489,7 @@ dependencies = [
  "reth-primitives-traits",
  "reth-rpc-api",
  "reth-rpc-engine-api",
+ "reth-transaction-pool",
  "revm-primitives",
  "serde",
 ]
diff --git a/examples/custom-node/Cargo.toml b/examples/custom-node/Cargo.toml
index 34724a0316..05f1b2a8d5 100644
--- a/examples/custom-node/Cargo.toml
+++ b/examples/custom-node/Cargo.toml
@@ -10,6 +10,7 @@ license.workspace = true
 reth-chain-state.workspace = true
 reth-chainspec.workspace = true
 reth-codecs.workspace = true
+reth-network.workspace = true
 reth-network-peers.workspace = true
 reth-node-api.workspace = true
 reth-node-builder.workspace = true
@@ -21,6 +22,7 @@ reth-payload-builder.workspace = true
 reth-primitives-traits.workspace = true
 reth-rpc-api.workspace = true
 reth-rpc-engine-api.workspace = true
+reth-transaction-pool.workspace = true
 
 # revm
 revm-primitives.workspace = true
@@ -55,4 +57,5 @@ arbitrary = [
     "reth-optimism-primitives/arbitrary",
     "reth-primitives-traits/arbitrary",
     "revm-primitives/arbitrary",
+    "reth-transaction-pool/arbitrary",
 ]
diff --git a/examples/custom-node/src/lib.rs b/examples/custom-node/src/lib.rs
index d84c708e3e..be82c639b1 100644
--- a/examples/custom-node/src/lib.rs
+++ b/examples/custom-node/src/lib.rs
@@ -20,6 +20,7 @@ use reth_optimism_node::{
 pub mod chainspec;
 pub mod engine;
 pub mod engine_api;
+pub mod network;
 pub mod primitives;
 
 #[derive(Debug, Clone)]
diff --git a/examples/custom-node/src/network.rs b/examples/custom-node/src/network.rs
new file mode 100644
index 0000000000..127ccbe533
--- /dev/null
+++ b/examples/custom-node/src/network.rs
@@ -0,0 +1,96 @@
+use crate::{
+    chainspec::CustomChainSpec,
+    primitives::{CustomHeader, CustomNodePrimitives},
+};
+use alloy_consensus::{Block, BlockBody};
+use eyre::Result;
+use op_alloy_consensus::OpPooledTransaction;
+use reth_chainspec::{EthChainSpec, Hardforks};
+use reth_network::{NetworkConfig, NetworkHandle, NetworkManager, NetworkPrimitives};
+use reth_node_api::{FullNodeTypes, NodeTypes, TxTy};
+use reth_node_builder::{components::NetworkBuilder, BuilderContext};
+use reth_optimism_primitives::{OpReceipt, OpTransactionSigned};
+use reth_transaction_pool::{PoolTransaction, TransactionPool};
+
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
+#[non_exhaustive]
+pub struct CustomNetworkPrimitives;
+
+impl NetworkPrimitives for CustomNetworkPrimitives {
+    type BlockHeader = CustomHeader;
+    type BlockBody = BlockBody<OpTransactionSigned, CustomHeader>;
+    type Block = Block<OpTransactionSigned, CustomHeader>;
+    type BroadcastedTransaction = OpTransactionSigned;
+    type PooledTransaction = OpPooledTransaction;
+    type Receipt = OpReceipt;
+}
+
+#[derive(Default)]
+pub struct CustomNetworkBuilder {}
+
+impl CustomNetworkBuilder {
+    fn network_config<Node>(
+        &self,
+        ctx: &BuilderContext<Node>,
+    ) -> eyre::Result<NetworkConfig<<Node as FullNodeTypes>::Provider, CustomNetworkPrimitives>>
+    where
+        Node: FullNodeTypes<Types: NodeTypes<ChainSpec: Hardforks>>,
+    {
+        let args = &ctx.config().network;
+        let network_builder = ctx
+            .network_config_builder()?
+            // apply discovery settings
+            .apply(|mut builder| {
+                let rlpx_socket = (args.addr, args.port).into();
+                if args.discovery.disable_discovery {
+                    builder = builder.disable_discv4_discovery();
+                }
+                if !args.discovery.disable_discovery {
+                    builder = builder.discovery_v5(
+                        args.discovery.discovery_v5_builder(
+                            rlpx_socket,
+                            ctx.config()
+                                .network
+                                .resolved_bootnodes()
+                                .or_else(|| ctx.chain_spec().bootnodes())
+                                .unwrap_or_default(),
+                        ),
+                    );
+                }
+
+                builder
+            });
+
+        let network_config = ctx.build_network_config(network_builder);
+
+        Ok(network_config)
+    }
+}
+
+impl<Node, Pool> NetworkBuilder<Node, Pool> for CustomNetworkBuilder
+where
+    Node: FullNodeTypes<
+        Types: NodeTypes<ChainSpec = CustomChainSpec, Primitives = CustomNodePrimitives>,
+    >,
+    Pool: TransactionPool<
+            Transaction: PoolTransaction<
+                Consensus = TxTy<Node::Types>,
+                Pooled = OpPooledTransaction,
+            >,
+        > + Unpin
+        + 'static,
+{
+    type Primitives = CustomNetworkPrimitives;
+
+    async fn build_network(
+        self,
+        ctx: &BuilderContext<Node>,
+        pool: Pool,
+    ) -> Result<NetworkHandle<Self::Primitives>> {
+        let network_config = self.network_config(ctx)?;
+        let network = NetworkManager::builder(network_config).await?;
+        let handle = ctx.start_network(network, pool);
+
+        Ok(handle)
+    }
+}