diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp index e72b8a2f6..2ea2c1d5b 100644 --- a/src/wallet/gtest/test_wallet.cpp +++ b/src/wallet/gtest/test_wallet.cpp @@ -11,11 +11,26 @@ #include "zcash/Note.hpp" #include "zcash/NoteEncryption.hpp" -using ::testing::_; using ::testing::Return; ZCJoinSplit* params = ZCJoinSplit::Unopened(); +ACTION(ThrowLogicError) { + throw std::logic_error("Boom"); +} + +class MockWalletDB { +public: + MOCK_METHOD0(TxnBegin, bool()); + MOCK_METHOD0(TxnCommit, bool()); + MOCK_METHOD0(TxnAbort, bool()); + + MOCK_METHOD2(WriteTx, bool(uint256 hash, const CWalletTx& wtx)); + MOCK_METHOD1(WriteWitnessCacheSize, bool(int64_t nWitnessCacheSize)); +}; + +template void CWallet::WriteWitnessCache(MockWalletDB& walletdb); + class TestWallet : public CWallet { public: TestWallet() : CWallet() { } @@ -28,6 +43,9 @@ public: void DecrementNoteWitnesses() { CWallet::DecrementNoteWitnesses(); } + void WriteWitnessCache(MockWalletDB& walletdb) { + CWallet::WriteWitnessCache(walletdb); + } bool UpdatedNoteData(const CWalletTx& wtxIn, CWalletTx& wtx) { return CWallet::UpdatedNoteData(wtxIn, wtx); } @@ -679,6 +697,50 @@ TEST(wallet_tests, ClearNoteWitnessCache) { EXPECT_FALSE((bool) witnesses[1]); } +TEST(wallet_tests, WriteWitnessCache) { + TestWallet wallet; + MockWalletDB walletdb; + + auto sk = libzcash::SpendingKey::random(); + wallet.AddSpendingKey(sk); + + auto wtx = GetValidReceive(sk, 10, true); + wallet.AddToWallet(wtx, true, NULL); + + EXPECT_CALL(walletdb, TxnBegin()) + .WillOnce(Return(false)) + .WillRepeatedly(Return(true)); + EXPECT_CALL(walletdb, TxnCommit()) + .WillOnce(Return(false)) + .WillRepeatedly(Return(true)); + EXPECT_CALL(walletdb, TxnAbort()) + .Times(4); + + EXPECT_CALL(walletdb, WriteTx(wtx.GetHash(), wtx)) + .WillOnce(Return(false)) + .WillOnce(ThrowLogicError()) + .WillRepeatedly(Return(true)); + EXPECT_CALL(walletdb, WriteWitnessCacheSize(0)) + .WillOnce(Return(false)) + .WillOnce(ThrowLogicError()) + .WillRepeatedly(Return(true)); + + // TxnBegin fails + wallet.WriteWitnessCache(walletdb); + // WriteTx fails + wallet.WriteWitnessCache(walletdb); + // WriteTx throws + wallet.WriteWitnessCache(walletdb); + // WriteWitnessCacheSize fails + wallet.WriteWitnessCache(walletdb); + // WriteWitnessCacheSize throws + wallet.WriteWitnessCache(walletdb); + // TxCommit fails + wallet.WriteWitnessCache(walletdb); + // Everything succeeds + wallet.WriteWitnessCache(walletdb); +} + TEST(wallet_tests, UpdatedNoteData) { TestWallet wallet; diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index fea1c048c..caf5d5818 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -15,7 +15,6 @@ #include "script/script.h" #include "script/sign.h" #include "timedata.h" -#include "util.h" #include "utilmoneystr.h" #include "zcash/Note.hpp" #include "crypter.h" @@ -697,7 +696,8 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex, nWitnessCacheSize += 1; } if (fFileBacked) { - WriteWitnessCache(); + CWalletDB walletdb(strWalletFile); + WriteWitnessCache(walletdb); } } } @@ -718,45 +718,12 @@ void CWallet::DecrementNoteWitnesses() // TODO: If nWitnessCache is zero, we need to regenerate the caches (#1302) assert(nWitnessCacheSize > 0); if (fFileBacked) { - WriteWitnessCache(); + CWalletDB walletdb(strWalletFile); + WriteWitnessCache(walletdb); } } } -void CWallet::WriteWitnessCache() { - CWalletDB walletdb(strWalletFile); - if (!walletdb.TxnBegin()) { - // This needs to be done atomically, so don't do it at all - LogPrintf("WriteWitnessCache(): Couldn't start atomic write\n"); - return; - } - try { - for (std::pair& wtxItem : mapWallet) { - if (!walletdb.WriteTx(wtxItem.first, wtxItem.second)) { - LogPrintf("WriteWitnessCache(): Failed to write CWalletTx, aborting atomic write\n"); - walletdb.TxnAbort(); - return; - } - } - if (!walletdb.WriteWitnessCacheSize(nWitnessCacheSize)) { - LogPrintf("WriteWitnessCache(): Failed to write nWitnessCacheSize, aborting atomic write\n"); - walletdb.TxnAbort(); - return; - } - } catch (const std::exception &exc) { - // Unexpected failure - LogPrintf("WriteWitnessCache(): Unexpected error during atomic write:\n"); - LogPrintf("%s\n", exc.what()); - walletdb.TxnAbort(); - return; - } - if (!walletdb.TxnCommit()) { - // Couldn't commit all to db, but in-memory state is fine - LogPrintf("WriteWitnessCache(): Couldn't commit atomic write\n"); - return; - } -} - bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase) { if (IsCrypted()) diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index c3f1fd63f..84706933e 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -15,6 +15,7 @@ #include "primitives/transaction.h" #include "tinyformat.h" #include "ui_interface.h" +#include "util.h" #include "utilstrencodings.h" #include "validationinterface.h" #include "wallet/crypter.h" @@ -605,7 +606,40 @@ protected: const CBlock* pblock, ZCIncrementalMerkleTree tree); void DecrementNoteWitnesses(); - void WriteWitnessCache(); + + template + void WriteWitnessCache(WalletDB& walletdb) { + if (!walletdb.TxnBegin()) { + // This needs to be done atomically, so don't do it at all + LogPrintf("WriteWitnessCache(): Couldn't start atomic write\n"); + return; + } + try { + for (std::pair& wtxItem : mapWallet) { + if (!walletdb.WriteTx(wtxItem.first, wtxItem.second)) { + LogPrintf("WriteWitnessCache(): Failed to write CWalletTx, aborting atomic write\n"); + walletdb.TxnAbort(); + return; + } + } + if (!walletdb.WriteWitnessCacheSize(nWitnessCacheSize)) { + LogPrintf("WriteWitnessCache(): Failed to write nWitnessCacheSize, aborting atomic write\n"); + walletdb.TxnAbort(); + return; + } + } catch (const std::exception &exc) { + // Unexpected failure + LogPrintf("WriteWitnessCache(): Unexpected error during atomic write:\n"); + LogPrintf("%s\n", exc.what()); + walletdb.TxnAbort(); + return; + } + if (!walletdb.TxnCommit()) { + // Couldn't commit all to db, but in-memory state is fine + LogPrintf("WriteWitnessCache(): Couldn't commit atomic write\n"); + return; + } + } private: template