diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp index 474129c34..d30944954 100644 --- a/src/wallet/gtest/test_wallet.cpp +++ b/src/wallet/gtest/test_wallet.cpp @@ -50,8 +50,8 @@ public: ZCIncrementalMerkleTree tree) { CWallet::IncrementNoteWitnesses(pindex, pblock, tree); } - void DecrementNoteWitnesses() { - CWallet::DecrementNoteWitnesses(); + void DecrementNoteWitnesses(const CBlockIndex* pindex) { + CWallet::DecrementNoteWitnesses(pindex); } void WriteWitnessCache(MockWalletDB& walletdb) { CWallet::WriteWitnessCache(walletdb); @@ -675,7 +675,7 @@ TEST(wallet_tests, cached_witnesses_empty_chain) { EXPECT_TRUE((bool) witnesses[1]); // Until #1302 is implemented, this should triggger an assertion - EXPECT_DEATH(wallet.DecrementNoteWitnesses(), + EXPECT_DEATH(wallet.DecrementNoteWitnesses(&index), "Assertion `nWitnessCacheSize > 0' failed."); } @@ -748,7 +748,7 @@ TEST(wallet_tests, cached_witnesses_chain_tip) { // Decrementing should give us the previous anchor uint256 anchor3; - wallet.DecrementNoteWitnesses(); + wallet.DecrementNoteWitnesses(&index2); witnesses.clear(); wallet.GetNoteWitnesses(notes, witnesses, anchor3); EXPECT_FALSE((bool) witnesses[0]); @@ -773,6 +773,101 @@ TEST(wallet_tests, cached_witnesses_chain_tip) { } } +TEST(wallet_tests, CachedWitnessesDecrementFirst) { + TestWallet wallet; + uint256 anchor2; + CBlock block2; + CBlockIndex index2(block2); + ZCIncrementalMerkleTree tree; + + auto sk = libzcash::SpendingKey::random(); + wallet.AddSpendingKey(sk); + + { + // First transaction (case tested in _empty_chain) + auto wtx = GetValidReceive(sk, 10, true); + auto note = GetNote(sk, wtx, 0, 1); + auto nullifier = note.nullifier(sk); + + mapNoteData_t noteData; + JSOutPoint jsoutpt {wtx.GetHash(), 0, 1}; + CNoteData nd {sk.address(), nullifier}; + noteData[jsoutpt] = nd; + wtx.SetNoteData(noteData); + wallet.AddToWallet(wtx, true, NULL); + + // First block (case tested in _empty_chain) + CBlock block1; + block1.vtx.push_back(wtx); + CBlockIndex index1(block1); + index1.nHeight = 1; + wallet.IncrementNoteWitnesses(&index1, &block1, tree); + } + + { + // Second transaction (case tested in _chain_tip) + auto wtx = GetValidReceive(sk, 50, true); + auto note = GetNote(sk, wtx, 0, 1); + auto nullifier = note.nullifier(sk); + + mapNoteData_t noteData; + JSOutPoint jsoutpt {wtx.GetHash(), 0, 1}; + CNoteData nd {sk.address(), nullifier}; + noteData[jsoutpt] = nd; + wtx.SetNoteData(noteData); + wallet.AddToWallet(wtx, true, NULL); + + std::vector notes {jsoutpt}; + std::vector> witnesses; + + // Second block (case tested in _chain_tip) + block2.vtx.push_back(wtx); + index2.nHeight = 2; + wallet.IncrementNoteWitnesses(&index2, &block2, tree); + // Called to fetch anchor + wallet.GetNoteWitnesses(notes, witnesses, anchor2); + } + + { + // Third transaction - never mined + auto wtx = GetValidReceive(sk, 20, true); + auto note = GetNote(sk, wtx, 0, 1); + auto nullifier = note.nullifier(sk); + + mapNoteData_t noteData; + JSOutPoint jsoutpt {wtx.GetHash(), 0, 1}; + CNoteData nd {sk.address(), nullifier}; + noteData[jsoutpt] = nd; + wtx.SetNoteData(noteData); + wallet.AddToWallet(wtx, true, NULL); + + std::vector notes {jsoutpt}; + std::vector> witnesses; + uint256 anchor3; + + wallet.GetNoteWitnesses(notes, witnesses, anchor3); + EXPECT_FALSE((bool) witnesses[0]); + + // Decrementing (before the transaction has ever seen an increment) + // should give us the previous anchor + uint256 anchor4; + wallet.DecrementNoteWitnesses(&index2); + witnesses.clear(); + wallet.GetNoteWitnesses(notes, witnesses, anchor4); + EXPECT_FALSE((bool) witnesses[0]); + // Should not equal second anchor because none of these notes had witnesses + EXPECT_NE(anchor2, anchor4); + + // Re-incrementing with the same block should give the same result + uint256 anchor5; + wallet.IncrementNoteWitnesses(&index2, &block2, tree); + witnesses.clear(); + wallet.GetNoteWitnesses(notes, witnesses, anchor5); + EXPECT_FALSE((bool) witnesses[0]); + EXPECT_EQ(anchor3, anchor5); + } +} + TEST(wallet_tests, ClearNoteWitnessCache) { TestWallet wallet; @@ -780,6 +875,7 @@ TEST(wallet_tests, ClearNoteWitnessCache) { wallet.AddSpendingKey(sk); auto wtx = GetValidReceive(sk, 10, true); + auto hash = wtx.GetHash(); auto note = GetNote(sk, wtx, 0, 0); auto nullifier = note.nullifier(sk); @@ -793,6 +889,8 @@ TEST(wallet_tests, ClearNoteWitnessCache) { // Pretend we mined the tx by adding a fake witness ZCIncrementalMerkleTree tree; wtx.mapNoteData[jsoutpt].witnesses.push_front(tree.witness()); + wtx.mapNoteData[jsoutpt].witnessHeight = 1; + wallet.nWitnessCacheSize = 1; wallet.AddToWallet(wtx, true, NULL); @@ -804,6 +902,8 @@ TEST(wallet_tests, ClearNoteWitnessCache) { wallet.GetNoteWitnesses(notes, witnesses, anchor2); EXPECT_TRUE((bool) witnesses[0]); EXPECT_FALSE((bool) witnesses[1]); + EXPECT_EQ(1, wallet.mapWallet[hash].mapNoteData[jsoutpt].witnessHeight); + EXPECT_EQ(1, wallet.nWitnessCacheSize); // After clearing, we should not have a witness for either note wallet.ClearNoteWitnessCache(); @@ -811,6 +911,8 @@ TEST(wallet_tests, ClearNoteWitnessCache) { wallet.GetNoteWitnesses(notes, witnesses, anchor2); EXPECT_FALSE((bool) witnesses[0]); EXPECT_FALSE((bool) witnesses[1]); + EXPECT_EQ(-1, wallet.mapWallet[hash].mapNoteData[jsoutpt].witnessHeight); + EXPECT_EQ(0, wallet.nWitnessCacheSize); } TEST(wallet_tests, WriteWitnessCache) { @@ -932,6 +1034,7 @@ TEST(wallet_tests, UpdatedNoteData) { // Pretend we mined the tx by adding a fake witness ZCIncrementalMerkleTree tree; wtx.mapNoteData[jsoutpt].witnesses.push_front(tree.witness()); + wtx.mapNoteData[jsoutpt].witnessHeight = 100; // Now pretend we added the key for the second note, and // the tx was "added" to the wallet again to update it. @@ -944,11 +1047,13 @@ TEST(wallet_tests, UpdatedNoteData) { // The txs should initially be different EXPECT_NE(wtx.mapNoteData, wtx2.mapNoteData); EXPECT_EQ(1, wtx.mapNoteData[jsoutpt].witnesses.size()); + EXPECT_EQ(100, wtx.mapNoteData[jsoutpt].witnessHeight); // After updating, they should be the same EXPECT_TRUE(wallet.UpdatedNoteData(wtx2, wtx)); EXPECT_EQ(wtx.mapNoteData, wtx2.mapNoteData); EXPECT_EQ(1, wtx.mapNoteData[jsoutpt].witnesses.size()); + EXPECT_EQ(100, wtx.mapNoteData[jsoutpt].witnessHeight); // TODO: The new note should get witnessed (but maybe not here) (#1350) } diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index e549e5484..51ed76f3a 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -373,7 +373,7 @@ void CWallet::ChainTip(const CBlockIndex *pindex, const CBlock *pblock, if (added) { IncrementNoteWitnesses(pindex, pblock, tree); } else { - DecrementNoteWitnesses(); + DecrementNoteWitnesses(pindex); } } @@ -630,8 +630,10 @@ void CWallet::ClearNoteWitnessCache() for (std::pair& wtxItem : mapWallet) { for (mapNoteData_t::value_type& item : wtxItem.second.mapNoteData) { item.second.witnesses.clear(); + item.second.witnessHeight = -1; } } + nWitnessCacheSize = 0; } void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex, @@ -648,7 +650,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex, // Only increment witnesses that are behind the current height if (nd->witnessHeight < pindex->nHeight) { // Witnesses being incremented should always be either -1 - // (never incremented) or one below pindex + // (never incremented or decremented) or one below pindex assert((nd->witnessHeight == -1) || (nd->witnessHeight == pindex->nHeight - 1)); // Copy the witness for the previous block if we have one @@ -746,7 +748,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex, } } -void CWallet::DecrementNoteWitnesses() +void CWallet::DecrementNoteWitnesses(const CBlockIndex* pindex) { { LOCK(cs_wallet); @@ -755,10 +757,16 @@ void CWallet::DecrementNoteWitnesses() CNoteData* nd = &(item.second); // Check the validity of the cache assert(nWitnessCacheSize >= nd->witnesses.size()); + // Witnesses being decremented should always be either -1 + // (never incremented or decremented) or equal to pindex + assert((nd->witnessHeight == -1) || + (nd->witnessHeight == pindex->nHeight)); if (nd->witnesses.size() > 0) { nd->witnesses.pop_front(); } - nd->witnessHeight -= 1; + // pindex is the block being removed, so the new witness cache + // height is one below it. + nd->witnessHeight = pindex->nHeight - 1; } } nWitnessCacheSize -= 1; @@ -1102,6 +1110,7 @@ bool CWallet::UpdatedNoteData(const CWalletTx& wtxIn, CWalletTx& wtx) tmp.at(nd.first).witnesses.assign( nd.second.witnesses.cbegin(), nd.second.witnesses.cend()); } + tmp.at(nd.first).witnessHeight = nd.second.witnessHeight; } // Now copy over the updated note data wtx.mapNoteData = tmp; diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index dfcbc1af8..85f960eb1 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -221,7 +221,15 @@ public: */ std::list witnesses; - /** Block height corresponding to the most current witness. */ + /** + * Block height corresponding to the most current witness. + * + * When we first create a CNoteData in CWallet::FindMyNotes, this is set to + * -1 as a placeholder. The next time CWallet::ChainTip is called, we can + * determine what height the witness cache for this note is valid for (even + * if no witnesses were cached), and so can set the correct value in + * CWallet::IncrementNoteWitnesses and CWallet::DecrementNoteWitnesses. + */ int witnessHeight; CNoteData() : address(), nullifier(), witnessHeight {-1} { } @@ -616,10 +624,16 @@ public: void ClearNoteWitnessCache(); protected: + /** + * pindex is the new tip being connected. + */ void IncrementNoteWitnesses(const CBlockIndex* pindex, const CBlock* pblock, ZCIncrementalMerkleTree& tree); - void DecrementNoteWitnesses(); + /** + * pindex is the old tip being disconnected. + */ + void DecrementNoteWitnesses(const CBlockIndex* pindex); template void WriteWitnessCache(WalletDB& walletdb) {