feat(net/manager): apply trust-nodes-only for incoming nodes (#7177)

Signed-off-by: jsvisa <delweng@gmail.com>

bin/reth/src/commands/p2p/mod.rs

@@ -116,7 +116,7 @@ impl Command {
             eyre::bail!("No trusted nodes. Set trusted peer with `--trusted-peer <enode record>` or set `--trusted-only` to `false`")
         }
 
-        config.peers.connect_trusted_nodes_only = self.trusted_only;
+        config.peers.trusted_nodes_only = self.trusted_only;
 
         let default_secret_key_path = data_dir.p2p_secret_path();
         let secret_key_path = self.p2p_secret_key.clone().unwrap_or(default_secret_key_path);

bin/reth/src/commands/stage/run.rs

@@ -172,7 +172,7 @@ impl Command {
                     let consensus = Arc::new(BeaconConsensus::new(self.chain.clone()));
 
                     let mut config = config;
-                    config.peers.connect_trusted_nodes_only = self.network.trusted_only;
+                    config.peers.trusted_nodes_only = self.network.trusted_only;
                     if !self.network.trusted_peers.is_empty() {
                         self.network.trusted_peers.iter().for_each(|peer| {
                             config.peers.trusted_nodes.insert(*peer);

crates/config/src/config.rs

@@ -671,4 +671,21 @@ nanos = 0
 #";
         let _conf: Config = toml::from_str(alpha_0_0_19).unwrap();
     }
+
+    #[test]
+    fn test_conf_trust_nodes_only() {
+        let trusted_nodes_only = r"#
+[peers]
+trusted_nodes_only = true
+#";
+        let conf: Config = toml::from_str(trusted_nodes_only).unwrap();
+        assert!(conf.peers.trusted_nodes_only);
+
+        let trusted_nodes_only = r"#
+[peers]
+connect_trusted_nodes_only = true
+#";
+        let conf: Config = toml::from_str(trusted_nodes_only).unwrap();
+        assert!(conf.peers.trusted_nodes_only);
+    }
 }

crates/net/network/src/peers/manager.rs

@@ -116,8 +116,9 @@ pub struct PeersManager {
     /// How long peers to which we could not connect for non-fatal reasons, e.g.
     /// [`DisconnectReason::TooManyPeers`], are put in time out.
     backoff_durations: PeerBackoffDurations,
-    /// If non-trusted peers should be connected to
-    connect_trusted_nodes_only: bool,
+    /// If non-trusted peers should be connected to, or the connection from non-trusted
+    /// incoming peers should be accepted.
+    trusted_nodes_only: bool,
     /// Timestamp of the last time [Self::tick] was called.
     last_tick: Instant,
     /// Maximum number of backoff attempts before we give up on a peer and dropping.
@@ -137,7 +138,7 @@ impl PeersManager {
             ban_duration,
             backoff_durations,
             trusted_nodes,
-            connect_trusted_nodes_only,
+            trusted_nodes_only,
             basic_nodes,
             max_backoff_count,
         } = config;
@@ -173,7 +174,7 @@ impl PeersManager {
             backed_off_peers: Default::default(),
             ban_duration,
             backoff_durations,
-            connect_trusted_nodes_only,
+            trusted_nodes_only,
             last_tick: Instant::now(),
             max_backoff_count,
             net_connection_state: NetworkConnectionState::default(),
@@ -269,30 +270,38 @@ impl PeersManager {
     /// If the reputation of the peer is below the `BANNED_REPUTATION` threshold, a disconnect will
     /// be scheduled.
     pub(crate) fn on_incoming_session_established(&mut self, peer_id: PeerId, addr: SocketAddr) {
+        self.connection_info.decr_pending_in();
+
         // we only need to check the peer id here as the ip address will have been checked at
-        // on_inbound_pending_session. We also check if the peer is in the backoff list here.
+        // on_incoming_pending_session. We also check if the peer is in the backoff list here.
         if self.ban_list.is_banned_peer(&peer_id) {
             self.queued_actions.push_back(PeerAction::DisconnectBannedIncoming { peer_id });
             return
         }
 
+        // check if the peer is trustable or not
+        let mut is_trusted = self.trusted_peer_ids.contains(&peer_id);
+        if self.trusted_nodes_only && !is_trusted {
+            self.queued_actions.push_back(PeerAction::DisconnectUntrustedIncoming { peer_id });
+            return
+        }
+
         // start a new tick, so the peer is not immediately rewarded for the time since last tick
         self.tick();
 
         let has_in_capacity = self.connection_info.has_in_capacity();
-        self.connection_info.decr_pending_in();
         self.connection_info.inc_in();
 
         match self.peers.entry(peer_id) {
             Entry::Occupied(mut entry) => {
-                let value = entry.get_mut();
-                if value.is_banned() {
+                let peer = entry.get_mut();
+                if peer.is_banned() {
                     self.queued_actions.push_back(PeerAction::DisconnectBannedIncoming { peer_id });
                     return
                 }
-                value.state = PeerConnectionState::In;
+                peer.state = PeerConnectionState::In;
 
-                let is_trusted = value.is_trusted() || self.trusted_peer_ids.contains(&peer_id);
+                is_trusted = is_trusted || peer.is_trusted();
 
                 // if a peer is not trusted and we don't have capacity for more inbound connections,
                 // disconnecting the peer
@@ -311,8 +320,6 @@ impl PeersManager {
                 entry.insert(peer);
                 self.queued_actions.push_back(PeerAction::PeerAdded(peer_id));
 
-                let is_trusted = self.trusted_peer_ids.contains(&peer_id);
-
                 // disconnect the peer if we don't have capacity for more inbound connections
                 if !is_trusted && !has_in_capacity {
                     self.queued_actions.push_back(PeerAction::Disconnect {
@@ -727,7 +734,7 @@ impl PeersManager {
     /// Peers that are `trusted`, see [PeerKind], are prioritized as long as they're not currently
     /// marked as banned or backed off.
     ///
-    /// If `connect_trusted_nodes_only` is enabled, see [PeersConfig], then this will only consider
+    /// If `trusted_nodes_only` is enabled, see [PeersConfig], then this will only consider
     /// `trusted` peers.
     ///
     /// Returns `None` if no peer is available.
@@ -736,7 +743,7 @@ impl PeersManager {
             !peer.is_backed_off() &&
                 !peer.is_banned() &&
                 peer.state.is_unconnected() &&
-                (!self.connect_trusted_nodes_only || peer.is_trusted())
+                (!self.trusted_nodes_only || peer.is_trusted())
         });
 
         // keep track of the best peer, if there's one
@@ -1190,6 +1197,11 @@ pub enum PeerAction {
         /// The peer ID of the established connection.
         peer_id: PeerId,
     },
+    /// Disconnect an untrusted incoming connection when trust-node-only is enabled.
+    DisconnectUntrustedIncoming {
+        /// The peer ID.
+        peer_id: PeerId,
+    },
     /// Ban the peer in discovery.
     DiscoveryBanPeerId {
         /// The peer ID.
@@ -1228,8 +1240,9 @@ pub struct PeersConfig {
     pub refill_slots_interval: Duration,
     /// Trusted nodes to connect to.
     pub trusted_nodes: HashSet<NodeRecord>,
-    /// Connect to trusted nodes only?
-    pub connect_trusted_nodes_only: bool,
+    /// Connect to or accept from trusted nodes only?
+    #[cfg_attr(feature = "serde", serde(alias = "connect_trusted_nodes_only"))]
+    pub trusted_nodes_only: bool,
     /// Maximum number of backoff attempts before we give up on a peer and dropping.
     ///
     /// The max time spent of a peer before it's removed from the set is determined by the
@@ -1271,7 +1284,7 @@ impl Default for PeersConfig {
             ban_duration: Duration::from_secs(60 * 60 * 12),
             backoff_durations: Default::default(),
             trusted_nodes: Default::default(),
-            connect_trusted_nodes_only: false,
+            trusted_nodes_only: false,
             basic_nodes: Default::default(),
             max_backoff_count: 5,
         }
@@ -1344,8 +1357,8 @@ impl PeersConfig {
     }
 
     /// Connect only to trusted nodes.
-    pub fn with_connect_trusted_nodes_only(mut self, trusted_only: bool) -> Self {
-        self.connect_trusted_nodes_only = trusted_only;
+    pub fn with_trusted_nodes_only(mut self, trusted_only: bool) -> Self {
+        self.trusted_nodes_only = trusted_only;
         self
     }
 
@@ -2346,7 +2359,9 @@ mod tests {
         match a {
             Ok(_) => panic!(),
             Err(err) => match err {
-                super::InboundConnectionError::IpBanned {} => {}
+                super::InboundConnectionError::IpBanned {} => {
+                    assert_eq!(peer_manager.connection_info.num_pending_in, 0)
+                }
             },
         }
     }
@@ -2359,7 +2374,14 @@ mod tests {
         let ban_list = BanList::new(vec![given_peer_id], HashSet::new());
         let config = PeersConfig::test().with_ban_list(ban_list);
         let mut peer_manager = PeersManager::new(config);
+        assert!(peer_manager.on_incoming_pending_session(socket_addr.ip()).is_ok());
+        // non-trusted nodes should also increase pending_in
+        assert_eq!(peer_manager.connection_info.num_pending_in, 1);
         peer_manager.on_incoming_session_established(given_peer_id, socket_addr);
+        // after the connection is established, the peer should be removed, the num_pending_in
+        // should be decreased, and the num_inbound should not be increased
+        assert_eq!(peer_manager.connection_info.num_pending_in, 0);
+        assert_eq!(peer_manager.connection_info.num_inbound, 0);
 
         let Some(PeerAction::DisconnectBannedIncoming { peer_id }) =
             peer_manager.queued_actions.pop_front()
@@ -2457,7 +2479,7 @@ mod tests {
                 udp_port: 8008,
                 id: trusted_peer,
             }]))
-            .with_connect_trusted_nodes_only(true);
+            .with_trusted_nodes_only(true);
         let mut peers = PeersManager::new(config);
 
         let basic_peer = PeerId::random();
@@ -2484,6 +2506,66 @@ mod tests {
         .await;
     }
 
+    #[tokio::test]
+    async fn test_incoming_with_trusted_nodes_only() {
+        let trusted_peer = PeerId::random();
+        let config = PeersConfig::test()
+            .with_trusted_nodes(HashSet::from([NodeRecord {
+                address: IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)),
+                tcp_port: 8008,
+                udp_port: 8008,
+                id: trusted_peer,
+            }]))
+            .with_trusted_nodes_only(true);
+        let mut peers = PeersManager::new(config);
+
+        let basic_peer = PeerId::random();
+        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
+        assert!(peers.on_incoming_pending_session(basic_sock.ip()).is_ok());
+        // non-trusted nodes should also increase pending_in
+        assert_eq!(peers.connection_info.num_pending_in, 1);
+        peers.on_incoming_session_established(basic_peer, basic_sock);
+        // after the connection is established, the peer should be removed, the num_pending_in
+        // should be decreased, and the num_inbound mut not be increased
+        assert_eq!(peers.connection_info.num_pending_in, 0);
+        assert_eq!(peers.connection_info.num_inbound, 0);
+
+        let Some(PeerAction::DisconnectUntrustedIncoming { peer_id }) =
+            peers.queued_actions.pop_front()
+        else {
+            panic!()
+        };
+        assert_eq!(basic_peer, peer_id);
+        assert!(!peers.peers.contains_key(&basic_peer));
+    }
+
+    #[tokio::test]
+    async fn test_incoming_without_trusted_nodes_only() {
+        let trusted_peer = PeerId::random();
+        let config = PeersConfig::test()
+            .with_trusted_nodes(HashSet::from([NodeRecord {
+                address: IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)),
+                tcp_port: 8008,
+                udp_port: 8008,
+                id: trusted_peer,
+            }]))
+            .with_trusted_nodes_only(false);
+        let mut peers = PeersManager::new(config);
+
+        let basic_peer = PeerId::random();
+        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
+        assert!(peers.on_incoming_pending_session(basic_sock.ip()).is_ok());
+
+        // non-trusted nodes should also increase pending_in
+        assert_eq!(peers.connection_info.num_pending_in, 1);
+        peers.on_incoming_session_established(basic_peer, basic_sock);
+        // after the connection is established, the peer should be removed, the num_pending_in
+        // should be decreased, and the num_inbound must be increased
+        assert_eq!(peers.connection_info.num_pending_in, 0);
+        assert_eq!(peers.connection_info.num_inbound, 1);
+        assert!(peers.peers.contains_key(&basic_peer));
+    }
+
     #[tokio::test]
     async fn test_tick() {
         let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));

crates/net/network/src/state.rs

@@ -319,6 +319,10 @@ where
                 self.state_fetcher.on_pending_disconnect(&peer_id);
                 self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason: None });
             }
+            PeerAction::DisconnectUntrustedIncoming { peer_id } => {
+                self.state_fetcher.on_pending_disconnect(&peer_id);
+                self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason: None });
+            }
             PeerAction::DiscoveryBanPeerId { peer_id, ip_addr } => {
                 self.ban_discovery(peer_id, ip_addr)
             }

crates/node-builder/src/builder.rs

@@ -159,7 +159,7 @@ impl<DB, State> NodeBuilder<DB, State> {
         info!(target: "reth::cli", path = ?config_path, "Configuration loaded");
 
         // Update the config with the command line arguments
-        config.peers.connect_trusted_nodes_only = self.config.network.trusted_only;
+        config.peers.trusted_nodes_only = self.config.network.trusted_only;
 
         if !self.config.network.trusted_peers.is_empty() {
             info!(target: "reth::cli", "Adding trusted nodes");