From 8db7e25c3f2248fcca2b57b0c5f59186f29e4e01 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Wed, 24 Aug 2016 15:51:56 +1200
Subject: [PATCH] Create mapping from nullifiers to received notes

This is used in the same way as CTxIn.prevout (e.g. to mark transactions dirty).
---
 src/wallet/gtest/test_wallet.cpp | 26 ++++++++++++++++++++++++++
 src/wallet/wallet.cpp            | 20 ++++++++++++++++++++
 src/wallet/wallet.h              |  2 ++
 3 files changed, 48 insertions(+)

diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp
index 56ec63373..a11d7db50 100644
--- a/src/wallet/gtest/test_wallet.cpp
+++ b/src/wallet/gtest/test_wallet.cpp
@@ -246,3 +246,29 @@ TEST(wallet_tests, nullifier_is_spent) {
     chainActive.SetTip(NULL);
     mapBlockIndex.erase(blockHash);
 }
+
+TEST(wallet_tests, navigate_from_nullifier_to_note) {
+    CWallet wallet;
+
+    auto sk = libzcash::SpendingKey::random();
+    wallet.AddSpendingKey(sk);
+
+    auto wtx = GetValidReceive(sk, 10, true);
+    auto note = GetNote(sk, wtx, 0, 1);
+    auto nullifier = note.nullifier(sk);
+
+    mapNoteData_t noteData;
+    JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+    CNoteData nd {sk.address(), nullifier};
+    noteData[jsoutpt] = nd;
+
+    wtx.SetNoteData(noteData);
+
+    EXPECT_EQ(0, wallet.mapNullifiers.count(nullifier));
+
+    wallet.AddToWallet(wtx, true, NULL);
+    EXPECT_EQ(1, wallet.mapNullifiers.count(nullifier));
+    EXPECT_EQ(wtx.GetTxid(), wallet.mapNullifiers[nullifier].hash);
+    EXPECT_EQ(0, wallet.mapNullifiers[nullifier].js);
+    EXPECT_EQ(1, wallet.mapNullifiers[nullifier].n);
+}
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 7da567eef..55e042ec6 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -722,6 +722,16 @@ void CWallet::MarkDirty()
     }
 }
 
+void CWallet::UpdateNullifierNoteMap(const CWalletTx& wtx)
+{
+    {
+        LOCK(cs_wallet);
+        for (const mapNoteData_t::value_type& item : wtx.mapNoteData) {
+            mapNullifiers[item.second.nullifier] = item.first;
+        }
+    }
+}
+
 bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletDB* pwalletdb)
 {
     uint256 hash = wtxIn.GetTxid();
@@ -730,6 +740,7 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletD
     {
         mapWallet[hash] = wtxIn;
         mapWallet[hash].BindWallet(this);
+        UpdateNullifierNoteMap(mapWallet[hash]);
         AddToSpends(hash);
     }
     else
@@ -739,6 +750,7 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletD
         pair<map<uint256, CWalletTx>::iterator, bool> ret = mapWallet.insert(make_pair(hash, wtxIn));
         CWalletTx& wtx = (*ret.first).second;
         wtx.BindWallet(this);
+        UpdateNullifierNoteMap(wtx);
         bool fInsertedNew = ret.second;
         if (fInsertedNew)
         {
@@ -894,6 +906,14 @@ void CWallet::SyncTransaction(const CTransaction& tx, const CBlock* pblock)
         if (mapWallet.count(txin.prevout.hash))
             mapWallet[txin.prevout.hash].MarkDirty();
     }
+    for (const JSDescription& jsdesc : tx.vjoinsplit) {
+        for (const uint256& nullifier : jsdesc.nullifiers) {
+            if (mapNullifiers.count(nullifier) &&
+                    mapWallet.count(mapNullifiers[nullifier].hash)) {
+                mapWallet[mapNullifiers[nullifier].hash].MarkDirty();
+            }
+        }
+    }
 }
 
 void CWallet::EraseFromWallet(const uint256 &hash)
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 5c94655e7..511bb9c02 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -625,6 +625,7 @@ public:
         fBroadcastTransactions = false;
     }
 
+    std::map<uint256, JSOutPoint> mapNullifiers;
     std::map<uint256, CWalletTx> mapWallet;
 
     int64_t nOrderPosNext;
@@ -727,6 +728,7 @@ public:
     TxItems OrderedTxItems(std::list<CAccountingEntry>& acentries, std::string strAccount = "");
 
     void MarkDirty();
+    void UpdateNullifierNoteMap(const CWalletTx& wtx);
     bool AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletDB* pwalletdb);
     void SyncTransaction(const CTransaction& tx, const CBlock* pblock);
     bool AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pblock, bool fUpdate);