feat(net): add peer_id/ip blacklist (#366)

crates/net/network/tests/it/connect.rs

@@ -1,12 +1,18 @@
 //! Connection tests
 
+use crate::NetworkEventStream;
+
 use super::testnet::Testnet;
+use enr::EnrPublicKey;
+use ethers_core::utils::Geth;
+use ethers_providers::{Http, Middleware, Provider};
 use futures::StreamExt;
 use reth_discv4::{bootnodes::mainnet_nodes, Discv4Config};
-use reth_network::{NetworkConfig, NetworkEvent, NetworkManager};
+use reth_network::{BanList, NetworkConfig, NetworkEvent, NetworkManager, PeersConfig};
+use reth_primitives::PeerId;
 use reth_provider::test_utils::TestApi;
 use secp256k1::SecretKey;
-use std::{collections::HashSet, sync::Arc};
+use std::{collections::HashSet, net::SocketAddr, sync::Arc};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_establish_connections() {
@@ -80,3 +86,50 @@ async fn test_connect_with_boot_nodes() {
         dbg!(ev);
     }
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_incoming_node_id_blacklist() {
+    reth_tracing::init_tracing();
+    let secret_key = SecretKey::new(&mut rand::thread_rng());
+
+    let reth_socket = SocketAddr::new([127, 0, 0, 1].into(), 30305);
+
+    // instantiate geth and add ourselves as a peer
+    let temp_dir = tempfile::tempdir().unwrap().into_path();
+    let geth = Geth::new().data_dir(temp_dir).disable_discovery().spawn();
+    let geth_endpoint = SocketAddr::new([127, 0, 0, 1].into(), geth.port());
+    let provider = Provider::<Http>::try_from(format!("http://{geth_endpoint}")).unwrap();
+
+    // get the peer id we should be expecting
+    let geth_peer_id: PeerId =
+        provider.node_info().await.unwrap().enr.public_key().encode_uncompressed().into();
+
+    let ban_list = BanList::new(HashSet::from_iter(vec![geth_peer_id]), HashSet::default());
+    let peer_config = PeersConfig::default().with_ban_list(ban_list);
+
+    let config = NetworkConfig::builder(Arc::new(TestApi::default()), secret_key)
+        .listener_addr(reth_socket)
+        .peer_config(peer_config)
+        .build();
+
+    let network = NetworkManager::new(config).await.unwrap();
+
+    let handle = network.handle().clone();
+    tokio::task::spawn(network);
+
+    // make geth connect to us
+    let our_peer_id = handle.peer_id();
+    let our_enode = format!("enode://{}@{}", hex::encode(our_peer_id.0), reth_socket);
+
+    provider.add_peer(our_enode).await.unwrap();
+
+    let events = handle.event_listener();
+    let mut event_stream = NetworkEventStream::new(events);
+
+    // check for session to be opened
+    let incoming_peer_id = event_stream.next_session_established().await.unwrap();
+    assert_eq!(incoming_peer_id, geth_peer_id);
+    // check to see that the session was closed
+    let incoming_peer_id = event_stream.next_session_closed().await.unwrap();
+    assert_eq!(incoming_peer_id, geth_peer_id);
+}

crates/net/network/tests/it/testnet.rs

@@ -289,6 +289,15 @@ impl NetworkEventStream {
         Self { inner }
     }
 
+    pub async fn next_session_closed(&mut self) -> Option<PeerId> {
+        while let Some(ev) = self.inner.next().await {
+            match ev {
+                NetworkEvent::SessionClosed { peer_id } => return Some(peer_id),
+                NetworkEvent::SessionEstablished { .. } => continue,
+            }
+        }
+        None
+    }
     /// Awaits the next event for an established session
     pub async fn next_session_established(&mut self) -> Option<PeerId> {
         while let Some(ev) = self.inner.next().await {