From c3a7307a69dcf94539e5354d0a1d723f6dd59e05 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Wed, 24 Aug 2016 15:51:09 +1200
Subject: [PATCH] Store mapping between notes and PaymentAddresses in CWalletTx

---
 src/wallet/gtest/test_wallet.cpp | 29 +++++++++++++++++++++++++++++
 src/wallet/wallet.cpp            | 29 ++++++++++++++++++++++++++++-
 src/wallet/wallet.h              |  5 +++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp
index bbb0501bb..8db093e7f 100644
--- a/src/wallet/gtest/test_wallet.cpp
+++ b/src/wallet/gtest/test_wallet.cpp
@@ -78,6 +78,35 @@ libzcash::Note GetNote(const libzcash::SpendingKey& sk,
     return note_pt.note(sk.address());
 }
 
+TEST(wallet_tests, set_note_addrs_in_cwallettx) {
+    auto sk = libzcash::SpendingKey::random();
+    auto wtx = GetValidReceive(sk, 10, true);
+    auto note = GetNote(sk, wtx, 0, 1);
+    auto nullifier = note.nullifier(sk);
+    EXPECT_EQ(0, wtx.mapNoteData.size());
+
+    mapNoteData_t noteData;
+    JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+    CNoteData nd {sk.address(), nullifier};
+    noteData[jsoutpt] = nd;
+
+    wtx.SetNoteData(noteData);
+    EXPECT_EQ(noteData, wtx.mapNoteData);
+}
+
+TEST(wallet_tests, set_invalid_note_addrs_in_cwallettx) {
+    CWalletTx wtx;
+    EXPECT_EQ(0, wtx.mapNoteData.size());
+
+    mapNoteData_t noteData;
+    auto sk = libzcash::SpendingKey::random();
+    JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+    CNoteData nd {sk.address(), uint256()};
+    noteData[jsoutpt] = nd;
+
+    EXPECT_THROW(wtx.SetNoteData(noteData), std::runtime_error);
+}
+
 TEST(wallet_tests, find_note_in_tx) {
     CWallet wallet;
 
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index cb977ec8e..0d287dac7 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -481,6 +481,8 @@ void CWallet::SyncMetaData(pair<TxSpends::iterator, TxSpends::iterator> range)
         CWalletTx* copyTo = &mapWallet[hash];
         if (copyFrom == copyTo) continue;
         copyTo->mapValue = copyFrom->mapValue;
+        // mapNoteData not copied on purpose
+        // (it is always set correctly for each CWalletTx)
         copyTo->vOrderForm = copyFrom->vOrderForm;
         // fTimeReceivedIsTxTime not copied on purpose
         // nTimeReceived not copied on purpose
@@ -758,6 +760,11 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletD
                 wtx.nIndex = wtxIn.nIndex;
                 fUpdated = true;
             }
+            if (!wtxIn.mapNoteData.empty() && wtxIn.mapNoteData != wtx.mapNoteData)
+            {
+                wtx.mapNoteData = wtxIn.mapNoteData;
+                fUpdated = true;
+            }
             if (wtxIn.fFromMe && wtxIn.fFromMe != wtx.fFromMe)
             {
                 wtx.fFromMe = wtxIn.fFromMe;
@@ -803,10 +810,14 @@ bool CWallet::AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pbl
         AssertLockHeld(cs_wallet);
         bool fExisted = mapWallet.count(tx.GetTxid()) != 0;
         if (fExisted && !fUpdate) return false;
-        if (fExisted || IsMine(tx) || IsFromMe(tx))
+        auto noteData = FindMyNotes(tx);
+        if (fExisted || IsMine(tx) || IsFromMe(tx) || noteData.size() > 0)
         {
             CWalletTx wtx(this,tx);
 
+            if (noteData.size() > 0)
+                wtx.SetNoteData(noteData);
+
             // Get merkle branch if transaction was found in a block
             if (pblock)
                 wtx.SetMerkleBranch(*pblock);
@@ -1007,6 +1018,22 @@ CAmount CWallet::GetChange(const CTransaction& tx) const
     return nChange;
 }
 
+void CWalletTx::SetNoteData(mapNoteData_t &noteData)
+{
+    mapNoteData.clear();
+    for (const std::pair<JSOutPoint, CNoteData> nd : noteData) {
+        if (nd.first.js < vjoinsplit.size() &&
+                nd.first.n < vjoinsplit[nd.first.js].ciphertexts.size()) {
+            // Store the address and nullifier for the Note
+            mapNoteData[nd.first] = nd.second;
+        } else {
+            // If FindMyNotes() was used to obtain noteData,
+            // this should never happen
+            throw std::runtime_error("CWalletTc::SetNoteData(): Invalid note");
+        }
+    }
+}
+
 int64_t CWalletTx::GetTxTime() const
 {
     int64_t n = nTimeSmart;
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 4bbbe2aa0..6b118b9c6 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -298,6 +298,7 @@ private:
 
 public:
     mapValue_t mapValue;
+    mapNoteData_t mapNoteData;
     std::vector<std::pair<std::string, std::string> > vOrderForm;
     unsigned int fTimeReceivedIsTxTime;
     unsigned int nTimeReceived; //! time received by this node
@@ -350,6 +351,7 @@ public:
     {
         pwallet = pwalletIn;
         mapValue.clear();
+        mapNoteData.clear();
         vOrderForm.clear();
         fTimeReceivedIsTxTime = false;
         nTimeReceived = 0;
@@ -399,6 +401,7 @@ public:
         std::vector<CMerkleTx> vUnused; //! Used to be vtxPrev
         READWRITE(vUnused);
         READWRITE(mapValue);
+        READWRITE(mapNoteData);
         READWRITE(vOrderForm);
         READWRITE(fTimeReceivedIsTxTime);
         READWRITE(nTimeReceived);
@@ -440,6 +443,8 @@ public:
         MarkDirty();
     }
 
+    void SetNoteData(mapNoteData_t &noteData);
+
     //! filter decides which addresses will count towards the debit
     CAmount GetDebit(const isminefilter& filter) const;
     CAmount GetCredit(const isminefilter& filter) const;