From 752d6b625247cce290a256ae2513fdcded8c0373 Mon Sep 17 00:00:00 2001
From: Alfredo Garcia <oxarbitrage@gmail.com>
Date: Tue, 8 Feb 2022 20:00:48 -0300
Subject: [PATCH] Add a test for peerset broadcast panic (#3470)

* add a test for peerset always broadcast while there are available peers

* fix minors from review

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* split the test into two

* simplify some code

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 zebra-network/src/peer_set/set/tests/prop.rs | 116 +++++++++++++++++++
 1 file changed, 116 insertions(+)

diff --git a/zebra-network/src/peer_set/set/tests/prop.rs b/zebra-network/src/peer_set/set/tests/prop.rs
index b803444c..f7ab9e12 100644
--- a/zebra-network/src/peer_set/set/tests/prop.rs
+++ b/zebra-network/src/peer_set/set/tests/prop.rs
@@ -166,6 +166,122 @@ proptest! {
             Ok::<_, TestCaseError>(())
         })?;
     }
+
+    /// Test the peerset will always broadcast iff there is at least one
+    /// peer in the set.
+    #[test]
+    fn peerset_always_broadcasts(
+        total_number_of_peers in (2..10usize)
+    ) {
+        // Get a dummy block hash to help us construct a valid request to be broadcasted
+        let block: block::Block = zebra_test::vectors::BLOCK_MAINNET_10_BYTES
+            .zcash_deserialize_into()
+            .unwrap();
+        let block_hash = block::Hash::from(&block);
+
+        // Start the runtime
+        let runtime = zebra_test::init_async();
+        let _guard = runtime.enter();
+
+        // All peers will have the current version
+        let peer_versions = vec![CURRENT_NETWORK_PROTOCOL_VERSION; total_number_of_peers];
+        let peer_versions = PeerVersions {
+            peer_versions,
+        };
+
+        // Get peers and handles
+        let (discovered_peers, mut handles) = peer_versions.mock_peer_discovery();
+        let (minimum_peer_version, _best_tip_height) =
+            MinimumPeerVersion::with_mock_chain_tip(Network::Mainnet);
+
+        runtime.block_on(async move {
+            // Build a peerset
+            let (mut peer_set, _peer_set_guard) = PeerSetBuilder::new()
+                .with_discover(discovered_peers)
+                .with_minimum_peer_version(minimum_peer_version.clone())
+                .build();
+
+            // Remove peers, test broadcast until there is only 1 peer left in the peerset
+            for port in 1u16..total_number_of_peers as u16 {
+                peer_set.remove(&SocketAddr::new([127, 0, 0, 1].into(), port));
+                handles.remove(0);
+
+                // poll the peers
+                check_if_only_up_to_date_peers_are_live(
+                    &mut peer_set,
+                    &mut handles,
+                    CURRENT_NETWORK_PROTOCOL_VERSION,
+                )?;
+
+                // Get the new number of active peers after removal
+                let number_of_peers_to_broadcast = peer_set.number_of_peers_to_broadcast();
+
+                // Send a request to all peers we have now
+                let _ = peer_set.route_broadcast(Request::AdvertiseBlock(block_hash));
+
+                // Check how many peers received the request
+                let mut received = 0;
+                for h in &mut handles {
+                    if let ReceiveRequestAttempt::Request(client_request) = h.try_to_receive_outbound_client_request() {
+                        prop_assert_eq!(client_request.request, Request::AdvertiseBlock(block_hash));
+                        received += 1;
+                    };
+                }
+
+                // Make sure the message is always broadcasted to the right number of peers
+                prop_assert_eq!(received, number_of_peers_to_broadcast);
+            }
+
+            Ok::<_, TestCaseError>(())
+        })?;
+    }
+
+    /// Test the peerset panics if a request is sent and no more peers are available.
+    #[test]
+    #[should_panic(expected = "requests must be routed to at least one peer")]
+    fn panics_when_broadcasting_to_no_peers(
+        total_number_of_peers in (2..10usize)
+    ) {
+        // Get a dummy block hash to help us construct a valid request to be broadcasted
+        let block: block::Block = zebra_test::vectors::BLOCK_MAINNET_10_BYTES
+            .zcash_deserialize_into()
+            .unwrap();
+        let block_hash = block::Hash::from(&block);
+
+        // Start the runtime
+        let runtime = zebra_test::init_async();
+        let _guard = runtime.enter();
+
+        // All peers will have the current version
+        let peer_versions = vec![CURRENT_NETWORK_PROTOCOL_VERSION; total_number_of_peers];
+        let peer_versions = PeerVersions {
+            peer_versions,
+        };
+
+        // Get peers and handles
+        let (discovered_peers, mut handles) = peer_versions.mock_peer_discovery();
+        let (minimum_peer_version, _best_tip_height) =
+            MinimumPeerVersion::with_mock_chain_tip(Network::Mainnet);
+
+        runtime.block_on(async move {
+            // Build a peerset
+            let (mut peer_set, _peer_set_guard) = PeerSetBuilder::new()
+                .with_discover(discovered_peers)
+                .with_minimum_peer_version(minimum_peer_version.clone())
+                .build();
+
+            // Remove peers
+            for port in 1u16..=total_number_of_peers as u16 {
+                peer_set.remove(&SocketAddr::new([127, 0, 0, 1].into(), port));
+                handles.remove(0);
+            }
+
+            // this will panic as expected
+            let _ = peer_set.route_broadcast(Request::AdvertiseBlock(block_hash));
+
+            Ok::<_, TestCaseError>(())
+        })?;
+    }
 }
 
 /// Check if only peers with up-to-date protocol versions are live.