From 3e0cd2eb3d1b880eee8d082dfb1fd44f11bb2aa9 Mon Sep 17 00:00:00 2001
From: int88 <106391185+int88@users.noreply.github.com>
Date: Sat, 1 Feb 2025 04:14:48 +0800
Subject: [PATCH] integration test for trusted peer only (#14127)

---
 crates/net/network/tests/it/connect.rs | 59 ++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/crates/net/network/tests/it/connect.rs b/crates/net/network/tests/it/connect.rs
index de73f9f28f..b54f2400fb 100644
--- a/crates/net/network/tests/it/connect.rs
+++ b/crates/net/network/tests/it/connect.rs
@@ -577,6 +577,65 @@ async fn test_shutdown() {
     assert_eq!(handle0.num_connected_peers(), 0);
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_trusted_peer_only() {
+    let net = Testnet::create(2).await;
+    let mut handles = net.handles();
+    let handle0 = handles.next().unwrap();
+    let handle1 = handles.next().unwrap();
+
+    drop(handles);
+    let _handle = net.spawn();
+
+    let secret_key = SecretKey::new(&mut rand::thread_rng());
+    let peers_config = PeersConfig::default().with_trusted_nodes_only(true);
+
+    let config = NetworkConfigBuilder::eth(secret_key)
+        .listener_port(0)
+        .disable_discovery()
+        .peer_config(peers_config)
+        .build(NoopProvider::default());
+
+    let network = NetworkManager::new(config).await.unwrap();
+
+    let handle = network.handle().clone();
+    tokio::task::spawn(network);
+
+    // create networkeventstream to get the next session event easily.
+    let events = handle.event_listener();
+    let mut event_stream = NetworkEventStream::new(events);
+
+    // only connect to trusted peers.
+
+    // connect to an untrusted peer should fail.
+    handle.add_peer(*handle0.peer_id(), handle0.local_addr());
+
+    // wait 2 seconds, the number of connection is still 0.
+    tokio::time::sleep(Duration::from_secs(2)).await;
+    assert_eq!(handle.num_connected_peers(), 0);
+
+    // add to trusted peer.
+    handle.add_trusted_peer(*handle0.peer_id(), handle0.local_addr());
+
+    let outgoing_peer_id = event_stream.next_session_established().await.unwrap();
+    assert_eq!(outgoing_peer_id, *handle0.peer_id());
+    assert_eq!(handle.num_connected_peers(), 1);
+
+    // only receive connections from trusted peers.
+
+    handle1.add_peer(*handle.peer_id(), handle0.local_addr());
+
+    // wait 2 seconds, the number of connections is still 1, because peer1 is untrusted.
+    tokio::time::sleep(Duration::from_secs(2)).await;
+    assert_eq!(handle.num_connected_peers(), 1);
+
+    handle1.add_trusted_peer(*handle.peer_id(), handle.local_addr());
+
+    let outgoing_peer_id1 = event_stream.next_session_established().await.unwrap();
+    assert_eq!(outgoing_peer_id1, *handle1.peer_id());
+    assert_eq!(handle.num_connected_peers(), 2);
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_exceed_outgoing_connections() {
     let net = Testnet::create(2).await;