From edeedd519e005e030b229ffdc89c4515209b998e Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Tue, 7 Mar 2023 17:56:05 +0100
Subject: [PATCH] fix(net): enforce soft limit when propagating tx hashes
 (#1660)

---
 crates/net/eth-wire/src/types/broadcast.rs | 29 ++++++++++++++++++++++
 crates/net/network/src/transactions.rs     | 19 +++++++++-----
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/crates/net/eth-wire/src/types/broadcast.rs b/crates/net/eth-wire/src/types/broadcast.rs
index 163f16bf8b..b7de075db2 100644
--- a/crates/net/eth-wire/src/types/broadcast.rs
+++ b/crates/net/eth-wire/src/types/broadcast.rs
@@ -152,6 +152,35 @@ impl NewPooledTransactionHashes {
             NewPooledTransactionHashes::Eth68(msg) => msg.hashes.into_iter(),
         }
     }
+
+    /// Shortens the number of hashes in the message, keeping the first `len` hashes and dropping
+    /// the rest. If `len` is greater than the number of hashes, this has no effect.
+    pub fn truncate(&mut self, len: usize) {
+        match self {
+            NewPooledTransactionHashes::Eth66(msg) => msg.0.truncate(len),
+            NewPooledTransactionHashes::Eth68(msg) => {
+                msg.types.truncate(len);
+                msg.sizes.truncate(len);
+                msg.hashes.truncate(len);
+            }
+        }
+    }
+
+    /// Returns true if the message is empty
+    pub fn is_empty(&self) -> bool {
+        match self {
+            NewPooledTransactionHashes::Eth66(msg) => msg.0.is_empty(),
+            NewPooledTransactionHashes::Eth68(msg) => msg.hashes.is_empty(),
+        }
+    }
+
+    /// Returns the number of hashes in the message
+    pub fn len(&self) -> usize {
+        match self {
+            NewPooledTransactionHashes::Eth66(msg) => msg.0.len(),
+            NewPooledTransactionHashes::Eth68(msg) => msg.hashes.len(),
+        }
+    }
 }
 
 impl From<NewPooledTransactionHashes> for EthMessage {
diff --git a/crates/net/network/src/transactions.rs b/crates/net/network/src/transactions.rs
index 39bf239df6..1cf1f0afe9 100644
--- a/crates/net/network/src/transactions.rs
+++ b/crates/net/network/src/transactions.rs
@@ -227,7 +227,7 @@ where
         let max_num_full = (self.peers.len() as f64).sqrt() as usize + 1;
 
         // Note: Assuming ~random~ order due to random state of the peers map hasher
-        for (idx, (peer_id, peer)) in self.peers.iter_mut().enumerate() {
+        for (peer_idx, (peer_id, peer)) in self.peers.iter_mut().enumerate() {
             // filter all transactions unknown to the peer
             let mut hashes = PooledTransactionsHashesBuilder::new(peer.version);
             let mut full_transactions = Vec::new();
@@ -237,20 +237,27 @@ where
                     full_transactions.push(Arc::clone(&tx.transaction));
                 }
             }
-            let hashes = hashes.build();
+            let mut new_pooled_hashes = hashes.build();
 
             if !full_transactions.is_empty() {
-                if idx > max_num_full {
-                    for hash in hashes.iter_hashes().copied() {
+                // determine whether to send full tx objects or hashes.
+                if peer_idx > max_num_full {
+                    // enforce tx soft limit per message for the (unlikely) event the number of
+                    // hashes exceeds it
+                    new_pooled_hashes.truncate(NEW_POOLED_TRANSACTION_HASHES_SOFT_LIMIT);
+
+                    for hash in new_pooled_hashes.iter_hashes().copied() {
                         propagated.0.entry(hash).or_default().push(PropagateKind::Hash(*peer_id));
                     }
                     // send hashes of transactions
-                    self.network.send_transactions_hashes(*peer_id, hashes);
+                    self.network.send_transactions_hashes(*peer_id, new_pooled_hashes);
                 } else {
+                    // TODO ensure max message size
+
                     // send full transactions
                     self.network.send_transactions(*peer_id, full_transactions);
 
-                    for hash in hashes.into_iter_hashes() {
+                    for hash in new_pooled_hashes.into_iter_hashes() {
                         propagated.0.entry(hash).or_default().push(PropagateKind::Full(*peer_id));
                     }
                 }