bls_tests.cpp

// Copyright (c) 2019-2024 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 <bls/bls.h>
#include <bls/bls_batchverifier.h>
#include <clientversion.h>
#include <random.h>
#include <streams.h>
#include <util/irange.h>
#include <util/strencodings.h>

#include <boost/test/unit_test.hpp>

BOOST_AUTO_TEST_SUITE(bls_tests)

void FuncSign(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    CBLSSecretKey sk1, sk2;
    sk1.MakeNewKey();
    sk2.MakeNewKey();

    uint256 msgHash1 = uint256::ONE;
    uint256 msgHash2 = uint256::TWO;

    auto sig1 = sk1.Sign(msgHash1, legacy_scheme);
    auto sig2 = sk2.Sign(msgHash1, legacy_scheme);
    BOOST_CHECK(sig1.VerifyInsecure(sk1.GetPublicKey(), msgHash1));
    BOOST_CHECK(!sig1.VerifyInsecure(sk1.GetPublicKey(), msgHash2));
    BOOST_CHECK(!sig2.VerifyInsecure(sk1.GetPublicKey(), msgHash1));
    BOOST_CHECK(!sig2.VerifyInsecure(sk2.GetPublicKey(), msgHash2));
    BOOST_CHECK(sig2.VerifyInsecure(sk2.GetPublicKey(), msgHash1));

    return;
}

void FuncSerialize(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    CBLSSecretKey sk;
    CDataStream ds2(SER_DISK, CLIENT_VERSION), ds3(SER_DISK, CLIENT_VERSION);
    uint256 msgHash = uint256::ONE;

    sk.MakeNewKey();
    CBLSSignature sig1 = sk.Sign(msgHash, legacy_scheme);
    ds2 << sig1;
    ds3 << CBLSSignatureVersionWrapper(const_cast<CBLSSignature&>(sig1), !legacy_scheme);

    CBLSSignature sig2;
    ds2 >> sig2;
    BOOST_CHECK(sig1 == sig2);

    CBLSSignature sig3;
    ds3 >> CBLSSignatureVersionWrapper(const_cast<CBLSSignature&>(sig3), !legacy_scheme);
    BOOST_CHECK(sig1 == sig3);

    return;
}

void FuncSetHexStr(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    // Note: 2nd bool argument for SetHexStr for bls::PrivateKey has a meaning modOrder, not is-legacy
    CBLSSecretKey sk;
    std::string strValidSecret = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";
    // Note: invalid string passed to SetHexStr() should cause it to fail and reset key internal data
    BOOST_CHECK(sk.SetHexStr(strValidSecret, false));
    BOOST_CHECK(!sk.SetHexStr("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1g", false)); // non-hex
    BOOST_CHECK(!sk.IsValid());
    BOOST_CHECK(sk == CBLSSecretKey());
    // Try few more invalid strings
    BOOST_CHECK(sk.SetHexStr(strValidSecret, false));
    BOOST_CHECK(!sk.SetHexStr("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e", false)); // hex but too short
    BOOST_CHECK(!sk.IsValid());
    BOOST_CHECK(sk.SetHexStr(strValidSecret, false));
    BOOST_CHECK(
        !sk.SetHexStr("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20", false)); // hex but too long
    BOOST_CHECK(!sk.IsValid());

    return;
}

void FuncKeyAgg(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    CBLSSecretKey sk1, sk2;
    sk1.MakeNewKey();
    sk2.MakeNewKey();

    CBLSPublicKey ag_pk = sk1.GetPublicKey();
    ag_pk.AggregateInsecure(sk2.GetPublicKey());

    CBLSSecretKey ag_sk = sk1;
    ag_sk.AggregateInsecure(sk2);

    BOOST_CHECK(ag_pk == ag_sk.GetPublicKey());

    uint256 msgHash1 = uint256::ONE;
    uint256 msgHash2 = uint256::TWO;

    auto sig = ag_sk.Sign(msgHash1, legacy_scheme);
    BOOST_CHECK(sig.VerifyInsecure(ag_pk, msgHash1));
    BOOST_CHECK(!sig.VerifyInsecure(ag_pk, msgHash2));
}

void FuncKeyAggVec(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    std::vector<CBLSSecretKey> vec_sk;
    std::vector<CBLSPublicKey> vec_pk;

    {
        auto ret = CBLSSecretKey::AggregateInsecure(vec_sk);
        BOOST_CHECK(ret == CBLSSecretKey());
    }
    {
        auto ret = CBLSPublicKey::AggregateInsecure(vec_pk);
        BOOST_CHECK(ret == CBLSPublicKey());
    }

    // In practice, we only aggregate 400 key shares at any given time, something substantially larger than that should
    // be good. Plus this is very very fast, so who cares!
    int key_count = 10000;
    vec_sk.reserve(key_count);
    vec_pk.reserve(key_count);
    CBLSSecretKey sk;
    for (int i = 0; i < key_count; i++) {
        sk.MakeNewKey();
        vec_sk.push_back(sk);
        vec_pk.push_back(sk.GetPublicKey());
    }

    CBLSSecretKey ag_sk = CBLSSecretKey::AggregateInsecure(vec_sk);
    CBLSPublicKey ag_pk = CBLSPublicKey::AggregateInsecure(vec_pk);

    BOOST_CHECK(ag_sk.IsValid());
    BOOST_CHECK(ag_pk.IsValid());

    uint256 msgHash1 = uint256::ONE;
    uint256 msgHash2 = uint256::TWO;

    auto sig = ag_sk.Sign(msgHash1, legacy_scheme);
    BOOST_CHECK(sig.VerifyInsecure(ag_pk, msgHash1));
    BOOST_CHECK(!sig.VerifyInsecure(ag_pk, msgHash2));
}


void FuncSigAggSub(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    int count = 20;
    std::vector<CBLSPublicKey> vec_pks;
    std::vector<uint256> vec_hashes;
    std::vector<CBLSSignature> vec_sigs;
    vec_pks.reserve(count);
    vec_hashes.reserve(count);
    vec_sigs.reserve(count);

    CBLSSignature sig;
    auto sk = CBLSSecretKey();
    uint256 hash;
    for (int i = 0; i < count; i++) {
        sk.MakeNewKey();
        vec_pks.push_back(sk.GetPublicKey());
        hash = GetRandHash();
        vec_hashes.push_back(hash);
        CBLSSignature sig_i = sk.Sign(hash, legacy_scheme);
        vec_sigs.push_back(sig_i);
        if (i == 0) {
            // first sig is assigned directly
            sig = sig_i;
        } else {
            // all other sigs are aggregated into the previously computed/stored sig
            sig.AggregateInsecure(sig_i);
        }
        BOOST_CHECK(sig.VerifyInsecureAggregated(vec_pks, vec_hashes));
    }
    // Create an aggregated signature from the vector of individual signatures
    auto vecSig = CBLSSignature::AggregateInsecure(vec_sigs);
    BOOST_CHECK(vecSig.VerifyInsecureAggregated(vec_pks, vec_hashes));
    // Check that these two signatures are equal
    BOOST_CHECK(sig == vecSig);

    // Test that the sig continues to be valid when subtracting sigs via `SubInsecure`

    for (int i = 0; i < count - 1; i++) {
        auto top_sig = vec_sigs.back();
        vec_pks.pop_back();
        vec_hashes.pop_back();
        BOOST_CHECK(!sig.VerifyInsecureAggregated(vec_pks, vec_hashes));
        sig.SubInsecure(top_sig);
        BOOST_CHECK(sig.VerifyInsecureAggregated(vec_pks, vec_hashes));
        vec_sigs.pop_back();
    }
    // Check that the final left-over sig validates
    BOOST_CHECK_EQUAL(vec_sigs.size(), 1);
    BOOST_CHECK_EQUAL(vec_pks.size(), 1);
    BOOST_CHECK_EQUAL(vec_hashes.size(), 1);
    BOOST_CHECK(vec_sigs[0].VerifyInsecure(vec_pks[0], vec_hashes[0]));
}


void FuncSigAggSecure(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    int count = 10;

    uint256 hash = GetRandHash();

    std::vector<CBLSSignature> vec_sigs;
    std::vector<CBLSPublicKey> vec_pks;

    CBLSSecretKey sk;
    for (int i = 0; i < count; i++) {
        sk.MakeNewKey();
        vec_pks.push_back(sk.GetPublicKey());
        vec_sigs.push_back(sk.Sign(hash, legacy_scheme));
    }

    auto sec_agg_sig = CBLSSignature::AggregateSecure(vec_sigs, vec_pks, hash);
    BOOST_CHECK(sec_agg_sig.IsValid());
    BOOST_CHECK(sec_agg_sig.VerifySecureAggregated(vec_pks, hash));
}

void FuncDHExchange(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    CBLSSecretKey sk1, sk2;
    sk1.MakeNewKey();
    sk2.MakeNewKey();

    CBLSPublicKey pk1, pk2;
    pk1 = sk1.GetPublicKey();
    pk2 = sk2.GetPublicKey();

    // Perform diffie-helman exchange
    CBLSPublicKey pke1, pke2;
    pke1.DHKeyExchange(sk1, pk2);
    pke2.DHKeyExchange(sk2, pk1);

    BOOST_CHECK(pke1.IsValid());
    BOOST_CHECK(pke2.IsValid());
    BOOST_CHECK(pke1 == pke2);
}

struct Message
{
    uint32_t sourceId;
    uint32_t msgId;
    uint256 msgHash;
    CBLSSecretKey sk;
    CBLSPublicKey pk;
    CBLSSignature sig;
    bool valid;
};

static void AddMessage(std::vector<Message>& vec, uint32_t sourceId, uint32_t msgId, uint8_t msgHash, bool valid)
{
    bool legacy_scheme = bls::bls_legacy_scheme.load();
    Message m;
    m.sourceId = sourceId;
    m.msgId = msgId;
    m.msgHash = uint256(msgHash);
    m.sk.MakeNewKey();
    m.pk = m.sk.GetPublicKey();
    m.sig = m.sk.Sign(m.msgHash, legacy_scheme);
    m.valid = valid;

    if (!valid) {
        CBLSSecretKey tmp;
        tmp.MakeNewKey();
        m.sig = tmp.Sign(m.msgHash, legacy_scheme);
    }

    vec.emplace_back(m);
}

static void Verify(std::vector<Message>& vec, bool secureVerification, bool perMessageFallback)
{
    CBLSBatchVerifier<uint32_t, uint32_t> batchVerifier(secureVerification, perMessageFallback);

    std::set<uint32_t> expectedBadMessages;
    std::set<uint32_t> expectedBadSources;
    for (auto& m : vec) {
        if (!m.valid) {
            expectedBadMessages.emplace(m.msgId);
            expectedBadSources.emplace(m.sourceId);
        }

        batchVerifier.PushMessage(m.sourceId, m.msgId, m.msgHash, m.sig, m.pk);
    }

    batchVerifier.Verify();

    BOOST_CHECK(batchVerifier.badSources == expectedBadSources);

    if (perMessageFallback) {
        BOOST_CHECK(batchVerifier.badMessages == expectedBadMessages);
    } else {
        BOOST_CHECK(batchVerifier.badMessages.empty());
    }
}

static void Verify(std::vector<Message>& vec)
{
    Verify(vec, false, false);
    Verify(vec, true, false);
    Verify(vec, false, true);
    Verify(vec, true, true);
}

void FuncBatchVerifier(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    std::vector<Message> msgs;

    // distinct messages from distinct sources
    AddMessage(msgs, 1, 1, 1, true);
    AddMessage(msgs, 2, 2, 2, true);
    AddMessage(msgs, 3, 3, 3, true);
    Verify(msgs);

    // distinct messages from same source
    AddMessage(msgs, 4, 4, 4, true);
    AddMessage(msgs, 4, 5, 5, true);
    AddMessage(msgs, 4, 6, 6, true);
    Verify(msgs);

    // invalid sig
    AddMessage(msgs, 7, 7, 7, false);
    Verify(msgs);

    // same message as before, but from another source and with valid sig
    AddMessage(msgs, 8, 8, 7, true);
    Verify(msgs);

    // same message as before, but from another source and signed with another key
    AddMessage(msgs, 9, 9, 7, true);
    Verify(msgs);

    msgs.clear();
    // same message, signed by multiple keys
    AddMessage(msgs, 1, 1, 1, true);
    AddMessage(msgs, 1, 2, 1, true);
    AddMessage(msgs, 1, 3, 1, true);
    AddMessage(msgs, 2, 4, 1, true);
    AddMessage(msgs, 2, 5, 1, true);
    AddMessage(msgs, 2, 6, 1, true);
    Verify(msgs);

    // last message invalid from one source
    AddMessage(msgs, 1, 7, 1, false);
    Verify(msgs);
}

void FuncThresholdSignature(const bool legacy_scheme)
{
    bls::bls_legacy_scheme.store(legacy_scheme);

    [[maybe_unused]] const uint256 hash = GetRandHash();

    constexpr size_t m_size = 20;
    constexpr size_t m_threshold = 15;

    std::vector<CBLSSecretKey> v_threshold_sks;
    std::vector<CBLSPublicKey> v_threshold_pks;
    for ([[maybe_unused]] const auto i : irange::range(m_threshold)) {
        CBLSSecretKey sk;
        sk.MakeNewKey();
        v_threshold_sks.push_back(sk);
        v_threshold_pks.emplace_back(sk.GetPublicKey());
    }

    CBLSSecretKey thr_sk = v_threshold_sks[0];
    CBLSPublicKey thr_pk = v_threshold_pks[0];
    CBLSSignature thr_sig = thr_sk.Sign(hash, legacy_scheme);

    std::vector<CBLSId> v_size_ids;
    std::vector<CBLSSecretKey> v_size_sk_shares;
    std::vector<CBLSPublicKey> v_size_pk_shares;
    for ([[maybe_unused]] const auto m_shares : irange::range(m_size)) {
        v_size_ids.emplace_back(GetRandHash());
        CBLSSecretKey sk;
        BOOST_CHECK(sk.SecretKeyShare(v_threshold_sks, v_size_ids[m_shares]));
        v_size_sk_shares.push_back(sk);
        CBLSPublicKey pk;
        BOOST_CHECK(pk.PublicKeyShare(v_threshold_pks, v_size_ids[m_shares]));
        v_size_pk_shares.push_back(pk);

        std::vector<CBLSSignature> v_share_sigs;
        std::vector<CBLSId> v_share_ids;
        for ([[maybe_unused]] const auto j : irange::range(m_shares)) {
            v_share_sigs.emplace_back(v_size_sk_shares[j].Sign(hash, legacy_scheme));
            BOOST_CHECK(v_share_sigs.back().VerifyInsecure(v_size_pk_shares[j], hash));
            v_share_ids.push_back(v_size_ids[j]);
        }

        CBLSSignature rec_share_sig;
        BOOST_CHECK_EQUAL(rec_share_sig.Recover(v_share_sigs, v_share_ids), m_shares >= 2);
        BOOST_CHECK_EQUAL(rec_share_sig.IsValid(), m_shares >= 2);
        BOOST_CHECK_EQUAL(rec_share_sig == thr_sig, m_shares >= m_threshold);
        BOOST_CHECK_EQUAL(rec_share_sig.VerifyInsecure(thr_pk, hash), m_shares >= m_threshold);
    }
}

BOOST_AUTO_TEST_CASE(bls_sethexstr_tests)
{
    FuncSetHexStr(true);
    FuncSetHexStr(false);
}

BOOST_AUTO_TEST_CASE(bls_serialize_tests)
{
    FuncSerialize(true);
    FuncSerialize(false);
}

BOOST_AUTO_TEST_CASE(bls_sig_tests)
{
    FuncSign(true);
    FuncSign(false);
}

BOOST_AUTO_TEST_CASE(bls_key_agg_tests)
{
    FuncKeyAgg(true);
    FuncKeyAgg(false);
}

BOOST_AUTO_TEST_CASE(bls_key_agg_vec_tests)
{
    FuncKeyAggVec(true);
    FuncKeyAggVec(false);
}

BOOST_AUTO_TEST_CASE(bls_sig_agg_sub_tests)
{
    FuncSigAggSub(true);
    FuncSigAggSub(false);
}

BOOST_AUTO_TEST_CASE(bls_sig_agg_secure_tests)
{
    FuncSigAggSecure(true);
    FuncSigAggSecure(false);
}

BOOST_AUTO_TEST_CASE(bls_dh_exchange_tests)
{
    FuncDHExchange(true);
    FuncDHExchange(false);
}

BOOST_AUTO_TEST_CASE(batch_verifier_tests)
{
    FuncBatchVerifier(true);
    FuncBatchVerifier(false);
}

BOOST_AUTO_TEST_CASE(bls_threshold_signature_tests)
{
    FuncThresholdSignature(true);
    FuncThresholdSignature(false);
}

// A dummy BLS object that satisfies the minimal interface expected by CBLSLazyWrapper.
class DummyBLS
{
public:
    // Define a fixed serialization size (for testing purposes).
    static const size_t SerSize = 4;
    std::array<uint8_t, SerSize> data{};

    DummyBLS() { data.fill(0); }

    // A dummy validity check: valid if any byte is non-zero.
    bool IsValid() const
    {
        return std::any_of(data.begin(), data.end(), [](uint8_t c) { return c != 0; });
    }

    // Convert to bytes; ignore the legacy flag for simplicity.
    std::array<uint8_t, SerSize> ToBytes(bool /*legacy*/) const { return data; }

    // Set from bytes; again, ignore the legacy flag.
    void SetBytes(const std::array<uint8_t, SerSize>& bytes, bool /*legacy*/) { data = bytes; }

    // A dummy malleability check: simply compares the stored data to the given bytes.
    bool CheckMalleable(const std::array<uint8_t, SerSize>& bytes, bool /*legacy*/) const { return data == bytes; }

    // Reset the object to an "empty" state.
    void Reset() { data.fill(0); }

    // Produce a string representation.
    std::string ToString(bool /*legacy*/) const { return HexStr(data); }

    // Equality operator.
    bool operator==(const DummyBLS& other) const { return data == other.data; }
};

// Define a type alias for our lazy wrapper instantiated with DummyBLS.
using LazyDummyBLS = CBLSLazyWrapper<DummyBLS>;

// Test 1: Two default (unset) wrappers should compare equal.
BOOST_AUTO_TEST_CASE(test_default_equality)
{
    LazyDummyBLS lazy1;
    LazyDummyBLS lazy2;
    // Neither instance has been set, so they represent the default/null object.
    BOOST_CHECK(lazy1 == lazy2);
}

// Test 2: A default wrapper and one initialized with a nonzero DummyBLS should compare unequal.
BOOST_AUTO_TEST_CASE(test_non_default_vs_default)
{
    LazyDummyBLS lazy_default;
    LazyDummyBLS lazy_set;
    DummyBLS obj;
    obj.data = {1, 0, 0, 0}; // nonzero data makes the object valid
    lazy_set.Set(obj, false);
    BOOST_CHECK(!(lazy_default == lazy_set));
    BOOST_CHECK(lazy_default != lazy_set);
}

// Test 2: A default wrapper and one initialized with a nonzero DummyBLS should compare unequal.
BOOST_AUTO_TEST_CASE(test_non_default_vs_different)
{
    LazyDummyBLS lazy_a;
    LazyDummyBLS lazy_b;
    DummyBLS obj;
    obj.data = {1, 2, 3, 4}; // nonzero data makes the object valid
    lazy_a.Set(obj, false);
    obj.data = {2, 2, 3, 4}; // nonzero data makes the object valid
    lazy_b.Set(obj, false);
    BOOST_CHECK(lazy_a != lazy_b);
}

// Test 3: Two wrappers set with the same underlying DummyBLS value compare equal.
BOOST_AUTO_TEST_CASE(test_equality_same_value)
{
    LazyDummyBLS lazy1;
    LazyDummyBLS lazy2;
    BOOST_CHECK(lazy1 == lazy2);
    DummyBLS obj;
    obj.data = {5, 6, 7, 8};
    lazy1.Set(obj, false);
    BOOST_CHECK(lazy1 != lazy2);
    lazy2.Set(obj, false);
    BOOST_CHECK(lazy1 == lazy2);
}

// Test 4: Serialization and unserialization preserve the wrapped value.
BOOST_AUTO_TEST_CASE(test_serialization_unserialization)
{
    LazyDummyBLS lazy1;
    DummyBLS obj;
    obj.data = {9, 10, 11, 12};
    // Set with a specific legacy flag (true in this case)
    lazy1.Set(obj, true);

    // Serialize the lazy object into a data stream.
    CDataStream ds(SER_DISK, CLIENT_VERSION);
    lazy1.Serialize(ds, true);

    // Create a new instance and unserialize the data into it.
    LazyDummyBLS lazy2;
    lazy2.Unserialize(ds, true);
    BOOST_CHECK(lazy1 == lazy2);
    BOOST_CHECK(lazy2.Get() == obj);
}

// Test 5: Two wrappers wrapping the same object should have the same hash.
BOOST_AUTO_TEST_CASE(test_get_hash_consistency)
{
    LazyDummyBLS lazy1;
    LazyDummyBLS lazy2;
    DummyBLS obj;
    obj.data = {13, 14, 15, 16};
    lazy1.Set(obj, false);
    lazy2.Set(obj, false);
    uint256 hash1 = lazy1.GetHash();
    uint256 hash2 = lazy2.GetHash();
    BOOST_CHECK(hash1 == hash2);
}

BOOST_AUTO_TEST_SUITE_END()