fix: bound pending DKG message retention per proTxHash

src/llmq/dkgsessionhandler.cpp

@@ -6,13 +6,14 @@
 
 #include <logging.h>
 #include <uint256.h>
+#include <util/check.h>
 
 #include <stdexcept>
 
 namespace llmq {
 CDKGSessionHandler::CDKGSessionHandler(const Consensus::LLMQParams& _params) :
     params{_params},
-    // we allow size*2 messages as we need to make sure we see bad behavior (double messages)
+    // we allow size*2 messages per proTx as we need to make sure we see bad behavior (double messages)
     pendingContributions{(size_t)_params.size * 2},
     pendingComplaints{(size_t)_params.size * 2},
     pendingJustifications{(size_t)_params.size * 2},
@@ -25,22 +26,32 @@ CDKGSessionHandler::CDKGSessionHandler(const Consensus::LLMQParams& _params) :
 
 CDKGSessionHandler::~CDKGSessionHandler() = default;
 
-void CDKGPendingMessages::PushPendingMessage(NodeId from, std::shared_ptr<CDataStream> pm, const uint256& hash)
+void CDKGPendingMessages::PushPendingMessage(NodeId from, const uint256& sender_protx,
+                                             std::shared_ptr<CDataStream> pm, const uint256& hash)
 {
     LOCK(cs_messages);
 
-    if (messagesPerNode[from] >= maxMessagesPerNode) {
-        // TODO ban?
-        LogPrint(BCLog::LLMQ_DKG, "CDKGPendingMessages::%s -- too many messages, peer=%d\n", __func__, from);
+    // Check duplicates before the quota so resent hashes don't burn budget
+    if (seenMessages.count(hash) != 0) {
+        LogPrint(BCLog::LLMQ_DKG, "CDKGPendingMessages::%s -- already seen %s, peer=%d\n", __func__, hash.ToString(), from);
         return;
     }
-    messagesPerNode[from]++;
 
-    if (!seenMessages.emplace(hash).second) {
-        LogPrint(BCLog::LLMQ_DKG, "CDKGPendingMessages::%s -- already seen %s, peer=%d\n", __func__, hash.ToString(), from);
+    // Callers always pass an identity (the MNAuth gate for remote, our own
+    // proTxHash for local); drop rather than share a null-keyed quota bucket
+    if (!Assume(!sender_protx.IsNull())) {
+        return;
+    }
+    auto& count = messagesPerProTx[sender_protx];
+    if (count >= maxMessagesPerProTx) {
+        // TODO ban?
+        LogPrint(BCLog::LLMQ_DKG, "CDKGPendingMessages::%s -- too many messages from %s, peer=%d\n", __func__,
+                 sender_protx.ToString(), from);
         return;
     }
+    count++;
 
+    seenMessages.emplace(hash);
     pendingMessages.emplace_back(std::make_pair(from, std::move(pm)));
 }
 
@@ -67,7 +78,7 @@ void CDKGPendingMessages::Clear()
 {
     LOCK(cs_messages);
     pendingMessages.clear();
-    messagesPerNode.clear();
+    messagesPerProTx.clear();
     seenMessages.clear();
 }
 

src/llmq/dkgsessionhandler.h

@@ -9,7 +9,6 @@
 #include <sync.h>
 
 #include <list>
-#include <map>
 #include <memory>
 #include <optional>
 #include <string_view>
@@ -55,24 +54,31 @@ class CDKGPendingMessages
     using BinaryMessage = std::pair<NodeId, std::shared_ptr<CDataStream>>;
 
 private:
-    const size_t maxMessagesPerNode;
+    const size_t maxMessagesPerProTx;
     mutable Mutex cs_messages;
     std::list<BinaryMessage> pendingMessages GUARDED_BY(cs_messages);
-    std::map<NodeId, size_t> messagesPerNode GUARDED_BY(cs_messages);
+    // Keyed by proTxHash rather than NodeId so the quota survives reconnects,
+    // and cumulative for the round (not refunded on pop) so draining the queue
+    // does not regain retention slots. MNAuth pins keys to registered MNs.
+    Uint256HashMap<size_t> messagesPerProTx GUARDED_BY(cs_messages);
     Uint256HashSet seenMessages GUARDED_BY(cs_messages);
 
 public:
-    explicit CDKGPendingMessages(size_t _maxMessagesPerNode) :
-        maxMessagesPerNode(_maxMessagesPerNode) {};
+    explicit CDKGPendingMessages(size_t _maxMessagesPerProTx) :
+        maxMessagesPerProTx(_maxMessagesPerProTx) {};
 
     /**
      * Enqueue a serialized DKG message under @p from with content hash @p hash.
-     * Caller is responsible for hashing the payload and (for real peers)
-     * routing the erase-request to PeerManager. Drops the message silently on
-     * per-node capacity overflow or duplicate hash.
+     * @p sender_protx keys the per-proTx quota: the sender's MNAuth-verified
+     * proTxHash for remote messages, or this node's own for messages it
+     * produced itself (@p from == -1). Drops the message silently on quota
+     * overflow or duplicate hash; quota-dropped messages are not marked seen,
+     * so another peer with budget can re-deliver them. Caller is responsible
+     * for hashing the payload and (for real peers) routing the erase-request
+     * to PeerManager.
      */
-    void PushPendingMessage(NodeId from, std::shared_ptr<CDataStream> pm, const uint256& hash)
-        EXCLUSIVE_LOCKS_REQUIRED(!cs_messages);
+    void PushPendingMessage(NodeId from, const uint256& sender_protx, std::shared_ptr<CDataStream> pm,
+                            const uint256& hash) EXCLUSIVE_LOCKS_REQUIRED(!cs_messages);
 
     std::list<BinaryMessage> PopPendingMessages(size_t maxCount) EXCLUSIVE_LOCKS_REQUIRED(!cs_messages);
     bool HasSeen(const uint256& hash) const EXCLUSIVE_LOCKS_REQUIRED(!cs_messages);

src/llmq/net_dkg.cpp

@@ -256,14 +256,14 @@ void RelayInvToParticipants(const CDKGSession& session, const CConnman& connman,
 }
 
 template <typename Message>
-void EnqueueOwn(CDKGPendingMessages& pending, const Message& msg)
+void EnqueueOwn(CDKGPendingMessages& pending, const uint256& own_protx, const Message& msg)
 {
     CDataStream ds(SER_NETWORK, PROTOCOL_VERSION);
     ds << msg;
     auto pm = std::make_shared<CDataStream>(std::move(ds));
     CHashWriter hw(SER_GETHASH, 0);
     hw.write(AsWritableBytes(Span{*pm}));
-    pending.PushPendingMessage(/*from=*/-1, std::move(pm), hw.GetHash());
+    pending.PushPendingMessage(/*from=*/-1, own_protx, std::move(pm), hw.GetHash());
 }
 
 template <typename Message>
@@ -392,7 +392,9 @@ void NetDKG::ProcessMessage(CNode& pfrom, const std::string& msg_type, CDataStre
     // attacker-controlled payloads, so they must originate from an MNAuth-verified
     // masternode. qwatch is unauthenticated (any peer can set it via QWATCH) and is
     // only meaningful for pull/observation paths; it must not bypass this gate.
-    if (pfrom.GetVerifiedProRegTxHash().IsNull()) {
+    // Read once: the same value keys the retention quota below
+    const uint256 sender_protx = pfrom.GetVerifiedProRegTxHash();
+    if (sender_protx.IsNull()) {
         m_peer_manager->PeerMisbehaving(pfrom.GetId(), 10, "DKG message from non-verified peer");
         return;
     }
@@ -524,7 +526,7 @@ void NetDKG::ProcessMessage(CNode& pfrom, const std::string& msg_type, CDataStre
             break;
         }
         Assume(pending != nullptr);
-        pending->PushPendingMessage(from, std::move(pm), hash);
+        pending->PushPendingMessage(from, sender_protx, std::move(pm), hash);
     });
     if (!dispatched) {
         LogPrintf("NetDKG -- no session handlers for quorumIndex [%d]\n", quorumIndex);
@@ -744,7 +746,7 @@ void NetDKG::HandleDKGRound(ActiveDKGSessionHandler& handler)
     // Contribute
     auto fContributeStart = [curSession, &handler]() {
         if (auto qc = curSession->Contribute(); qc) {
-            EnqueueOwn(handler.pendingContributions, *qc);
+            EnqueueOwn(handler.pendingContributions, curSession->ProTx(), *qc);
         }
     };
     auto fContributeWait = [this, curSession, &handler, &active] {
@@ -757,7 +759,7 @@ void NetDKG::HandleDKGRound(ActiveDKGSessionHandler& handler)
     // Complain
     auto fComplainStart = [curSession, &handler, &active]() {
         if (auto qc = curSession->VerifyAndComplain(active.connman); qc) {
-            EnqueueOwn(handler.pendingComplaints, *qc);
+            EnqueueOwn(handler.pendingComplaints, curSession->ProTx(), *qc);
         }
     };
     auto fComplainWait = [this, curSession, &handler, &active] {
@@ -769,7 +771,7 @@ void NetDKG::HandleDKGRound(ActiveDKGSessionHandler& handler)
     // Justify
     auto fJustifyStart = [curSession, &handler]() {
         if (auto qj = curSession->VerifyAndJustify(); qj) {
-            EnqueueOwn(handler.pendingJustifications, *qj);
+            EnqueueOwn(handler.pendingJustifications, curSession->ProTx(), *qj);
         }
     };
     auto fJustifyWait = [this, curSession, &handler, &active] {
@@ -781,7 +783,7 @@ void NetDKG::HandleDKGRound(ActiveDKGSessionHandler& handler)
     // Commit
     auto fCommitStart = [curSession, &handler]() {
         if (auto qc = curSession->VerifyAndCommit(); qc) {
-            EnqueueOwn(handler.pendingPrematureCommitments, *qc);
+            EnqueueOwn(handler.pendingPrematureCommitments, curSession->ProTx(), *qc);
         }
     };
     auto fCommitWait = [this, curSession, &handler, &active] {

src/test/llmq_dkg_tests.cpp

@@ -3,6 +3,9 @@
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
 #include <llmq/dkgsession.h>
+#include <llmq/dkgsessionhandler.h>
+#include <protocol.h>
+#include <streams.h>
 #include <util/helpers.h>
 #include <util/std23.h>
 
@@ -23,4 +26,98 @@ BOOST_AUTO_TEST_CASE(llmq_dkgerror)
     BOOST_REQUIRE(GetSimulatedErrorRate(llmq::DKGError::type::_COUNT) == 0.0);
 }
 
+namespace {
+std::shared_ptr<CDataStream> MakeDKGMessage()
+{
+    return std::make_shared<CDataStream>(SER_NETWORK, PROTOCOL_VERSION);
+}
+
+uint256 MakeTestHash(uint8_t value)
+{
+    uint256 hash;
+    hash.begin()[0] = value;
+    return hash;
+}
+} // namespace
+
+BOOST_AUTO_TEST_CASE(pending_messages_own_messages_share_quota_path)
+{
+    using namespace llmq;
+
+    const uint256 own_protx = MakeTestHash(0xee);
+
+    // Own messages (from=-1) are enqueued under this node's proTxHash and
+    // charged like any other sender's.
+    CDKGPendingMessages pending{/*_maxMessagesPerProTx=*/2};
+    pending.PushPendingMessage(/*from=*/-1, own_protx, MakeDKGMessage(), MakeTestHash(1));
+    pending.PushPendingMessage(/*from=*/-1, own_protx, MakeDKGMessage(), MakeTestHash(2));
+    pending.PushPendingMessage(/*from=*/-1, own_protx, MakeDKGMessage(), MakeTestHash(3));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(2)));
+    BOOST_CHECK(!pending.HasSeen(MakeTestHash(3)));
+
+    BOOST_CHECK_EQUAL(pending.PopPendingMessages(3).size(), 2U);
+}
+
+BOOST_AUTO_TEST_CASE(pending_messages_quota_survives_reconnect)
+{
+    using namespace llmq;
+
+    const uint256 protx_a = MakeTestHash(0xa1);
+
+    CDKGPendingMessages pending{/*_maxMessagesPerProTx=*/2};
+    pending.PushPendingMessage(/*from=*/1, protx_a, MakeDKGMessage(), MakeTestHash(1));
+    pending.PushPendingMessage(/*from=*/1, protx_a, MakeDKGMessage(), MakeTestHash(2));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(1)));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(2)));
+
+    // Reconnecting mints a fresh NodeId but keeps the proTxHash, so the quota
+    // is already spent.
+    pending.PushPendingMessage(/*from=*/2, protx_a, MakeDKGMessage(), MakeTestHash(3));
+    pending.PushPendingMessage(/*from=*/3, protx_a, MakeDKGMessage(), MakeTestHash(4));
+    BOOST_CHECK(!pending.HasSeen(MakeTestHash(3)));
+    BOOST_CHECK(!pending.HasSeen(MakeTestHash(4)));
+
+    // Draining frees queue slots but does not refund the per-proTx quota.
+    BOOST_CHECK_EQUAL(pending.PopPendingMessages(5).size(), 2U);
+    pending.PushPendingMessage(/*from=*/4, protx_a, MakeDKGMessage(), MakeTestHash(5));
+    BOOST_CHECK(!pending.HasSeen(MakeTestHash(5)));
+
+    // A new round resets everything.
+    pending.Clear();
+    pending.PushPendingMessage(/*from=*/4, protx_a, MakeDKGMessage(), MakeTestHash(5));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(5)));
+}
+
+BOOST_AUTO_TEST_CASE(pending_messages_quota_is_per_protx)
+{
+    using namespace llmq;
+
+    const uint256 protx_a = MakeTestHash(0xa1);
+    const uint256 protx_b = MakeTestHash(0xb1);
+
+    CDKGPendingMessages pending{/*_maxMessagesPerProTx=*/2};
+    pending.PushPendingMessage(/*from=*/1, protx_a, MakeDKGMessage(), MakeTestHash(1));
+    pending.PushPendingMessage(/*from=*/1, protx_a, MakeDKGMessage(), MakeTestHash(2));
+
+    // Duplicates are rejected before charging the quota.
+    pending.PushPendingMessage(/*from=*/2, protx_b, MakeDKGMessage(), MakeTestHash(1));
+    pending.PushPendingMessage(/*from=*/2, protx_b, MakeDKGMessage(), MakeTestHash(2));
+
+    // One proTx's spent quota has no effect on another's.
+    pending.PushPendingMessage(/*from=*/2, protx_b, MakeDKGMessage(), MakeTestHash(3));
+    pending.PushPendingMessage(/*from=*/2, protx_b, MakeDKGMessage(), MakeTestHash(4));
+    pending.PushPendingMessage(/*from=*/2, protx_b, MakeDKGMessage(), MakeTestHash(5));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(3)));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(4)));
+    BOOST_CHECK(!pending.HasSeen(MakeTestHash(5)));
+
+    // A quota-dropped hash is not marked seen, so a sender with remaining
+    // budget can still deliver it.
+    const uint256 protx_c = MakeTestHash(0xc1);
+    pending.PushPendingMessage(/*from=*/3, protx_c, MakeDKGMessage(), MakeTestHash(5));
+    BOOST_CHECK(pending.HasSeen(MakeTestHash(5)));
+
+    BOOST_CHECK_EQUAL(pending.PopPendingMessages(6).size(), 5U);
+}
+
 BOOST_AUTO_TEST_SUITE_END()