From a4ecd0fa72684f9cf44145ad6a9d8929d835e4f8 Mon Sep 17 00:00:00 2001 From: Eirik Ogilvie-Wigley Date: Mon, 1 Oct 2018 09:34:25 -0600 Subject: [PATCH] Add newly discovered sapling addresses to the wallet --- src/keystore.cpp | 18 +++++++++++++++--- src/keystore.h | 6 ++++++ src/wallet/gtest/test_wallet.cpp | 18 +++++++++--------- src/wallet/wallet.cpp | 18 +++++++++++++++--- src/wallet/wallet.h | 2 +- 5 files changed, 46 insertions(+), 16 deletions(-) diff --git a/src/keystore.cpp b/src/keystore.cpp index be5e0b82f..50766518f 100644 --- a/src/keystore.cpp +++ b/src/keystore.cpp @@ -156,9 +156,21 @@ bool CBasicKeyStore::AddSaplingFullViewingKey( LOCK(cs_SpendingKeyStore); auto ivk = fvk.in_viewing_key(); mapSaplingFullViewingKeys[ivk] = fvk; - - // Add defaultAddr -> SaplingIncomingViewing to SaplingIncomingViewingKeyMap - mapSaplingIncomingViewingKeys[defaultAddr] = ivk; + + return AddSaplingIncomingViewingKey(ivk, defaultAddr); +} + +// This function updates the wallet's internal address->ivk map. +// If we add an address that is already in the map, the map will +// remain unchanged as each address only has one ivk. +bool CBasicKeyStore::AddSaplingIncomingViewingKey( + const libzcash::SaplingIncomingViewingKey &ivk, + const libzcash::SaplingPaymentAddress &addr) +{ + LOCK(cs_SpendingKeyStore); + + // Add addr -> SaplingIncomingViewing to SaplingIncomingViewingKeyMap + mapSaplingIncomingViewingKeys[addr] = ivk; return true; } diff --git a/src/keystore.h b/src/keystore.h index 5b064f27e..2cf34ea2c 100644 --- a/src/keystore.h +++ b/src/keystore.h @@ -82,6 +82,9 @@ public: libzcash::SaplingFullViewingKey& fvkOut) const =0; //! Sapling incoming viewing keys + virtual bool AddSaplingIncomingViewingKey( + const libzcash::SaplingIncomingViewingKey &ivk, + const libzcash::SaplingPaymentAddress &addr) =0; virtual bool HaveSaplingIncomingViewingKey(const libzcash::SaplingPaymentAddress &addr) const =0; virtual bool GetSaplingIncomingViewingKey( const libzcash::SaplingPaymentAddress &addr, @@ -269,6 +272,9 @@ public: const libzcash::SaplingIncomingViewingKey &ivk, libzcash::SaplingFullViewingKey& fvkOut) const; + virtual bool AddSaplingIncomingViewingKey( + const libzcash::SaplingIncomingViewingKey &ivk, + const libzcash::SaplingPaymentAddress &addr); virtual bool HaveSaplingIncomingViewingKey(const libzcash::SaplingPaymentAddress &addr) const; virtual bool GetSaplingIncomingViewingKey( const libzcash::SaplingPaymentAddress &addr, diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp index fa066061d..f4b9ddc55 100644 --- a/src/wallet/gtest/test_wallet.cpp +++ b/src/wallet/gtest/test_wallet.cpp @@ -512,13 +512,13 @@ TEST(WalletTests, FindMySaplingNotes) { // No Sapling notes can be found in tx which does not belong to the wallet CWalletTx wtx {&wallet, tx}; ASSERT_FALSE(wallet.HaveSaplingSpendingKey(fvk)); - auto noteMap = wallet.FindMySaplingNotes(wtx); + auto noteMap = wallet.FindMySaplingNotes(wtx).first; EXPECT_EQ(0, noteMap.size()); // Add spending key to wallet, so Sapling notes can be found ASSERT_TRUE(wallet.AddSaplingZKey(sk, pk)); ASSERT_TRUE(wallet.HaveSaplingSpendingKey(fvk)); - noteMap = wallet.FindMySaplingNotes(wtx); + noteMap = wallet.FindMySaplingNotes(wtx).first; EXPECT_EQ(2, noteMap.size()); // Revert to default @@ -664,7 +664,7 @@ TEST(WalletTests, GetConflictedSaplingNotes) { EXPECT_EQ(0, chainActive.Height()); // Simulate SyncTransaction which calls AddToWalletIfInvolvingMe - auto saplingNoteData = wallet.FindMySaplingNotes(wtx); + auto saplingNoteData = wallet.FindMySaplingNotes(wtx).first; ASSERT_TRUE(saplingNoteData.size() > 0); wtx.SetSaplingNoteData(saplingNoteData); wtx.SetMerkleBranch(block); @@ -938,7 +938,7 @@ TEST(WalletTests, NavigateFromSaplingNullifierToNote) { // Simulate SyncTransaction which calls AddToWalletIfInvolvingMe wtx.SetMerkleBranch(block); - auto saplingNoteData = wallet.FindMySaplingNotes(wtx); + auto saplingNoteData = wallet.FindMySaplingNotes(wtx).first; ASSERT_TRUE(saplingNoteData.size() > 0); wtx.SetSaplingNoteData(saplingNoteData); wallet.AddToWallet(wtx, true, NULL); @@ -1064,7 +1064,7 @@ TEST(WalletTests, SpentSaplingNoteIsFromMe) { EXPECT_TRUE(chainActive.Contains(&fakeIndex)); EXPECT_EQ(0, chainActive.Height()); - auto saplingNoteData = wallet.FindMySaplingNotes(wtx); + auto saplingNoteData = wallet.FindMySaplingNotes(wtx).first; ASSERT_TRUE(saplingNoteData.size() > 0); wtx.SetSaplingNoteData(saplingNoteData); wtx.SetMerkleBranch(block); @@ -1141,7 +1141,7 @@ TEST(WalletTests, SpentSaplingNoteIsFromMe) { EXPECT_TRUE(chainActive.Contains(&fakeIndex2)); EXPECT_EQ(1, chainActive.Height()); - auto saplingNoteData2 = wallet.FindMySaplingNotes(wtx2); + auto saplingNoteData2 = wallet.FindMySaplingNotes(wtx2).first; ASSERT_TRUE(saplingNoteData2.size() > 0); wtx2.SetSaplingNoteData(saplingNoteData2); wtx2.SetMerkleBranch(block2); @@ -1769,7 +1769,7 @@ TEST(WalletTests, UpdatedSaplingNoteData) { EXPECT_EQ(0, chainActive.Height()); // Simulate SyncTransaction which calls AddToWalletIfInvolvingMe - auto saplingNoteData = wallet.FindMySaplingNotes(wtx); + auto saplingNoteData = wallet.FindMySaplingNotes(wtx).first; ASSERT_TRUE(saplingNoteData.size() == 1); // wallet only has key for change output wtx.SetSaplingNoteData(saplingNoteData); wtx.SetMerkleBranch(block); @@ -1787,7 +1787,7 @@ TEST(WalletTests, UpdatedSaplingNoteData) { ASSERT_TRUE(wallet.AddSaplingZKey(sk2, pk2)); ASSERT_TRUE(wallet.HaveSaplingSpendingKey(fvk2)); CWalletTx wtx2 = wtx; - auto saplingNoteData2 = wallet.FindMySaplingNotes(wtx2); + auto saplingNoteData2 = wallet.FindMySaplingNotes(wtx2).first; ASSERT_TRUE(saplingNoteData2.size() == 2); wtx2.SetSaplingNoteData(saplingNoteData2); @@ -1923,7 +1923,7 @@ TEST(WalletTests, MarkAffectedSaplingTransactionsDirty) { EXPECT_EQ(0, chainActive.Height()); // Simulate SyncTransaction which calls AddToWalletIfInvolvingMe - auto saplingNoteData = wallet.FindMySaplingNotes(wtx); + auto saplingNoteData = wallet.FindMySaplingNotes(wtx).first; ASSERT_TRUE(saplingNoteData.size() > 0); wtx.SetSaplingNoteData(saplingNoteData); wtx.SetMerkleBranch(block); diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 2093fadb9..72872f5e1 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -1483,7 +1483,14 @@ bool CWallet::AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pbl bool fExisted = mapWallet.count(tx.GetHash()) != 0; if (fExisted && !fUpdate) return false; auto sproutNoteData = FindMySproutNotes(tx); - auto saplingNoteData = FindMySaplingNotes(tx); + auto saplingNoteDataAndAddressesToAdd = FindMySaplingNotes(tx); + auto saplingNoteData = saplingNoteDataAndAddressesToAdd.first; + auto addressesToAdd = saplingNoteDataAndAddressesToAdd.second; + for (const auto &addressToAdd : addressesToAdd) { + if (!AddSaplingIncomingViewingKey(addressToAdd.second, addressToAdd.first)) { + return false; + } + } if (fExisted || IsMine(tx) || IsFromMe(tx) || sproutNoteData.size() > 0 || saplingNoteData.size() > 0) { CWalletTx wtx(this,tx); @@ -1644,12 +1651,13 @@ mapSproutNoteData_t CWallet::FindMySproutNotes(const CTransaction &tx) const * the result of FindMySaplingNotes (for the addresses available at the time) will * already have been cached in CWalletTx.mapSaplingNoteData. */ -mapSaplingNoteData_t CWallet::FindMySaplingNotes(const CTransaction &tx) const +std::pair CWallet::FindMySaplingNotes(const CTransaction &tx) const { LOCK(cs_SpendingKeyStore); uint256 hash = tx.GetHash(); mapSaplingNoteData_t noteData; + SaplingIncomingViewingKeyMap viewingKeysToAdd; // Protocol Spec: 4.19 Block Chain Scanning (Sapling) for (uint32_t i = 0; i < tx.vShieldedOutput.size(); ++i) { @@ -1660,6 +1668,10 @@ mapSaplingNoteData_t CWallet::FindMySaplingNotes(const CTransaction &tx) const if (!result) { continue; } + auto address = ivk.address(result.get().d); + if (address && mapSaplingIncomingViewingKeys.count(address.get()) == 0) { + viewingKeysToAdd[address.get()] = ivk; + } // We don't cache the nullifier here as computing it requires knowledge of the note position // in the commitment tree, which can only be determined when the transaction has been mined. SaplingOutPoint op {hash, i}; @@ -1670,7 +1682,7 @@ mapSaplingNoteData_t CWallet::FindMySaplingNotes(const CTransaction &tx) const } } - return noteData; + return std::make_pair(noteData, viewingKeysToAdd); } bool CWallet::IsSproutNullifierFromMe(const uint256& nullifier) const diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index 1deb9b3aa..770650d48 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -1132,7 +1132,7 @@ public: const uint256& hSig, uint8_t n) const; mapSproutNoteData_t FindMySproutNotes(const CTransaction& tx) const; - mapSaplingNoteData_t FindMySaplingNotes(const CTransaction& tx) const; + std::pair FindMySaplingNotes(const CTransaction& tx) const; bool IsSproutNullifierFromMe(const uint256& nullifier) const; bool IsSaplingNullifierFromMe(const uint256& nullifier) const;