llmq_utils_tests.cpp

// Copyright (c) 2025 The Dash Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <test/util/llmq_tests.h>
#include <test/util/setup_common.h>

#include <consensus/params.h>
#include <llmq/net_signing.h>
#include <llmq/params.h>
#include <llmq/signing_shares.h>
#include <llmq/utils.h>
#include <netaddress.h>
#include <random.h>
#include <streams.h>

#include <boost/test/unit_test.hpp>

#include <map>
#include <set>

using namespace llmq;
using namespace llmq::testutils;

BOOST_FIXTURE_TEST_SUITE(llmq_utils_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(trivially_passes) { BOOST_CHECK(true); }

static CSigSesAnn MakeSigSesAnn(uint32_t session_id, uint32_t nonce, Consensus::LLMQType llmq_type = Consensus::LLMQType::LLMQ_50_60)
{
    return CSigSesAnn{session_id, llmq_type, GetTestQuorumHash(1), GetTestQuorumHash(2), GetTestQuorumHash(nonce)};
}

static CSigShare MakeSigShare(uint32_t nonce, Consensus::LLMQType llmq_type = Consensus::LLMQType::LLMQ_50_60)
{
    CSigShare sig_share{llmq_type, GetTestQuorumHash(1), GetTestQuorumHash(2), GetTestQuorumHash(nonce), 1, CBLSLazySignature{}};
    sig_share.UpdateKey();
    return sig_share;
}

BOOST_AUTO_TEST_CASE(sig_ses_ann_respects_session_limit_but_allows_refresh)
{
    CSigSharesNodeState node_state;

    const CSigSesAnn ann1{MakeSigSesAnn(1, 1)};
    const CSigSesAnn ann2{MakeSigSesAnn(2, 2)};
    const CSigSesAnn ann3{MakeSigSesAnn(3, 3)};
    constexpr size_t max_sessions{2};

    BOOST_CHECK(node_state.CanCreateSessionFromAnn(ann1, max_sessions));
    node_state.GetOrCreateSessionFromAnn(ann1);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(), 1U);
    BOOST_CHECK_EQUAL(node_state.GetAnnouncementSessionCount(Consensus::LLMQType::LLMQ_50_60), 1U);

    BOOST_CHECK(node_state.CanCreateSessionFromAnn(ann2, max_sessions));
    node_state.GetOrCreateSessionFromAnn(ann2);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(), max_sessions);
    BOOST_CHECK_EQUAL(node_state.GetAnnouncementSessionCount(Consensus::LLMQType::LLMQ_50_60), max_sessions);

    BOOST_CHECK(!node_state.CanCreateSessionFromAnn(ann3, max_sessions));

    const CSigSesAnn ann1_refresh{4, Consensus::LLMQType::LLMQ_50_60, ann1.getQuorumHash(), ann1.getId(), ann1.getMsgHash()};
    BOOST_CHECK(node_state.CanCreateSessionFromAnn(ann1_refresh, max_sessions));
    node_state.GetOrCreateSessionFromAnn(ann1_refresh);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(), max_sessions);
    BOOST_CHECK_EQUAL(node_state.GetAnnouncementSessionCount(Consensus::LLMQType::LLMQ_50_60), max_sessions);
}

BOOST_AUTO_TEST_CASE(sig_ses_ann_limit_ignores_send_only_sessions)
{
    CSigSharesNodeState node_state;

    constexpr size_t max_sessions{1};
    const CSigShare sig_share{MakeSigShare(1)};
    const CSigSesAnn ann{MakeSigSesAnn(1, 2)};

    node_state.GetOrCreateSessionFromShare(sig_share);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(Consensus::LLMQType::LLMQ_50_60), 1U);
    BOOST_CHECK_EQUAL(node_state.GetAnnouncementSessionCount(Consensus::LLMQType::LLMQ_50_60), 0U);

    BOOST_CHECK(node_state.CanCreateSessionFromAnn(ann, max_sessions));
    node_state.GetOrCreateSessionFromAnn(ann);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(Consensus::LLMQType::LLMQ_50_60), 2U);
    BOOST_CHECK_EQUAL(node_state.GetAnnouncementSessionCount(Consensus::LLMQType::LLMQ_50_60), 1U);
}

BOOST_AUTO_TEST_CASE(sig_ses_ann_limit_is_per_llmq_type)
{
    CSigSharesNodeState node_state;

    constexpr size_t max_sessions{1};
    const CSigSesAnn ann1{MakeSigSesAnn(1, 1)};
    const CSigSesAnn ann2{MakeSigSesAnn(2, 2)};
    const CSigSesAnn other_type_ann{MakeSigSesAnn(3, 3, Consensus::LLMQType::LLMQ_400_60)};

    BOOST_CHECK(node_state.CanCreateSessionFromAnn(ann1, max_sessions));
    node_state.GetOrCreateSessionFromAnn(ann1);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(), 1U);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(Consensus::LLMQType::LLMQ_50_60), 1U);

    BOOST_CHECK(!node_state.CanCreateSessionFromAnn(ann2, max_sessions));
    BOOST_CHECK(node_state.CanCreateSessionFromAnn(other_type_ann, max_sessions));
    node_state.GetOrCreateSessionFromAnn(other_type_ann);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(), 2U);
    BOOST_CHECK_EQUAL(node_state.GetSessionCount(Consensus::LLMQType::LLMQ_400_60), 1U);
}

BOOST_AUTO_TEST_CASE(sig_share_map_size_tracks_mutations)
{
    SigShareMap<CSigShare> sig_share_map;
    const CSigShare sig_share1{MakeSigShare(1)};
    const CSigShare sig_share2{MakeSigShare(2)};

    BOOST_CHECK(sig_share_map.Add(sig_share1.GetKey(), sig_share1));
    BOOST_CHECK(!sig_share_map.Add(sig_share1.GetKey(), sig_share1));
    BOOST_CHECK(sig_share_map.Add(sig_share2.GetKey(), sig_share2));
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 2U);

    sig_share_map.Erase(sig_share1.GetKey());
    sig_share_map.Erase(sig_share1.GetKey());
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 1U);

    sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash());
    sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash());
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 0U);
    BOOST_CHECK(sig_share_map.Empty());

    BOOST_CHECK(sig_share_map.Add(sig_share1.GetKey(), sig_share1));
    BOOST_CHECK(sig_share_map.Add(sig_share2.GetKey(), sig_share2));
    sig_share_map.EraseIf([&](const SigShareKey& k, const CSigShare&) { return k == sig_share1.GetKey(); });
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 1U);

    sig_share_map.Clear();
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 0U);
}

BOOST_AUTO_TEST_CASE(sig_share_map_bucket_erase_updates_size)
{
    SigShareMap<CSigShare> sig_share_map;
    const auto sign_hash = MakeSigShare(1).GetSignHash();

    for (uint16_t member = 0; member < 5; ++member) {
        CSigShare s{Consensus::LLMQType::LLMQ_50_60, GetTestQuorumHash(1), GetTestQuorumHash(2), GetTestQuorumHash(1),
                    member, CBLSLazySignature{}};
        s.UpdateKey();
        BOOST_CHECK_EQUAL(s.GetSignHash(), sign_hash);
        BOOST_CHECK(sig_share_map.Add(s.GetKey(), s));
    }
    BOOST_CHECK_EQUAL(sig_share_map.Size(), 5U);

    sig_share_map.EraseAllForSignHash(sign_hash);
    BOOST_CHECK(sig_share_map.Empty());
}

BOOST_AUTO_TEST_CASE(pending_sig_shares_session_removal_updates_count)
{
    CSigSharesNodeState node_state;
    const CSigShare sig_share1{MakeSigShare(1)};
    const CSigShare sig_share2{MakeSigShare(2)};

    BOOST_CHECK(node_state.pendingIncomingSigShares.Add(sig_share1.GetKey(), sig_share1));
    BOOST_CHECK(node_state.pendingIncomingSigShares.Add(sig_share2.GetKey(), sig_share2));
    BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 2U);

    node_state.RemoveSession(sig_share1.GetSignHash());
    BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 1U);
    BOOST_CHECK(!node_state.pendingIncomingSigShares.Has(sig_share1.GetKey()));
    BOOST_CHECK(node_state.pendingIncomingSigShares.Has(sig_share2.GetKey()));

    // Removing the same session twice, or a session with no pending shares, is a no-op.
    node_state.RemoveSession(sig_share1.GetSignHash());
    node_state.RemoveSession(MakeSigShare(3).GetSignHash());
    BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 1U);
}

BOOST_AUTO_TEST_CASE(batched_sig_shares_rejects_oversized_inner_vector)
{
    CDataStream stream{SER_NETWORK, PROTOCOL_VERSION};
    stream << VARINT(uint32_t{1});
    WriteCompactSize(stream, MAX_MSGS_TOTAL_BATCHED_SIGS + 1);

    CBatchedSigShares batched_sig_shares;
    BOOST_CHECK_THROW(stream >> batched_sig_shares, std::ios_base::failure);
    BOOST_CHECK(batched_sig_shares.sigShares.empty());
}

BOOST_AUTO_TEST_CASE(batched_sig_shares_accepts_max_inner_vector)
{
    CBatchedSigShares batched;
    batched.sessionId = 1;
    batched.sigShares.resize(MAX_MSGS_TOTAL_BATCHED_SIGS); // exactly the cap must be accepted

    CDataStream stream{SER_NETWORK, PROTOCOL_VERSION};
    stream << batched;

    CBatchedSigShares roundtripped;
    BOOST_CHECK_NO_THROW(stream >> roundtripped);
    BOOST_CHECK_EQUAL(roundtripped.sigShares.size(), MAX_MSGS_TOTAL_BATCHED_SIGS);
}

static CBatchedSigShares MakeBatch(uint32_t session_id, size_t share_count)
{
    CBatchedSigShares batch;
    batch.sessionId = session_id;
    batch.sigShares.resize(share_count); // default pairs are enough to exercise the count invariant
    return batch;
}

// Exercise the production QBSIGSHARES decoder directly: the outer batch count is
// bounded regardless of the inner contents. Every batch here is empty, so the
// aggregate running-total guard never fires; only the outer-count guard can
// reject the stream, which keeps this case a genuine test of that guard rather
// than an end-of-stream artifact.
BOOST_AUTO_TEST_CASE(qbsigshares_rejects_oversized_batch_count)
{
    std::vector<CBatchedSigShares> msgs(MAX_MSGS_TOTAL_BATCHED_SIGS + 1); // count over cap, 0 inner shares each

    CDataStream stream{SER_NETWORK, PROTOCOL_VERSION};
    stream << msgs;

    BOOST_CHECK_THROW(UnserializeBatchedSigShares(stream), std::ios_base::failure);
}

// The regression this targets: each batch is individually within the per-batch
// LIMITED_VECTOR cap, but their running aggregate exceeds MAX_MSGS_TOTAL_BATCHED_SIGS,
// so the decoder must abort mid-stream rather than accept the cross product.
BOOST_AUTO_TEST_CASE(qbsigshares_rejects_oversized_aggregate_total)
{
    std::vector<CBatchedSigShares> msgs;
    msgs.push_back(MakeBatch(1, 300));
    msgs.push_back(MakeBatch(2, 200)); // 300 + 200 = 500 > 400, though each batch is <= the cap

    CDataStream stream{SER_NETWORK, PROTOCOL_VERSION};
    stream << msgs;

    BOOST_CHECK_THROW(UnserializeBatchedSigShares(stream), std::ios_base::failure);
}

// Batches whose aggregate is exactly at the cap must be accepted and decoded intact.
BOOST_AUTO_TEST_CASE(qbsigshares_accepts_aggregate_at_cap)
{
    std::vector<CBatchedSigShares> msgs;
    msgs.push_back(MakeBatch(1, 200));
    msgs.push_back(MakeBatch(2, MAX_MSGS_TOTAL_BATCHED_SIGS - 200)); // aggregate == cap

    CDataStream stream{SER_NETWORK, PROTOCOL_VERSION};
    stream << msgs;

    std::vector<CBatchedSigShares> decoded;
    BOOST_CHECK_NO_THROW(decoded = UnserializeBatchedSigShares(stream));
    BOOST_CHECK_EQUAL(decoded.size(), 2U);
    BOOST_CHECK_EQUAL(decoded[0].sigShares.size(), 200U);
    BOOST_CHECK_EQUAL(decoded[1].sigShares.size(), MAX_MSGS_TOTAL_BATCHED_SIGS - 200);
}

BOOST_AUTO_TEST_CASE(deterministic_outbound_connection_test)
{
    // Test deterministic behavior
    // DeterministicOutboundConnection returns one of the two input hashes based on a deterministic calculation
    uint256 proTxHash1 = GetTestQuorumHash(1);
    uint256 proTxHash2 = GetTestQuorumHash(2);

    // Same inputs should produce same output
    uint256 conn1a = llmq::utils::DeterministicOutboundConnection(proTxHash1, proTxHash2);
    uint256 conn1b = llmq::utils::DeterministicOutboundConnection(proTxHash1, proTxHash2);
    BOOST_CHECK(conn1a == conn1b);
    // Result should be one of the input hashes
    BOOST_CHECK(conn1a == proTxHash1 || conn1a == proTxHash2);

    // Swapped inputs should produce the same result (commutative)
    // The function deterministically selects which node initiates the connection
    uint256 conn2 = llmq::utils::DeterministicOutboundConnection(proTxHash2, proTxHash1);
    BOOST_CHECK(conn1a == conn2);

    // The result should consistently be the same node regardless of order
    BOOST_CHECK(llmq::utils::DeterministicOutboundConnection(proTxHash1, proTxHash2) ==
                llmq::utils::DeterministicOutboundConnection(proTxHash2, proTxHash1));
}

BOOST_AUTO_TEST_CASE(deterministic_outbound_connection_edge_cases_test)
{
    // Test with null hashes
    uint256 nullHash;
    uint256 validHash = GetTestQuorumHash(1);

    // DeterministicOutboundConnection returns one of the input hashes
    uint256 conn1 = llmq::utils::DeterministicOutboundConnection(nullHash, validHash);
    uint256 conn2 = llmq::utils::DeterministicOutboundConnection(validHash, nullHash);
    uint256 conn3 = llmq::utils::DeterministicOutboundConnection(nullHash, nullHash);

    // With null and valid hash, should return one of them
    BOOST_CHECK(conn1 == nullHash || conn1 == validHash);
    BOOST_CHECK(conn2 == nullHash || conn2 == validHash);
    // Since the function is order-independent, conn1 and conn2 should be the same
    BOOST_CHECK(conn1 == conn2);

    // With two null hashes, should return null
    BOOST_CHECK(conn3 == nullHash);

    // Test with same source and destination
    uint256 sameHash = GetTestQuorumHash(42);
    uint256 connSame = llmq::utils::DeterministicOutboundConnection(sameHash, sameHash);
    // Should return the same hash
    BOOST_CHECK(connSame == sameHash);
    BOOST_CHECK(!connSame.IsNull());
}

// Note: CalcDeterministicWatchConnections requires CBlockIndex which is complex to mock
// Testing is deferred to functional tests

// Note: InitQuorumsCache requires specific cache types with LLMQ consensus parameters
// Testing is deferred to integration tests

BOOST_AUTO_TEST_CASE(deterministic_connection_symmetry_test)
{
    // Test interesting properties of DeterministicOutboundConnection
    uint256 proTxHash1 = GetTestQuorumHash(1);
    uint256 proTxHash2 = GetTestQuorumHash(2);
    uint256 proTxHash3 = GetTestQuorumHash(3);

    // Create a "network" of connections
    // DeterministicOutboundConnection is symmetric - order doesn't matter
    uint256 conn12 = llmq::utils::DeterministicOutboundConnection(proTxHash1, proTxHash2);
    uint256 conn21 = llmq::utils::DeterministicOutboundConnection(proTxHash2, proTxHash1);
    uint256 conn13 = llmq::utils::DeterministicOutboundConnection(proTxHash1, proTxHash3);
    uint256 conn31 = llmq::utils::DeterministicOutboundConnection(proTxHash3, proTxHash1);
    uint256 conn23 = llmq::utils::DeterministicOutboundConnection(proTxHash2, proTxHash3);
    uint256 conn32 = llmq::utils::DeterministicOutboundConnection(proTxHash3, proTxHash2);

    // Verify symmetry - swapped inputs produce same output
    BOOST_CHECK(conn12 == conn21);
    BOOST_CHECK(conn13 == conn31);
    BOOST_CHECK(conn23 == conn32);

    // Each connection returns one of the two nodes
    BOOST_CHECK(conn12 == proTxHash1 || conn12 == proTxHash2);
    BOOST_CHECK(conn13 == proTxHash1 || conn13 == proTxHash3);
    BOOST_CHECK(conn23 == proTxHash2 || conn23 == proTxHash3);

    // The function deterministically picks which node initiates the connection
    // Verify we get consistent results for each pair
    std::set<uint256> uniqueResults;
    uniqueResults.insert(conn12);
    uniqueResults.insert(conn13);
    uniqueResults.insert(conn23);
    // Each pair should produce one of its members, but pairs may have overlapping results
    BOOST_CHECK(uniqueResults.size() >= 2 && uniqueResults.size() <= 3);
}

BOOST_AUTO_TEST_SUITE_END()