feat: add NetworkProtocols trait (#5543)

bin/reth/src/cli/components.rs

@@ -1,6 +1,6 @@
 //! Components that are used by the node command.
 
-use reth_network::NetworkEvents;
+use reth_network::{NetworkEvents, NetworkProtocols};
 use reth_network_api::{NetworkInfo, Peers};
 use reth_primitives::ChainSpec;
 use reth_provider::{
@@ -48,7 +48,7 @@ pub trait RethNodeComponents: Clone + Send + Sync + 'static {
     /// The transaction pool type
     type Pool: TransactionPool + Clone + Unpin + 'static;
     /// The network type used to communicate with p2p.
-    type Network: NetworkInfo + Peers + NetworkEvents + Clone + 'static;
+    type Network: NetworkInfo + Peers + NetworkProtocols + NetworkEvents + Clone + 'static;
     /// The events type used to create subscriptions.
     type Events: CanonStateSubscriptions + Clone + 'static;
     /// The type that is used to spawn tasks.
@@ -117,7 +117,7 @@ where
     Provider: FullProvider + Clone + 'static,
     Tasks: TaskSpawner + Clone + Unpin + 'static,
     Pool: TransactionPool + Clone + Unpin + 'static,
-    Network: NetworkInfo + Peers + NetworkEvents + Clone + 'static,
+    Network: NetworkInfo + Peers + NetworkProtocols + NetworkEvents + Clone + 'static,
     Events: CanonStateSubscriptions + Clone + 'static,
 {
     type Provider = Provider;

crates/net/network/src/lib.rs

@@ -141,7 +141,7 @@ pub use discovery::{Discovery, DiscoveryEvent};
 pub use fetch::FetchClient;
 pub use manager::{NetworkEvent, NetworkManager};
 pub use message::PeerRequest;
-pub use network::{NetworkEvents, NetworkHandle};
+pub use network::{NetworkEvents, NetworkHandle, NetworkProtocols};
 pub use peers::PeersConfig;
 pub use session::{
     ActiveSessionHandle, ActiveSessionMessage, Direction, PeerInfo, PendingSessionEvent,

crates/net/network/src/manager.rs

@@ -604,6 +604,7 @@ where
                 let peers = self.swarm.state().peers().peers_by_kind(kind);
                 let _ = tx.send(self.swarm.sessions().get_peer_infos_by_ids(peers));
             }
+            NetworkHandleMessage::AddRlpxSubProtocol(proto) => self.add_rlpx_sub_protocol(proto),
         }
     }
 }

crates/net/network/src/network.rs

@@ -1,6 +1,6 @@
 use crate::{
     config::NetworkMode, discovery::DiscoveryEvent, manager::NetworkEvent, message::PeerRequest,
-    peers::PeersHandle, FetchClient,
+    peers::PeersHandle, protocol::RlpxSubProtocol, FetchClient,
 };
 use async_trait::async_trait;
 use parking_lot::Mutex;
@@ -155,6 +155,8 @@ impl NetworkHandle {
     }
 }
 
+// === API Implementations ===
+
 impl NetworkEvents for NetworkHandle {
     fn event_listener(&self) -> UnboundedReceiverStream<NetworkEvent> {
         let (tx, rx) = mpsc::unbounded_channel();
@@ -169,7 +171,11 @@ impl NetworkEvents for NetworkHandle {
     }
 }
 
-// === API Implementations ===
+impl NetworkProtocols for NetworkHandle {
+    fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol) {
+        self.send_message(NetworkHandleMessage::AddRlpxSubProtocol(protocol))
+    }
+}
 
 impl PeersInfo for NetworkHandle {
     fn num_connected_peers(&self) -> usize {
@@ -353,6 +359,12 @@ pub trait NetworkEvents: Send + Sync {
     fn discovery_listener(&self) -> UnboundedReceiverStream<DiscoveryEvent>;
 }
 
+/// Provides access to modify the network's additional protocol handlers.
+pub trait NetworkProtocols: Send + Sync {
+    /// Adds an additional protocol handler to the RLPx sub-protocol list.
+    fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol);
+}
+
 /// Internal messages that can be passed to the  [`NetworkManager`](crate::NetworkManager).
 #[allow(missing_docs)]
 #[derive(Debug)]
@@ -400,4 +412,6 @@ pub(crate) enum NetworkHandleMessage {
     Shutdown(oneshot::Sender<()>),
     /// Add a new listener for `DiscoveryEvent`.
     DiscoveryListener(UnboundedSender<DiscoveryEvent>),
+    /// Add an additional [RlpxSubProtocol].
+    AddRlpxSubProtocol(RlpxSubProtocol),
 }