1  // Copyright (c) 2025 ADnet Contributors
 2  // This file is part of the AlphaOS library.
 3  
 4  // Licensed under the Apache License, Version 2.0 (the "License");
 5  // you may not use this file except in compliance with the License.
 6  // You may obtain a copy of the License at:
 7  
 8  // http://www.apache.org/licenses/LICENSE-2.0
 9  
10  // Unless required by applicable law or agreed to in writing, software
11  // distributed under the License is distributed on an "AS IS" BASIS,
12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  // See the License for the specific language governing permissions and
14  // limitations under the License.
15  
16  #[allow(dead_code)]
17  mod common;
18  use common::test_peer::TestPeer;
19  
20  use alphaos_node_network::PeerPoolHandling;
21  use alphaos_node_router::messages::{Message, PeerResponse};
22  use alphaos_node_tcp::P2P;
23  
24  use deadline::deadline;
25  use paste::paste;
26  use pea2pea::{Pea2Pea, protocols::Writing};
27  use std::time::Duration;
28  
29  macro_rules! test_reject_unsolicited_peer_response {
30      ($($node_type:ident),*) => {
31          $(
32              paste! {
33                  #[tokio::test]
34                  async fn [<$node_type _rejects_unsolicited_peer_response>]() {
35                      // Spin up a full node.
36                      let node = $crate::common::node::$node_type().await;
37  
38                      // Spin up a test peer (synthetic node), it doesn't really matter what type it is.
39                      let peer = TestPeer::validator().await;
40                      let peer_addr = peer.node().listening_addr().unwrap();
41  
42                      // Connect the node to the test peer.
43                      node.router().connect(peer_addr).unwrap().await.unwrap();
44  
45                      // Check the peer counts.
46                      let node_clone = node.clone();
47                      deadline!(Duration::from_secs(5), move || node_clone.router().number_of_connected_peers() == 1);
48                      let node_clone = node.clone();
49                      deadline!(Duration::from_secs(5), move || node_clone.tcp().num_connected() == 1);
50                      let peer_clone = peer.clone();
51                      deadline!(Duration::from_secs(5), move || peer_clone.node().num_connected() == 1);
52  
53                      // Check the candidate peers.
54                      assert_eq!(node.router().number_of_candidate_peers(), 0);
55  
56                      let peers = vec![
57                          ("1.1.1.1:1111".parse().unwrap(), None),
58                          ("2.2.2.2:2222".parse().unwrap(), None),
59                      ];
60  
61                      // Send a `PeerResponse` to the node.
62                      assert!(
63                          peer.unicast(
64                              *peer.node().connected_addrs().first().unwrap(),
65                              Message::PeerResponse(PeerResponse { peers: peers.clone() })
66                          )
67                          .is_ok()
68                      );
69  
70                      // Wait for the peer to be disconnected for a protocol violation.
71                      let node_clone = node.clone();
72                      deadline!(Duration::from_secs(5), move || node_clone.router().number_of_connected_peers() == 0);
73  
74                      // Make sure the sent addresses weren't inserted in the candidate peers.
75                      let candidate_peer_addrs = node.router().get_candidate_peers().into_iter().map(|peer| peer.listener_addr).collect::<Vec<_>>();
76                      for (peer, _) in peers {
77                          assert!(!candidate_peer_addrs.contains(&peer));
78                      }
79                  }
80              }
81          )*
82      };
83  }
84  
85  test_reject_unsolicited_peer_response!(client, prover, validator);