feat: block endless peer stream (#5951)

Description --- Block endless peer stream from remote peer. Block duplicate peers. Motivation and Context --- See #5811 Fixes #5811
tari-project · Nov 14, 2023 · 16b325d · 16b325d
1 parent b99d6f5
commit 16b325d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/comms/dht/src/network_discovery/discovering.rs b/comms/dht/src/network_discovery/discovering.rs
@@ -20,7 +20,7 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use std::convert::TryInto;
+use std::{collections::HashSet, convert::TryInto};
 
 use futures::{stream::FuturesUnordered, Stream, StreamExt};
 use log::*;
@@ -170,14 +170,24 @@ impl Discovering {
                     }),
             })
             .await?;
-
+        let mut counter = 0;
+        let mut peers_received = HashSet::new();
         while let Some(resp) = stream.next().await {
+            counter += 1;
+            if counter > self.params.num_peers_to_request {
+                warn!(target: LOG_TARGET, "Remote peer sent more peers than we requested.");
+                return Err(NetworkDiscoveryError::TooManyPeersReceived);
+            }
             let GetPeersResponse { peer } = resp?;
 
             let peer = peer.ok_or_else(|| NetworkDiscoveryError::EmptyPeerMessageReceived)?;
-            let new_peer = peer
+            let new_peer: UnvalidatedPeerInfo = peer
                 .try_into()
                 .map_err(NetworkDiscoveryError::InvalidPeerDataReceived)?;
+            if !peers_received.insert(new_peer.public_key.clone()) {
+                warn!(target: LOG_TARGET, "Remote peer sent duplicate peer.");
+                return Err(NetworkDiscoveryError::DuplicatePeerReceived);
+            }
             self.validate_and_add_peer(sync_peer, new_peer).await?;
         }
 
@@ -230,7 +240,9 @@ impl Discovering {
                 match &err {
                     NetworkDiscoveryError::EmptyPeerMessageReceived |
                     NetworkDiscoveryError::InvalidPeerDataReceived(_) |
-                    NetworkDiscoveryError::PeerValidationError(_) => {
+                    NetworkDiscoveryError::PeerValidationError(_) |
+                    NetworkDiscoveryError::DuplicatePeerReceived |
+                    NetworkDiscoveryError::TooManyPeersReceived => {
                         self.ban_peer(peer, OffenceSeverity::High, &err).await;
                     },
                     NetworkDiscoveryError::RpcError(rpc_err) if rpc_err.is_caused_by_server() => {

diff --git a/comms/dht/src/network_discovery/error.rs b/comms/dht/src/network_discovery/error.rs
@@ -44,6 +44,10 @@ pub enum NetworkDiscoveryError {
     PeerValidationError(#[from] DhtPeerValidatorError),
     #[error("Sync peer sent empty peer message")]
     EmptyPeerMessageReceived,
+    #[error("Sync peer sent too many peers")]
+    TooManyPeersReceived,
+    #[error("Sync peer sent duplicate peer")]
+    DuplicatePeerReceived,
     #[error("Sync peer sent invalid peer data: {0}")]
     InvalidPeerDataReceived(anyhow::Error),
 }