From 7b3e418950a43ee944faf581c90e1c0b054e727a Mon Sep 17 00:00:00 2001 From: Delweng Date: Mon, 18 Mar 2024 18:51:14 +0800 Subject: [PATCH] feat(net/manager): apply trust-nodes-only for incoming nodes (#7177) Signed-off-by: jsvisa --- bin/reth/src/commands/p2p/mod.rs | 2 +- bin/reth/src/commands/stage/run.rs | 2 +- crates/config/src/config.rs | 17 ++++ crates/net/network/src/peers/manager.rs | 124 ++++++++++++++++++++---- crates/net/network/src/state.rs | 4 + crates/node-builder/src/builder.rs | 2 +- 6 files changed, 127 insertions(+), 24 deletions(-) diff --git a/bin/reth/src/commands/p2p/mod.rs b/bin/reth/src/commands/p2p/mod.rs index 9a429610c..02dce49c4 100644 --- a/bin/reth/src/commands/p2p/mod.rs +++ b/bin/reth/src/commands/p2p/mod.rs @@ -116,7 +116,7 @@ impl Command { eyre::bail!("No trusted nodes. Set trusted peer with `--trusted-peer ` 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); diff --git a/bin/reth/src/commands/stage/run.rs b/bin/reth/src/commands/stage/run.rs index 78792c830..c9ceaefee 100644 --- a/bin/reth/src/commands/stage/run.rs +++ b/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); diff --git a/crates/config/src/config.rs b/crates/config/src/config.rs index 2f603edb1..1ce450824 100644 --- a/crates/config/src/config.rs +++ b/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); + } } diff --git a/crates/net/network/src/peers/manager.rs b/crates/net/network/src/peers/manager.rs index 567f2c4eb..c6e072fcf 100644 --- a/crates/net/network/src/peers/manager.rs +++ b/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, - /// 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)); diff --git a/crates/net/network/src/state.rs b/crates/net/network/src/state.rs index b487a7845..d75a1aaa5 100644 --- a/crates/net/network/src/state.rs +++ b/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) } diff --git a/crates/node-builder/src/builder.rs b/crates/node-builder/src/builder.rs index 9c03f3289..94b1f2692 100644 --- a/crates/node-builder/src/builder.rs +++ b/crates/node-builder/src/builder.rs @@ -159,7 +159,7 @@ impl NodeBuilder { 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");