From e008afdb3c5dbeb7aaff8ed13b84520a159b38fc Mon Sep 17 00:00:00 2001
From: glozow <gloriajzhao@gmail.com>
Date: Mon, 19 Aug 2024 10:58:20 +0100
Subject: [PATCH] [unit test] orphan dos

---
 src/test/fuzz/txdownloadman.cpp |   1 -
 src/test/orphanage_tests.cpp    | 116 +++++++++++++++++++++++++++++++-
 2 files changed, 115 insertions(+), 2 deletions(-)

diff --git a/src/test/fuzz/txdownloadman.cpp b/src/test/fuzz/txdownloadman.cpp
index 24ec00cc12762d..b3ac1af8e59c9f 100644
--- a/src/test/fuzz/txdownloadman.cpp
+++ b/src/test/fuzz/txdownloadman.cpp
@@ -259,7 +259,6 @@ static void CheckInvariants(const node::TxDownloadImpl& txdownload_impl, size_t
             // We should never have more than the maximum in-flight requests out for a peer.
             Assert(txdownload_impl.m_txrequest.CountInFlight(peer) <= node::MAX_PEER_TX_REQUEST_IN_FLIGHT);
             Assert(txdownload_impl.m_orphan_resolution_tracker.Count(peer) <= node::MAX_ORPHAN_RESOLUTIONS);
-            Assert(txdownload_impl.m_orphanage.BytesFromPeer(peer) <= node::MAX_ORPHAN_BYTES_PREFERRED);
         }
     }
 }
diff --git a/src/test/orphanage_tests.cpp b/src/test/orphanage_tests.cpp
index 6fb0ff96ee3352..ebe6069f331ec2 100644
--- a/src/test/orphanage_tests.cpp
+++ b/src/test/orphanage_tests.cpp
@@ -5,6 +5,8 @@
 #include <arith_uint256.h>
 #include <primitives/transaction.h>
 #include <consensus/validation.h>
+#include <node/txdownload_impl.h>
+#include <node/txdownloadman.h>
 #include <pubkey.h>
 #include <script/sign.h>
 #include <script/signingprovider.h>
@@ -45,15 +47,36 @@ static void MakeNewKeyWithFastRandomContext(CKey& key, FastRandomContext& rand_c
     assert(key.IsValid());
 }
 
+static CTransactionRef MakeLargeOrphan(FastRandomContext& det_rand)
+{
+    CKey key;
+    MakeNewKeyWithFastRandomContext(key, det_rand);
+    CMutableTransaction tx;
+    tx.vout.resize(1);
+    tx.vout[0].nValue = CENT;
+    tx.vout[0].scriptPubKey = GetScriptForDestination(PKHash(key.GetPubKey()));
+    tx.vin.resize(80);
+    for (unsigned int j = 0; j < tx.vin.size(); j++) {
+        tx.vin[j].prevout.n = j;
+        tx.vin[j].prevout.hash = Txid::FromUint256(det_rand.rand256());
+        tx.vin[j].scriptWitness.stack.reserve(100);
+        for (int i = 0; i < 100; ++i) {
+            tx.vin[j].scriptWitness.stack.push_back(std::vector<unsigned char>(j));
+        }
+    }
+    return MakeTransactionRef(tx);
+}
+
 // Creates a transaction with 2 outputs. Spends all outpoints. If outpoints is empty, spends a random one.
 static CTransactionRef MakeTransactionSpending(const std::vector<COutPoint>& outpoints, FastRandomContext& det_rand)
 {
+    static uint32_t num = 0;
     CKey key;
     MakeNewKeyWithFastRandomContext(key, det_rand);
     CMutableTransaction tx;
     // If no outpoints are given, create a random one.
     if (outpoints.empty()) {
-        tx.vin.emplace_back(Txid::FromUint256(det_rand.rand256()), 0);
+        tx.vin.emplace_back(Txid::FromUint256(det_rand.rand256()), num++);
     } else {
         for (const auto& outpoint : outpoints) {
             tx.vin.emplace_back(outpoint);
@@ -556,4 +579,95 @@ BOOST_AUTO_TEST_CASE(peer_worksets)
         }
     }
 }
+
+BOOST_AUTO_TEST_CASE(orphan_peer_dos)
+{
+    const NodeId peer_normal_pref{1};
+    const NodeId peer_normal_nonpref{2};
+    std::vector<NodeId> peer_spammers{3, 4};
+
+    const unsigned int max_orphan_count = 100;
+    FastRandomContext det_rand{true};
+    node::TxDownloadImpl txdownload_impl{node::TxDownloadOptions{*m_node.mempool, det_rand, max_orphan_count}};
+
+    txdownload_impl.ConnectedPeer(peer_normal_pref, node::TxDownloadConnectionInfo{/*m_preferred=*/true, /*m_relay_permissions=*/false, /*m_wtxid_relay=*/true});
+    txdownload_impl.ConnectedPeer(peer_normal_nonpref, node::TxDownloadConnectionInfo{/*m_preferred=*/false, /*m_relay_permissions=*/false, /*m_wtxid_relay=*/true});
+
+    for (auto peer_dos : peer_spammers) {
+        txdownload_impl.ConnectedPeer(peer_dos, node::TxDownloadConnectionInfo{/*m_preferred=*/false, /*m_relay_permissions=*/false, /*m_wtxid_relay=*/true});
+    }
+
+    // Resuable TxValidationState indicating the transaction is an orphan.
+    TxValidationState state_missing_inputs;
+    state_missing_inputs.Invalid(TxValidationResult::TX_MISSING_INPUTS, "");
+
+    // Set time to now
+    auto start_time = GetTime<std::chrono::seconds>();
+    SetMockTime(start_time);
+
+    // Add a "normal" orphan. Updates requests_to_expect for later checking.
+    auto add_orphan = [&](NodeId peer, FastRandomContext& det_rand, std::vector<GenTxid>& requests_to_expect) {
+        const auto parent_txid = det_rand.rand256();
+        const auto orphan_tx = MakeTransactionSpending({{Txid::FromUint256(parent_txid), 0}}, det_rand);
+        txdownload_impl.MempoolRejectedTx(orphan_tx, state_missing_inputs, peer, /*maybe_add_new_orphan=*/true);
+        BOOST_CHECK(txdownload_impl.m_orphanage.HaveTxAndPeer(orphan_tx->GetWitnessHash(), peer));
+        requests_to_expect.emplace_back(GenTxid::Txid(parent_txid));
+    };
+
+    // Send orphans from normal peers
+    std::vector<GenTxid> requests_pref;
+    std::vector<GenTxid> requests_nonpref;
+
+    add_orphan(peer_normal_pref, det_rand, requests_pref);
+    add_orphan(peer_normal_nonpref, det_rand, requests_nonpref);
+
+    // Send spam:
+    for (auto peer_dos : peer_spammers) {
+        if (peer_dos % 2) {
+            // Odd peers spam by sending a lot of orphans
+            for (unsigned int i = 0; i < max_orphan_count; ++i) {
+                const auto fake_orphan = MakeTransactionSpending({}, det_rand);
+                txdownload_impl.MempoolRejectedTx(fake_orphan, state_missing_inputs, peer_dos, /*maybe_add_new_orphan=*/true);
+                // orphans are added until the limit is reached
+                BOOST_CHECK_EQUAL(txdownload_impl.m_orphanage.HaveTx(fake_orphan->GetWitnessHash()), i < node::MAX_ORPHAN_RESOLUTIONS);
+            }
+        } else {
+            // Even peers spam by sending a large amount of orphan bytes
+            unsigned int total_orphan_bytes = 0;
+            for (int i = 0; i < 20; ++i) {
+                auto large_orphan = MakeLargeOrphan(det_rand);
+                txdownload_impl.MempoolRejectedTx(large_orphan, state_missing_inputs, peer_dos, /*maybe_add_new_orphan=*/true);
+
+                // Ensure this tx is within max standard size but is large, i.e. will reach the
+                // MAX_ORPHAN_BYTES_NONPREFERRED limit before the MAX_ORPHAN_RESOLUTIONS limit.
+                auto orphan_bytes = large_orphan->GetTotalSize();
+                BOOST_CHECK(orphan_bytes <= MAX_STANDARD_TX_WEIGHT);
+                BOOST_CHECK(orphan_bytes * node::MAX_ORPHAN_RESOLUTIONS > node::MAX_ORPHAN_BYTES_NONPREFERRED);
+
+                // all orphans are added except the one that busts the limit.
+                BOOST_CHECK_EQUAL(txdownload_impl.m_orphanage.HaveTx(large_orphan->GetWitnessHash()),
+                                  total_orphan_bytes <= node::MAX_ORPHAN_BYTES_NONPREFERRED);
+                total_orphan_bytes += orphan_bytes;
+            }
+        }
+
+        // After each spam round, send another orphan from each normal peer.
+        add_orphan(peer_normal_pref, det_rand, requests_pref);
+        add_orphan(peer_normal_nonpref, det_rand, requests_nonpref);
+    }
+
+    add_orphan(peer_normal_pref, det_rand, requests_pref);
+    add_orphan(peer_normal_nonpref, det_rand, requests_nonpref);
+
+    // No evictions happen in this test.
+    BOOST_CHECK(max_orphan_count > txdownload_impl.m_orphanage.Size());
+
+    // Check that txdownload still remembers to schedule the "normal" orphan resolutions after the DoSy peers' spam.
+    const auto requests1 = txdownload_impl.GetRequestsToSend(peer_normal_pref, start_time + 10s);
+    BOOST_CHECK(requests1 == requests_pref);
+
+    const auto requests2 = txdownload_impl.GetRequestsToSend(peer_normal_nonpref, start_time + 10s);
+    BOOST_CHECK(requests2 == requests_nonpref);
+
+}
 BOOST_AUTO_TEST_SUITE_END()