// Copyright (c) 2016-2024 The Hush developers
// Distributed under the GPLv3 software license, see the accompanying
// file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html

#include <gtest/gtest.h>

#include <deque>
#include <functional>

#include "chain.h"
#include "coins.h"
#include "main.h"
#include "primitives/block.h"
#include "primitives/transaction.h"
#include "random.h"
#include "wallet/wallet.h"
#include "zcash/Address.hpp"
#include "zcash/Note.hpp"
#include "zcash/NoteEncryption.hpp"
#include "zcash/zip32.h"

int32_t hush_dpowconfs(int32_t height, int32_t numconfs);
int tx_height(const uint256 &hash);

using namespace libzcash;

namespace {

typedef std::function<int(const uint256&)> TxHeightFn;

// CWallet::GetFilteredNotes as of 69bb035e7, verbatim except that tx_height() is passed in and its two
// LogPrintf lines are dropped. This is the reference the optimized walk must reproduce entry for
// entry and in the same order (z_sendmany sorts the result with an unstable sort, so the order is
// observable).
void LegacyFilteredNotes(
    CWallet& wallet,
    std::vector<SaplingNoteEntry>& saplingEntries,
    std::set<PaymentAddress>& filterAddresses,
    int minDepth,
    int maxDepth,
    bool ignoreSpent,
    bool requireSpendingKey,
    bool ignoreLocked,
    const TxHeightFn& tx_height)
{
    LOCK2(cs_main, wallet.cs_wallet);

    for (auto & p : wallet.mapWallet) {
        CWalletTx wtx = p.second;

        // Filter the transactions before checking for notes
        if (!CheckFinalTx(wtx) || wtx.GetBlocksToMaturity() > 0)
            continue;

        if (minDepth > 1) {
            int nHeight    = tx_height(wtx.GetHash());
            if ( nHeight == 0 ) {
                continue;
            }
            int nDepth     = wtx.GetDepthInMainChain();
            int dpowconfs  = hush_dpowconfs(nHeight,nDepth);
            if ( dpowconfs < minDepth || dpowconfs > maxDepth) {
                continue;
            }
        } else {
            if ( wtx.GetDepthInMainChain() < minDepth ||
                wtx.GetDepthInMainChain() > maxDepth) {
                continue;
            }
        }

        for (auto & pair : wtx.mapSaplingNoteData) {
            SaplingOutPoint op = pair.first;
            SaplingNoteData nd = pair.second;

            auto maybe_pt = SaplingNotePlaintext::decrypt(
                wtx.vShieldedOutput[op.n].encCiphertext,
                nd.ivk,
                wtx.vShieldedOutput[op.n].ephemeralKey,
                wtx.vShieldedOutput[op.n].cm);
            assert(static_cast<bool>(maybe_pt));
            auto notePt = maybe_pt.get();

            auto maybe_pa = nd.ivk.address(notePt.d);
            assert(static_cast<bool>(maybe_pa));
            auto pa = maybe_pa.get();

            // skip notes which belong to a different payment address in the wallet
            if (!(filterAddresses.empty() || filterAddresses.count(pa))) {
                continue;
            }

            if (ignoreSpent && nd.nullifier && wallet.IsSaplingSpent(*nd.nullifier)) {
                continue;
            }

            // skip notes which cannot be spent
            if (requireSpendingKey) {
                libzcash::SaplingIncomingViewingKey ivk;
                libzcash::SaplingFullViewingKey fvk;
                if (!(wallet.GetSaplingIncomingViewingKey(pa, ivk) &&
                    wallet.GetSaplingFullViewingKey(ivk, fvk) &&
                    wallet.HaveSaplingSpendingKey(fvk))) {
                    continue;
                }
            }

            // skip locked notes
            if (ignoreLocked && wallet.IsLockedNote(op)) {
                continue;
            }

            auto note = notePt.note(nd.ivk).get();
            saplingEntries.push_back(SaplingNoteEntry {
                op, pa, note, notePt.memo(), wtx.GetDepthInMainChain() });
        }
    }
}

class NoteCacheWallet : public CWallet {
public:
    bool EncryptKeys(CKeyingMaterial& vMasterKeyIn) { return CCryptoKeyStore::EncryptKeys(vMasterKeyIn); }
    bool Unlock(const CKeyingMaterial& vMasterKeyIn) { return CCryptoKeyStore::Unlock(vMasterKeyIn); }
    bool UpdatedNoteData(const CWalletTx& wtxIn, CWalletTx& wtx, bool fSameBlock = true) {
        return CWallet::UpdatedNoteData(wtxIn, wtx, fSameBlock);
    }

    // Drops every note's in-memory plaintext, as a restart does.
    void ClearPlaintextCache() {
        for (auto& p : mapWallet)
            for (auto& item : p.second.mapSaplingNoteData)
                item.second.cachedPlaintext = boost::none;
    }
    size_t CachedNotes() const {
        size_t n = 0;
        for (const auto& p : mapWallet)
            for (const auto& item : p.second.mapSaplingNoteData)
                n += item.second.cachedPlaintext ? 1 : 0;
        return n;
    }
};

void ExpectCacheMatchesFreshDecrypt(const SaplingNoteData& nd, const OutputDescription& od)
{
    ASSERT_TRUE(static_cast<bool>(nd.cachedPlaintext));
    auto fresh = SaplingNotePlaintext::decrypt_with_note(od.encCiphertext, nd.ivk, od.ephemeralKey, od.cm);
    ASSERT_TRUE(static_cast<bool>(fresh));
    const SaplingNote& cached = nd.cachedPlaintext->note;
    EXPECT_TRUE(cached.d == fresh->second.d);
    EXPECT_EQ(fresh->second.pk_d, cached.pk_d);
    EXPECT_EQ(fresh->second.r, cached.r);
    EXPECT_EQ(fresh->second.value(), cached.value());
    EXPECT_TRUE(nd.cachedPlaintext->memo == fresh->first.memo());
}

// The fake active chain is 0..TIP_HEIGHT; STALE_HEIGHT is a side block forked off it.
constexpr int TIP_HEIGHT = 30;
constexpr int STALE_HEIGHT = 25;

std::array<unsigned char, HUSH_MEMO_SIZE> MemoFor(uint64_t value)
{
    std::array<unsigned char, HUSH_MEMO_SIZE> memo;
    memo.fill(0);
    std::string s = "memo-" + std::to_string(value);
    memcpy(memo.data(), s.data(), s.size());
    return memo;
}

// A real encrypted Sapling output to pa. Its proof and value commitment are never checked by the
// wallet, so they are left empty / random.
OutputDescription MakeOutput(const SaplingPaymentAddress& pa, uint64_t value)
{
    SaplingNote note(pa, value);
    SaplingNotePlaintext pt(note, MemoFor(value));
    auto enc = pt.encrypt(note.pk_d);
    assert(static_cast<bool>(enc));
    OutputDescription od;
    od.cv = GetRandHash();
    od.cm = note.cm().get();
    od.ephemeralKey = enc->second.get_epk();
    od.encCiphertext = enc->first;
    return od;
}

CTransaction MakeTx(const std::vector<OutputDescription>& outputs, const std::vector<uint256>& spentNullifiers = {})
{
    CMutableTransaction mtx;
    mtx.fOverwintered = true;
    mtx.nVersion = SAPLING_TX_VERSION;
    mtx.nVersionGroupId = SAPLING_VERSION_GROUP_ID;
    mtx.vShieldedOutput = outputs;
    for (const uint256& nf : spentNullifiers) {
        SpendDescription sd;
        sd.nullifier = nf;
        mtx.vShieldedSpend.push_back(sd);
    }
    return CTransaction(mtx);
}

SaplingExtendedSpendingKey TestKey(uint32_t account)
{
    std::vector<unsigned char, secure_allocator<unsigned char>> raw(32, 0x5a);
    HDSeed seed(raw);
    return SaplingExtendedSpendingKey::Master(seed).Derive(account | ZIP32_HARDENED_KEY_LIMIT);
}

std::string Describe(int minDepth, int maxDepth, bool ignoreSpent, bool requireSpendingKey, bool ignoreLocked, const std::string& filter)
{
    return strprintf("minDepth=%d maxDepth=%d ignoreSpent=%d requireSpendingKey=%d ignoreLocked=%d filter=%s",
                     minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked, filter);
}

void ExpectSameEntries(const std::vector<SaplingNoteEntry>& want, const std::vector<SaplingNoteEntry>& got, const std::string& ctx)
{
    ASSERT_EQ(want.size(), got.size()) << ctx;
    for (size_t i = 0; i < want.size(); i++) {
        EXPECT_TRUE(want[i].op == got[i].op) << ctx << " entry " << i;
        EXPECT_TRUE(want[i].address == got[i].address) << ctx << " entry " << i;
        EXPECT_TRUE(want[i].note.d == got[i].note.d) << ctx << " entry " << i;
        EXPECT_EQ(want[i].note.pk_d, got[i].note.pk_d) << ctx << " entry " << i;
        EXPECT_EQ(want[i].note.r, got[i].note.r) << ctx << " entry " << i;
        EXPECT_EQ(want[i].note.value(), got[i].note.value()) << ctx << " entry " << i;
        EXPECT_TRUE(want[i].memo == got[i].memo) << ctx << " entry " << i;
        EXPECT_EQ(want[i].confirmations, got[i].confirmations) << ctx << " entry " << i;
    }
}

}

// A wallet with notes at several depths, spent and unspent, locked and unlocked, spendable and
// viewing-key-only, one mined only in a stale block, plus txs that carry no wallet notes; mined into a
// fake active chain. TearDown restores the process-wide chain state it changed (the chainActive tip,
// its own mapBlockIndex entries, pcoinsTip), because every gtest suite runs in the same process.
class WalletNoteCache : public ::testing::Test {
protected:
    NoteCacheWallet wallet;
    SaplingPaymentAddress paA, paB, paC;    // A and C spendable, B viewing key only
    std::map<uint256, int> depthOf;         // txid -> depth to mine it at (absent: never mined)
    std::vector<uint256> txids;
    SaplingOutPoint lockedOp;

    CCoinsView coinsBase;
    CCoinsViewCache coins{&coinsBase};
    CCoinsViewCache* savedCoinsTip = nullptr;
    std::deque<CBlock> blocks;
    std::deque<CBlockIndex> index;
    std::vector<uint256> insertedBlocks;    // mapBlockIndex keys this fixture added
    CBlockIndex* savedTip = nullptr;
    std::map<uint256, int> minedHeight;     // txid -> height of the block holding it (stale ones too)
    uint256 staleTx;                        // only ever mined in the side block at STALE_HEIGHT

    void SetUp() override {
        // tx_height()/GetTransaction dereference pcoinsTip; an empty view answers "not found".
        savedCoinsTip = pcoinsTip;
        pcoinsTip = &coins;
        {
            LOCK(cs_main);
            savedTip = chainActive.Tip();
        }

        LOCK2(cs_main, wallet.cs_wallet);
        auto skA = TestKey(0), skB = TestKey(1), skC = TestKey(2);
        paA = skA.DefaultAddress();
        paB = skB.DefaultAddress();
        paC = skC.DefaultAddress();
        ASSERT_TRUE(wallet.AddSaplingZKey(skA, paA));
        ASSERT_TRUE(wallet.AddSaplingIncomingViewingKey(skB.expsk.full_viewing_key().in_viewing_key(), paB));
        ASSERT_TRUE(wallet.AddSaplingZKey(skC, paC));
        SaplingPaymentAddress foreign = TestKey(9).DefaultAddress();   // not in the wallet

        AddTx(MakeTx({MakeOutput(paA, 1000), MakeOutput(paB, 2000)}), 1);
        uint256 tx2 = AddTx(MakeTx({MakeOutput(paA, 3000)}), 5);
        AddTx(MakeTx({MakeOutput(paA, 4000), MakeOutput(paC, 5000), MakeOutput(paB, 6000)}), 11);
        AddTx(MakeTx({MakeOutput(paC, 7000)}), 21);
        AddTx(MakeTx({MakeOutput(paA, 8000)}), -1);                      // never mined: depth -1
        AddTx(MakeTx({MakeOutput(foreign, 50)}), 3);                     // no wallet notes at all
        staleTx = AddTx(MakeTx({MakeOutput(paA, 11000)}), -1);         // mined in a stale block only

        // A note spent by a mined tx, and one whose spending tx never confirmed.
        const uint256 nfSpent = uint256S("5a1d"), nfPending = uint256S("9e9d");
        uint256 tx6 = AddTx(MakeTx({MakeOutput(paA, 9000)}), 12);
        uint256 tx7 = AddTx(MakeTx({MakeOutput(paC, 10000)}), 3);
        wallet.mapWallet.at(tx6).mapSaplingNoteData.at(SaplingOutPoint(tx6, 0)).nullifier = nfSpent;
        wallet.mapWallet.at(tx7).mapSaplingNoteData.at(SaplingOutPoint(tx7, 0)).nullifier = nfPending;
        AddTx(MakeTx({}, {nfSpent}), 2);
        AddTx(MakeTx({}, {nfPending}), -1);

        lockedOp = SaplingOutPoint(tx2, 0);
        wallet.LockNote(lockedOp);

        AddMoreTxs();
        Mine();
    }

    // Lets a test add txs before the chain is built.
    virtual void AddMoreTxs() {}

    void TearDown() override {
        LOCK(cs_main);
        chainActive.SetTip(savedTip);
        for (const uint256& hash : insertedBlocks)
            mapBlockIndex.erase(hash);
        pcoinsTip = savedCoinsTip;
    }

    // Adds tx the way LoadWallet does (no wallet.dat), with the note data FindMySaplingNotes finds.
    uint256 AddTx(const CTransaction& tx, int depth) {
        CWalletTx wtx(&wallet, tx);
        mapSaplingNoteData_t found = wallet.FindMySaplingNotes(tx).first;
        wtx.SetSaplingNoteData(found);
        EXPECT_TRUE(wallet.AddToWallet(wtx, true, nullptr));
        if (depth > 0)
            depthOf[tx.GetHash()] = depth;
        txids.push_back(tx.GetHash());
        return tx.GetHash();
    }

    CBlockIndex& AddBlock(CBlock& block, int height, CBlockIndex* pprev) {
        block.hashMerkleRoot = block.BuildMerkleTree();
        index.emplace_back(block);
        CBlockIndex& bi = index.back();
        bi.SetHeight(height);
        bi.pprev = pprev;
        auto inserted = mapBlockIndex.insert(std::make_pair(block.GetHash(), &bi));
        EXPECT_TRUE(inserted.second) << "block hash collision at height " << height;
        insertedBlocks.push_back(block.GetHash());
        bi.phashBlock = &inserted.first->first;
        return bi;
    }

    // Builds an active chain 0..TIP_HEIGHT and puts each wallet tx into the block at its depth, plus a
    // side block at STALE_HEIGHT (not on the active chain) holding staleTx.
    void Mine() {
        for (int h = 0; h <= TIP_HEIGHT; h++) {
            blocks.emplace_back();
            CBlock& block = blocks.back();
            block.nTime = 1700000000 + h;
            block.nNonce = uint256S(strprintf("%x", h + 1));
            if (h > 0)
                block.hashPrevBlock = blocks[h - 1].GetHash();
            for (const auto& d : depthOf) {
                if (TIP_HEIGHT - d.second + 1 == h)
                    block.vtx.push_back(wallet.mapWallet.at(d.first));
            }
            AddBlock(block, h, h > 0 ? &index[h - 1] : nullptr);
        }
        chainActive.SetTip(&index[TIP_HEIGHT]);
        for (const auto& d : depthOf) {
            int h = TIP_HEIGHT - d.second + 1;
            wallet.mapWallet.at(d.first).SetMerkleBranch(blocks[h]);
            minedHeight[d.first] = h;
        }

        blocks.emplace_back();
        CBlock& side = blocks.back();
        side.nTime = 1800000000;
        side.hashPrevBlock = blocks[STALE_HEIGHT - 1].GetHash();
        side.vtx.push_back(wallet.mapWallet.at(staleTx));
        AddBlock(side, STALE_HEIGHT, &index[STALE_HEIGHT - 1]);
        wallet.mapWallet.at(staleTx).SetMerkleBranch(side);
        minedHeight[staleTx] = STALE_HEIGHT;   // a txindex can still point at a stale block
    }

    // What tx_height() returns on a -txindex node: the height of the block the txindex points at
    // (possibly stale), 0 if none.
    TxHeightFn TxIndexHeight() const {
        return [this](const uint256& txid) {
            auto it = minedHeight.find(txid);
            return it == minedHeight.end() ? 0 : it->second;
        };
    }

    std::vector<SaplingNoteEntry> Filtered(const std::set<PaymentAddress>& filter, int minDepth, int maxDepth,
                                           bool ignoreSpent, bool requireSpendingKey, bool ignoreLocked) {
        std::set<PaymentAddress> f = filter;
        std::vector<SaplingNoteEntry> entries;
        wallet.GetFilteredNotes(entries, f, minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked);
        return entries;
    }

    std::vector<SaplingNoteEntry> Legacy(const std::set<PaymentAddress>& filter, int minDepth, int maxDepth,
                                         bool ignoreSpent, bool requireSpendingKey, bool ignoreLocked) {
        std::set<PaymentAddress> f = filter;
        std::vector<SaplingNoteEntry> entries;
        LegacyFilteredNotes(wallet, entries, f, minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked,
                            TxIndexHeight());
        return entries;
    }

    // cold: drop the plaintext cache before every call, so each one decrypts every note it returns.
    void ExpectEquivalent(const std::vector<int>& minDepths, bool cold = false) {
        const std::vector<std::pair<std::string, std::set<PaymentAddress>>> filters = {
            {"none", {}}, {"A", {paA}}, {"B", {paB}}, {"A+C", {paA, paC}},
        };
        for (int minDepth : minDepths)
        for (int maxDepth : {INT_MAX, 11})
        for (bool ignoreSpent : {false, true})
        for (bool requireSpendingKey : {false, true})
        for (bool ignoreLocked : {false, true})
        for (const auto& filter : filters) {
            std::string ctx = Describe(minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked, filter.first);
            if (cold) {
                LOCK(wallet.cs_wallet);
                wallet.ClearPlaintextCache();
            }
            ExpectSameEntries(Legacy(filter.second, minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked),
                              Filtered(filter.second, minDepth, maxDepth, ignoreSpent, requireSpendingKey, ignoreLocked),
                              ctx);
        }
    }
};

// The fixture must exercise every filter, or the equivalence check below could pass vacuously.
TEST_F(WalletNoteCache, FixtureExercisesEveryFilter) {
    // Mined, unspent notes: tx1 (2), tx2 (1), tx3 (3), tx4 (1), tx7 (1); tx6's note is spent.
    EXPECT_EQ(8u, Filtered({}, 1, INT_MAX, true, false, false).size());
    EXPECT_EQ(9u, Filtered({}, 1, INT_MAX, false, false, false).size()) << "ignoreSpent drops tx6's note";
    EXPECT_EQ(8u, Filtered({}, 0, INT_MAX, true, false, false).size()) << "a never-mined tx is at depth -1, not 0";
    EXPECT_EQ(6u, Filtered({}, 1, INT_MAX, true, true, false).size()) << "requireSpendingKey drops B's 2 notes";
    EXPECT_EQ(7u, Filtered({}, 1, INT_MAX, true, false, true).size()) << "ignoreLocked drops tx2's note";
    EXPECT_EQ(7u, Filtered({}, 1, 11, true, false, false).size()) << "maxDepth 11 drops tx4";
    EXPECT_EQ(2u, Filtered({paB}, 1, INT_MAX, true, false, false).size());

    auto onlyC = Filtered({paC}, 1, INT_MAX, true, false, false);
    ASSERT_EQ(3u, onlyC.size());
    for (const SaplingNoteEntry& e : onlyC) {
        EXPECT_TRUE(e.address == paC);
        EXPECT_TRUE(e.memo == MemoFor(e.note.value()));
    }
}

// The walk returns exactly what the 69bb035e7 loop returned on a -txindex node, in the same order,
// for every filter combination. minDepth 11/12/22 sit on the fixture's depth-11 and depth-21 txs.
TEST_F(WalletNoteCache, FilteredNotesMatchLegacy) {
    ExpectEquivalent({0, 1, 2, 11, 12, 22});
}

// minDepth > 1 (sweep and consolidation use 11) locates each tx's block from its depth. The old
// tx_height() re-read the tx through the txindex, and without one (as here, and on a -txindex=0
// node) it never found a fully shielded tx, so nothing was returned at all.
TEST_F(WalletNoteCache, MinDepthAboveOneUsesWalletHeight) {
    // Unspent notes at depth >= 11: tx3's 3 (depth 11) and tx4's 1 (depth 21); tx6 (12) is spent.
    EXPECT_EQ(4u, Filtered({}, 11, INT_MAX, true, false, false).size());
    EXPECT_EQ(5u, Filtered({}, 11, INT_MAX, false, false, false).size());
    EXPECT_EQ(1u, Filtered({}, 12, INT_MAX, true, false, false).size()) << "depth 11 is below minDepth 12";
    EXPECT_EQ(3u, Filtered({}, 11, 11, true, false, false).size()) << "depth 21 is above maxDepth 11";
    EXPECT_EQ(0u, Filtered({}, 22, INT_MAX, false, false, false).size());
    for (const SaplingNoteEntry& e : Filtered({}, 2, INT_MAX, false, false, false))
        EXPECT_GE(e.confirmations, 2) << "unmined and shallower txs stay out";

    std::set<PaymentAddress> none;
    std::vector<SaplingNoteEntry> legacy;
    LegacyFilteredNotes(wallet, legacy, none, 11, INT_MAX, true, false, false, tx_height);
    EXPECT_EQ(0u, legacy.size()) << "the old code found nothing here without a txindex";
}

TEST_F(WalletNoteCache, HeightFromDepth) {
    LOCK2(cs_main, wallet.cs_wallet);
    EXPECT_EQ(-1, wallet.mapWallet.at(staleTx).GetDepthInMainChain()) << "a stale block is not the active chain";
    EXPECT_EQ(TIP_HEIGHT, HeightFromDepth(1));
    EXPECT_EQ(TIP_HEIGHT - 10, HeightFromDepth(11));
    EXPECT_EQ(0, HeightFromDepth(0)) << "mempool";
    EXPECT_EQ(0, HeightFromDepth(-1)) << "not in the chain or mempool";
    for (const auto& d : depthOf)
        EXPECT_EQ(minedHeight.at(d.first), HeightFromDepth(wallet.mapWallet.at(d.first).GetDepthInMainChain()));
}

// Two notes recorded under A's viewing key that it can't actually decrypt (real outputs to someone
// else): output 0 is spent by a mined tx, output 1 is locked. If the walk decrypted either one, the
// decrypt assert would abort. A third, locked too, names an output the tx doesn't have.
class WalletNoteCacheUndecryptable : public WalletNoteCache {
protected:
    uint256 txBad;

    void AddMoreTxs() override {
        SaplingPaymentAddress foreign = TestKey(9).DefaultAddress();
        CTransaction tx = MakeTx({MakeOutput(foreign, 1), MakeOutput(foreign, 2)});
        txBad = tx.GetHash();
        const SaplingIncomingViewingKey ivkA = TestKey(0).expsk.full_viewing_key().in_viewing_key();
        const uint256 nf = uint256S("bad5");

        CWalletTx wtx(&wallet, tx);
        wtx.mapSaplingNoteData[SaplingOutPoint(txBad, 0)] = SaplingNoteData(ivkA, nf);
        wtx.mapSaplingNoteData[SaplingOutPoint(txBad, 1)] = SaplingNoteData(ivkA);
        wtx.mapSaplingNoteData[SaplingOutPoint(txBad, 7)] = SaplingNoteData(ivkA);
        EXPECT_TRUE(wallet.AddToWallet(wtx, true, nullptr));
        depthOf[txBad] = 4;
        txids.push_back(txBad);

        AddTx(MakeTx({}, {nf}), 4);
        wallet.LockNote(SaplingOutPoint(txBad, 1));
        wallet.LockNote(SaplingOutPoint(txBad, 7));
    }
};

// Spent and locked notes are dropped before the decryption (~0.5-0.8ms per note), not after it.
TEST_F(WalletNoteCacheUndecryptable, SpentAndLockedNotesAreNotDecrypted) {
    auto entries = Filtered({}, 1, INT_MAX, true, false, true);
    EXPECT_EQ(7u, entries.size()) << "the fixture's 8 notes less tx2's locked one";
    for (const SaplingNoteEntry& e : entries)
        EXPECT_FALSE(e.op.hash == txBad);
}

// The same wallet with its keystore encrypted and locked: spendability then comes from the
// encrypted-key map (CCryptoKeyStore::HaveSaplingSpendingKey), and nothing else may change.
class WalletNoteCacheEncrypted : public WalletNoteCache {
protected:
    void AddMoreTxs() override {
        CKeyingMaterial masterKey(32, 0x42);
        ASSERT_TRUE(wallet.EncryptKeys(masterKey));
        ASSERT_TRUE(wallet.IsCrypted());
        ASSERT_TRUE(wallet.IsLocked());
    }
};

TEST_F(WalletNoteCacheEncrypted, LockedWalletMatchesLegacy) {
    EXPECT_TRUE(wallet.HaveSaplingSpendingKey(TestKey(0).expsk.full_viewing_key()));
    EXPECT_FALSE(wallet.HaveSaplingSpendingKey(TestKey(1).expsk.full_viewing_key())) << "B is viewing key only";
    EXPECT_FALSE(wallet.HaveSaplingSpendingKey(TestKey(9).expsk.full_viewing_key()));
    EXPECT_EQ(6u, Filtered({}, 1, INT_MAX, true, true, false).size()) << "A's and C's keys are still known";
    ExpectEquivalent({1, 11});
}

TEST_F(WalletNoteCacheEncrypted, UnlockedWalletFindsSpendingKeys) {
    libzcash::SaplingExtendedSpendingKey sk;
    EXPECT_FALSE(wallet.GetSaplingSpendingKey(TestKey(0).expsk.full_viewing_key(), sk)) << "locked";
    ASSERT_TRUE(wallet.Unlock(CKeyingMaterial(32, 0x42)));
    ASSERT_TRUE(wallet.GetSaplingSpendingKey(TestKey(0).expsk.full_viewing_key(), sk));
    EXPECT_TRUE(sk.DefaultAddress() == paA);
    ASSERT_TRUE(wallet.GetSaplingSpendingKey(TestKey(2).expsk.full_viewing_key(), sk));
    EXPECT_TRUE(sk.DefaultAddress() == paC);
    EXPECT_FALSE(wallet.GetSaplingSpendingKey(TestKey(9).expsk.full_viewing_key(), sk));
}

namespace {

// The default address plus a few diversified addresses of the same key.
std::vector<SaplingPaymentAddress> AddressesOf(const SaplingExtendedSpendingKey& sk)
{
    std::vector<SaplingPaymentAddress> addrs = {sk.DefaultAddress()};
    for (unsigned char start : {7, 90, 201}) {
        diversifier_index_t j;
        j.begin()[0] = start;
        auto found = sk.ToXFVK().Address(j);
        if (found)
            addrs.push_back(found->second);
    }
    return addrs;
}

}

// decrypt_with_note returns exactly decrypt() plus decrypt()->note(ivk), and fails exactly when they do.
TEST(SaplingDecryptWithNote, MatchesDecryptThenNote) {
    auto sk = TestKey(0);
    const SaplingIncomingViewingKey ivk = sk.expsk.full_viewing_key().in_viewing_key();
    const SaplingIncomingViewingKey otherIvk = TestKey(1).expsk.full_viewing_key().in_viewing_key();
    auto addrs = AddressesOf(sk);
    ASSERT_GE(addrs.size(), 3u) << "need diversified addresses too";

    uint64_t value = 100;
    for (const SaplingPaymentAddress& pa : addrs) {
        OutputDescription od = MakeOutput(pa, ++value);
        auto legacyPt = SaplingNotePlaintext::decrypt(od.encCiphertext, ivk, od.ephemeralKey, od.cm);
        ASSERT_TRUE(static_cast<bool>(legacyPt));
        auto legacyNote = legacyPt->note(ivk);
        ASSERT_TRUE(static_cast<bool>(legacyNote));

        auto both = SaplingNotePlaintext::decrypt_with_note(od.encCiphertext, ivk, od.ephemeralKey, od.cm);
        ASSERT_TRUE(static_cast<bool>(both));
        const SaplingNotePlaintext& pt = both->first;
        const SaplingNote& note = both->second;
        EXPECT_TRUE(pt.d == legacyPt->d);
        EXPECT_EQ(legacyPt->rcm, pt.rcm);
        EXPECT_EQ(legacyPt->value(), pt.value());
        EXPECT_TRUE(pt.memo() == legacyPt->memo());
        EXPECT_TRUE(note.d == legacyNote->d);
        EXPECT_EQ(legacyNote->pk_d, note.pk_d);
        EXPECT_EQ(legacyNote->r, note.r);
        EXPECT_EQ(legacyNote->value(), note.value());
        EXPECT_TRUE(SaplingPaymentAddress(note.d, note.pk_d) == pa);
        EXPECT_EQ(value, note.value());

        EXPECT_FALSE(static_cast<bool>(SaplingNotePlaintext::decrypt_with_note(od.encCiphertext, otherIvk, od.ephemeralKey, od.cm)));
        uint256 badCm = od.cm;
        badCm.begin()[0] ^= 1;
        EXPECT_FALSE(static_cast<bool>(SaplingNotePlaintext::decrypt(od.encCiphertext, ivk, od.ephemeralKey, badCm)));
        EXPECT_FALSE(static_cast<bool>(SaplingNotePlaintext::decrypt_with_note(od.encCiphertext, ivk, od.ephemeralKey, badCm)));
    }
}

// What FindMySaplingNotes discovers: notes to a spending key's default and diversified addresses, and to
// a viewing-key-only address, each with the key that found it; a diversified address the wallet hasn't
// seen yet is reported for registration; outputs to other wallets are ignored.
TEST(WalletNoteCacheDiscovery, FindMySaplingNotes) {
    NoteCacheWallet wallet;
    auto skA = TestKey(0), skB = TestKey(1);
    const SaplingIncomingViewingKey ivkA = skA.expsk.full_viewing_key().in_viewing_key();
    const SaplingIncomingViewingKey ivkB = skB.expsk.full_viewing_key().in_viewing_key();
    auto addrsA = AddressesOf(skA);
    const SaplingPaymentAddress paB = skB.DefaultAddress();

    LOCK2(cs_main, wallet.cs_wallet);
    ASSERT_TRUE(wallet.AddSaplingZKey(skA, addrsA[0]));
    ASSERT_TRUE(wallet.AddSaplingIncomingViewingKey(ivkB, paB));

    CTransaction tx = MakeTx({MakeOutput(addrsA[0], 1), MakeOutput(addrsA[1], 2), MakeOutput(paB, 3),
                              MakeOutput(TestKey(9).DefaultAddress(), 4), MakeOutput(addrsA[2], 5)});
    auto found = wallet.FindMySaplingNotes(tx);
    const mapSaplingNoteData_t& notes = found.first;
    ASSERT_EQ(4u, notes.size());
    EXPECT_TRUE(notes.at(SaplingOutPoint(tx.GetHash(), 0)).ivk == ivkA);
    EXPECT_TRUE(notes.at(SaplingOutPoint(tx.GetHash(), 1)).ivk == ivkA);
    EXPECT_TRUE(notes.at(SaplingOutPoint(tx.GetHash(), 2)).ivk == ivkB);
    EXPECT_EQ(0u, notes.count(SaplingOutPoint(tx.GetHash(), 3)));
    EXPECT_TRUE(notes.at(SaplingOutPoint(tx.GetHash(), 4)).ivk == ivkA);
    for (const auto& item : notes) {
        EXPECT_FALSE(static_cast<bool>(item.second.nullifier));
        EXPECT_TRUE(item.second.witnesses.empty());
        EXPECT_EQ(-1, item.second.witnessHeight);
    }

    const SaplingIncomingViewingKeyMap& toAdd = found.second;
    ASSERT_EQ(2u, toAdd.size()) << "the two diversified addresses, not the registered default one";
    EXPECT_TRUE(toAdd.at(addrsA[1]) == ivkA);
    EXPECT_TRUE(toAdd.at(addrsA[2]) == ivkA);
}

// Every note FindMySaplingNotes finds comes with its plaintext, for both the spending-key loop and the
// viewing-key-only loop.
TEST(WalletNoteCacheDiscovery, FindMySaplingNotesFillsCache) {
    NoteCacheWallet wallet;
    auto skA = TestKey(0), skB = TestKey(1);
    auto addrsA = AddressesOf(skA);
    LOCK2(cs_main, wallet.cs_wallet);
    ASSERT_TRUE(wallet.AddSaplingZKey(skA, addrsA[0]));
    ASSERT_TRUE(wallet.AddSaplingIncomingViewingKey(skB.expsk.full_viewing_key().in_viewing_key(), skB.DefaultAddress()));
    CTransaction tx = MakeTx({MakeOutput(addrsA[0], 11), MakeOutput(addrsA[1], 12), MakeOutput(skB.DefaultAddress(), 13)});
    auto notes = wallet.FindMySaplingNotes(tx).first;
    ASSERT_EQ(3u, notes.size());
    for (const auto& item : notes)
        ExpectCacheMatchesFreshDecrypt(item.second, tx.vShieldedOutput[item.first.n]);
}

// With the cache dropped before every call (a restart), the walk still returns exactly the old loop's
// entries: the lazy fill decrypts the same notes the same way.
TEST_F(WalletNoteCache, ColdCacheMatchesLegacy) {
    ExpectEquivalent({0, 1, 11}, true);
}

// The cache holds exactly what a fresh decryption gives, wherever it was filled.
TEST_F(WalletNoteCache, CachedPlaintextMatchesFreshDecrypt) {
    LOCK2(cs_main, wallet.cs_wallet);
    for (int pass = 0; pass < 2; pass++) {
        if (pass == 1) {
            wallet.ClearPlaintextCache();
            Filtered({}, 0, INT_MAX, false, false, false);   // lazy fill of every confirmed note
        }
        size_t checked = 0;
        for (const auto& p : wallet.mapWallet) {
            for (const auto& item : p.second.mapSaplingNoteData) {
                if (!item.second.cachedPlaintext)
                    continue;
                ExpectCacheMatchesFreshDecrypt(item.second, p.second.vShieldedOutput[item.first.n]);
                checked++;
            }
        }
        EXPECT_GE(checked, 9u) << "pass " << pass;
    }
}

// A warm walk decrypts nothing; a cold one decrypts each note it needs exactly once, and spent notes
// only when the caller asks for them.
TEST_F(WalletNoteCache, WarmPassDoesNotDecrypt) {
    LOCK2(cs_main, wallet.cs_wallet);
    const uint64_t base = wallet.nSaplingNoteCacheFills;
    Filtered({}, 1, INT_MAX, true, true, true);
    EXPECT_EQ(base, wallet.nSaplingNoteCacheFills) << "FindMySaplingNotes already filled every note";

    wallet.ClearPlaintextCache();
    EXPECT_EQ(0u, wallet.CachedNotes());
    Filtered({}, 1, INT_MAX, true, false, true);
    // tx1 (2), tx3 (3), tx4 (1), tx7 (1): tx2's note is locked and tx6's is spent, so neither is decrypted.
    EXPECT_EQ(base + 7, wallet.nSaplingNoteCacheFills);
    Filtered({}, 1, INT_MAX, true, false, true);
    Filtered({paA}, 0, INT_MAX, true, true, true);
    EXPECT_EQ(base + 7, wallet.nSaplingNoteCacheFills) << "warm";

    Filtered({}, 1, INT_MAX, false, false, false);
    EXPECT_EQ(base + 9, wallet.nSaplingNoteCacheFills) << "the locked and the spent note, once";
    Filtered({}, 1, INT_MAX, false, false, false);
    EXPECT_EQ(base + 9, wallet.nSaplingNoteCacheFills);
}

// The cache never reaches wallet.dat, the network or a tx hash: serialization is byte-identical with it
// filled or empty, and a deserialized entry comes back empty.
TEST_F(WalletNoteCache, CacheNeverSerialized) {
    LOCK2(cs_main, wallet.cs_wallet);
    const uint256 txid = txids[2];   // tx3: three notes
    CWalletTx filled = wallet.mapWallet.at(txid);
    ASSERT_EQ(3u, filled.mapSaplingNoteData.size());
    CWalletTx empty = filled;
    for (auto& item : empty.mapSaplingNoteData) {
        ASSERT_TRUE(static_cast<bool>(item.second.cachedPlaintext));
        item.second.cachedPlaintext = boost::none;
    }
    for (int type : {SER_DISK, SER_NETWORK, SER_GETHASH}) {
        CDataStream a(type, CLIENT_VERSION), b(type, CLIENT_VERSION);
        a << filled;
        b << empty;
        EXPECT_TRUE(a.str() == b.str()) << "CWalletTx, type " << type;
        CDataStream c(type, CLIENT_VERSION), d(type, CLIENT_VERSION);
        c << filled.mapSaplingNoteData.begin()->second;
        d << empty.mapSaplingNoteData.begin()->second;
        EXPECT_TRUE(c.str() == d.str()) << "SaplingNoteData, type " << type;
    }
    EXPECT_EQ(txid, filled.GetHash());

    CDataStream ss(SER_DISK, CLIENT_VERSION);
    ss << filled;
    CWalletTx back;
    ss >> back;
    ASSERT_EQ(3u, back.mapSaplingNoteData.size());
    for (const auto& item : back.mapSaplingNoteData)
        EXPECT_FALSE(static_cast<bool>(item.second.cachedPlaintext));
}

// UpdatedNoteData: a difference in the cache alone is no update (no rewrite to disk); the replace branch
// keeps a plaintext the wallet already had; and nothing is carried across a different ivk.
TEST_F(WalletNoteCache, UpdatedNoteDataAndTheCache) {
    LOCK2(cs_main, wallet.cs_wallet);
    const uint256 txid = txids[2];
    CWalletTx& wtx = wallet.mapWallet.at(txid);
    const SaplingOutPoint op0(txid, 0), op1(txid, 1);

    // Same note data, but arriving without plaintext: unchanged, and the wallet keeps its cache.
    CWalletTx replay = wtx;
    for (auto& item : replay.mapSaplingNoteData)
        item.second.cachedPlaintext = boost::none;
    EXPECT_FALSE(wallet.UpdatedNoteData(replay, wtx));
    EXPECT_TRUE(static_cast<bool>(wtx.mapSaplingNoteData.at(op0).cachedPlaintext));

    // A real update (an extra note) without plaintext for the existing ones: replaced, cache carried.
    CWalletTx incoming = replay;
    incoming.mapSaplingNoteData.erase(op1);
    wtx.mapSaplingNoteData.erase(op1);
    incoming.mapSaplingNoteData[op1] = SaplingNoteData(TestKey(2).expsk.full_viewing_key().in_viewing_key());
    EXPECT_TRUE(wallet.UpdatedNoteData(incoming, wtx));
    ExpectCacheMatchesFreshDecrypt(wtx.mapSaplingNoteData.at(op0), wtx.vShieldedOutput[0]);
    EXPECT_FALSE(static_cast<bool>(wtx.mapSaplingNoteData.at(op1).cachedPlaintext)) << "nothing to carry";

    // Incoming note data under another ivk: its (absent) plaintext is not replaced by ours.
    CWalletTx otherKey = wtx;
    for (auto& item : otherKey.mapSaplingNoteData) {
        item.second.ivk = TestKey(9).expsk.full_viewing_key().in_viewing_key();
        item.second.cachedPlaintext = boost::none;
    }
    EXPECT_TRUE(wallet.UpdatedNoteData(otherKey, wtx));
    for (const auto& item : wtx.mapSaplingNoteData)
        EXPECT_FALSE(static_cast<bool>(item.second.cachedPlaintext));
}

// Spendability is decided per call, never cached: a note to a viewing-key-only address joins the
// requireSpendingKey results as soon as its spending key is added, with its plaintext already cached.
TEST_F(WalletNoteCache, SpendabilityNotCached) {
    LOCK2(cs_main, wallet.cs_wallet);
    auto before = Filtered({paB}, 1, INT_MAX, true, true, false);
    EXPECT_EQ(0u, before.size());
    EXPECT_EQ(2u, Filtered({paB}, 1, INT_MAX, true, false, false).size());
    const uint64_t fills = wallet.nSaplingNoteCacheFills;
    ASSERT_TRUE(wallet.AddSaplingZKey(TestKey(1), paB));
    EXPECT_EQ(2u, Filtered({paB}, 1, INT_MAX, true, true, false).size());
    EXPECT_EQ(fills, wallet.nSaplingNoteCacheFills);
}

// The startup warm-up decrypts every note the wallet holds: spent, locked, unconfirmed and stale ones too
// (z_listreceivedbyaddress asks for spent notes), after which a walk decrypts nothing.
TEST_F(WalletNoteCache, WarmUpFillsEveryNote) {
    LOCK2(cs_main, wallet.cs_wallet);
    size_t total = 0;
    for (const auto& p : wallet.mapWallet)
        total += p.second.mapSaplingNoteData.size();
    ASSERT_EQ(11u, total);   // tx1 2, tx2 1, tx3 3, tx4 1, unmined 1, stale 1, tx6 1, tx7 1

    wallet.ClearPlaintextCache();
    const uint64_t base = wallet.nSaplingNoteCacheFills;
    wallet.WarmSaplingNoteCache();
    EXPECT_EQ(total, wallet.CachedNotes());
    EXPECT_EQ(base + total, wallet.nSaplingNoteCacheFills);

    wallet.WarmSaplingNoteCache();
    EXPECT_EQ(base + total, wallet.nSaplingNoteCacheFills) << "nothing left to do";
    Filtered({}, 0, INT_MAX, false, false, false);
    Filtered({paC}, 11, INT_MAX, true, true, true);
    EXPECT_EQ(base + total, wallet.nSaplingNoteCacheFills) << "walks find everything cached";
}

// The warm-up never aborts: notes that don't decrypt, or name an output the tx doesn't have, are left
// unfilled and everything else is filled.
TEST_F(WalletNoteCacheUndecryptable, WarmUpSkipsBadNotes) {
    LOCK2(cs_main, wallet.cs_wallet);
    wallet.ClearPlaintextCache();
    wallet.WarmSaplingNoteCache();
    const auto& bad = wallet.mapWallet.at(txBad).mapSaplingNoteData;
    ASSERT_EQ(3u, bad.size());
    for (const auto& item : bad)
        EXPECT_FALSE(static_cast<bool>(item.second.cachedPlaintext)) << item.first.n;
    size_t cached = 0, total = 0;
    for (const auto& p : wallet.mapWallet) {
        for (const auto& item : p.second.mapSaplingNoteData) {
            total++;
            cached += item.second.cachedPlaintext ? 1 : 0;
        }
    }
    EXPECT_EQ(total - 3, cached);
    EXPECT_EQ(7u, Filtered({}, 1, INT_MAX, true, false, true).size()) << "the bad notes are spent or locked";
}