From c6e016aa139c8363e9b38bbc1ba0dca55700b8a7 Mon Sep 17 00:00:00 2001
From: glozow <gloriajzhao@gmail.com>
Date: Thu, 15 Jul 2021 07:18:18 +0100
Subject: [PATCH] [mempool] check ancestor/descendant limits for packages

When calculating ancestor/descendant counts for transactions in the
package, as a heuristic, count every transaction in the package as an
ancestor and descendant of every other transaction in the package.

This may overestimate, but will not underestimate, the
ancestor/descendant counts. This shortcut still produces an accurate
count for packages of 1 parent + 1 child.
---
 src/txmempool.cpp | 35 +++++++++++++++++++++++++++++++++++
 src/txmempool.h   | 23 +++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index 4a992bf2a44ed..d5a888ac67099 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -200,6 +200,41 @@ bool CTxMemPool::CalculateAncestorsAndCheckLimits(size_t entry_size,
     return true;
 }
 
+bool CTxMemPool::CheckPackageLimits(const Package& package,
+                                    uint64_t limitAncestorCount,
+                                    uint64_t limitAncestorSize,
+                                    uint64_t limitDescendantCount,
+                                    uint64_t limitDescendantSize,
+                                    std::string &errString) const
+{
+    CTxMemPoolEntry::Parents staged_ancestors;
+    size_t total_size = 0;
+    for (const auto& tx : package) {
+        total_size += GetVirtualTransactionSize(*tx);
+        for (const auto& input : tx->vin) {
+            std::optional<txiter> piter = GetIter(input.prevout.hash);
+            if (piter) {
+                staged_ancestors.insert(**piter);
+                if (staged_ancestors.size() + package.size() > limitAncestorCount) {
+                    errString = strprintf("too many unconfirmed parents [limit: %u]", limitAncestorCount);
+                    return false;
+                }
+            }
+        }
+    }
+    // When multiple transactions are passed in, the ancestors and descendants of all transactions
+    // considered together must be within limits even if they are not interdependent. This may be
+    // stricter than the limits for each individual transaction.
+    setEntries setAncestors;
+    const auto ret = CalculateAncestorsAndCheckLimits(total_size, package.size(),
+                                                      setAncestors, staged_ancestors,
+                                                      limitAncestorCount, limitAncestorSize,
+                                                      limitDescendantCount, limitDescendantSize, errString);
+    // It's possible to overestimate the ancestor/descendant totals.
+    if (!ret) errString.insert(0, "possibly ");
+    return ret;
+}
+
 bool CTxMemPool::CalculateMemPoolAncestors(const CTxMemPoolEntry &entry,
                                            setEntries &setAncestors,
                                            uint64_t limitAncestorCount,
diff --git a/src/txmempool.h b/src/txmempool.h
index 71345ffb5d68a..0a84a6e6b12df 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -18,6 +18,7 @@
 #include <coins.h>
 #include <indirectmap.h>
 #include <policy/feerate.h>
+#include <policy/packages.h>
 #include <primitives/transaction.h>
 #include <random.h>
 #include <sync.h>
@@ -700,6 +701,28 @@ class CTxMemPool
      */
     bool CalculateMemPoolAncestors(const CTxMemPoolEntry& entry, setEntries& setAncestors, uint64_t limitAncestorCount, uint64_t limitAncestorSize, uint64_t limitDescendantCount, uint64_t limitDescendantSize, std::string& errString, bool fSearchForParents = true) const EXCLUSIVE_LOCKS_REQUIRED(cs);
 
+    /** Calculate all in-mempool ancestors of a set of transactions not already in the mempool and
+     * check ancestor and descendant limits. Heuristics are used to estimate the ancestor and
+     * descendant count of all entries if the package were to be added to the mempool.  The limits
+     * are applied to the union of all package transactions. For example, if the package has 3
+     * transactions and limitAncestorCount = 25, the union of all 3 sets of ancestors (including the
+     * transactions themselves) must be <= 22.
+     * @param[in]       package                 Transaction package being evaluated for acceptance
+     *                                          to mempool. The transactions need not be direct
+     *                                          ancestors/descendants of each other.
+     * @param[in]       limitAncestorCount      Max number of txns including ancestors.
+     * @param[in]       limitAncestorSize       Max virtual size including ancestors.
+     * @param[in]       limitDescendantCount    Max number of txns including descendants.
+     * @param[in]       limitDescendantSize     Max virtual size including descendants.
+     * @param[out]      errString               Populated with error reason if a limit is hit.
+     */
+    bool CheckPackageLimits(const Package& package,
+                            uint64_t limitAncestorCount,
+                            uint64_t limitAncestorSize,
+                            uint64_t limitDescendantCount,
+                            uint64_t limitDescendantSize,
+                            std::string &errString) const EXCLUSIVE_LOCKS_REQUIRED(cs);
+
     /** Populate setDescendants with all in-mempool descendants of hash.
      *  Assumes that setDescendants includes all in-mempool descendants of anything
      *  already in it.  */