From ad041fceec7a86dcfd405be2bd9b18d268839643 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Thu, 29 Sep 2016 11:34:59 +1300
Subject: [PATCH] Cache note decryptors in encrypted keystore

---
 src/gtest/test_keystore.cpp | 15 +++++++++++++++
 src/wallet/crypter.cpp      |  9 ++++++---
 src/wallet/crypter.h        |  4 +++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/gtest/test_keystore.cpp b/src/gtest/test_keystore.cpp
index a82cd711a..26fde42d9 100644
--- a/src/gtest/test_keystore.cpp
+++ b/src/gtest/test_keystore.cpp
@@ -56,19 +56,26 @@ TEST(keystore_tests, store_and_retrieve_spending_key_in_encrypted_store) {
     uint256 r {GetRandHash()};
     CKeyingMaterial vMasterKey (r.begin(), r.end());
     libzcash::SpendingKey keyOut;
+    ZCNoteDecryption decOut;
     std::set<libzcash::PaymentAddress> addrs;
 
     // 1) Test adding a key to an unencrypted key store, then encrypting it
     auto sk = libzcash::SpendingKey::random();
     auto addr = sk.address();
+    EXPECT_FALSE(keyStore.GetNoteDecryptor(addr, decOut));
+
     keyStore.AddSpendingKey(sk);
     ASSERT_TRUE(keyStore.HaveSpendingKey(addr));
     ASSERT_TRUE(keyStore.GetSpendingKey(addr, keyOut));
     ASSERT_EQ(sk, keyOut);
+    EXPECT_TRUE(keyStore.GetNoteDecryptor(addr, decOut));
+    EXPECT_EQ(ZCNoteDecryption(sk.viewing_key()), decOut);
 
     ASSERT_TRUE(keyStore.EncryptKeys(vMasterKey));
     ASSERT_TRUE(keyStore.HaveSpendingKey(addr));
     ASSERT_FALSE(keyStore.GetSpendingKey(addr, keyOut));
+    EXPECT_TRUE(keyStore.GetNoteDecryptor(addr, decOut));
+    EXPECT_EQ(ZCNoteDecryption(sk.viewing_key()), decOut);
 
     // Unlocking with a random key should fail
     uint256 r2 {GetRandHash()};
@@ -92,18 +99,26 @@ TEST(keystore_tests, store_and_retrieve_spending_key_in_encrypted_store) {
     // 2) Test adding a spending key to an already-encrypted key store
     auto sk2 = libzcash::SpendingKey::random();
     auto addr2 = sk2.address();
+    EXPECT_FALSE(keyStore.GetNoteDecryptor(addr2, decOut));
+
     keyStore.AddSpendingKey(sk2);
     ASSERT_TRUE(keyStore.HaveSpendingKey(addr2));
     ASSERT_TRUE(keyStore.GetSpendingKey(addr2, keyOut));
     ASSERT_EQ(sk2, keyOut);
+    EXPECT_TRUE(keyStore.GetNoteDecryptor(addr2, decOut));
+    EXPECT_EQ(ZCNoteDecryption(sk2.viewing_key()), decOut);
 
     ASSERT_TRUE(keyStore.Lock());
     ASSERT_TRUE(keyStore.HaveSpendingKey(addr2));
     ASSERT_FALSE(keyStore.GetSpendingKey(addr2, keyOut));
+    EXPECT_TRUE(keyStore.GetNoteDecryptor(addr2, decOut));
+    EXPECT_EQ(ZCNoteDecryption(sk2.viewing_key()), decOut);
 
     ASSERT_TRUE(keyStore.Unlock(vMasterKey));
     ASSERT_TRUE(keyStore.GetSpendingKey(addr2, keyOut));
     ASSERT_EQ(sk2, keyOut);
+    EXPECT_TRUE(keyStore.GetNoteDecryptor(addr2, decOut));
+    EXPECT_EQ(ZCNoteDecryption(sk2.viewing_key()), decOut);
 
     keyStore.GetPaymentAddresses(addrs);
     ASSERT_EQ(2, addrs.size());
diff --git a/src/wallet/crypter.cpp b/src/wallet/crypter.cpp
index 5572ebe2f..5a066d75d 100644
--- a/src/wallet/crypter.cpp
+++ b/src/wallet/crypter.cpp
@@ -318,13 +318,15 @@ bool CCryptoKeyStore::AddSpendingKey(const libzcash::SpendingKey &sk)
         if (!EncryptSecret(vMasterKey, vchSecret, address.GetHash(), vchCryptedSecret))
             return false;
 
-        if (!AddCryptedSpendingKey(address, vchCryptedSecret))
+        if (!AddCryptedSpendingKey(address, sk.viewing_key(), vchCryptedSecret))
             return false;
     }
     return true;
 }
 
-bool CCryptoKeyStore::AddCryptedSpendingKey(const libzcash::PaymentAddress &address, const std::vector<unsigned char> &vchCryptedSecret)
+bool CCryptoKeyStore::AddCryptedSpendingKey(const libzcash::PaymentAddress &address,
+                                            const libzcash::ViewingKey &vk,
+                                            const std::vector<unsigned char> &vchCryptedSecret)
 {
     {
         LOCK(cs_KeyStore);
@@ -332,6 +334,7 @@ bool CCryptoKeyStore::AddCryptedSpendingKey(const libzcash::PaymentAddress &addr
             return false;
 
         mapCryptedSpendingKeys[address] = vchCryptedSecret;
+        mapNoteDecryptors.insert(std::make_pair(address, ZCNoteDecryption(vk)));
     }
     return true;
 }
@@ -383,7 +386,7 @@ bool CCryptoKeyStore::EncryptKeys(CKeyingMaterial& vMasterKeyIn)
             std::vector<unsigned char> vchCryptedSecret;
             if (!EncryptSecret(vMasterKeyIn, vchSecret, address.GetHash(), vchCryptedSecret))
                 return false;
-            if (!AddCryptedSpendingKey(address, vchCryptedSecret))
+            if (!AddCryptedSpendingKey(address, sk.viewing_key(), vchCryptedSecret))
                 return false;
         }
         mapSpendingKeys.clear();
diff --git a/src/wallet/crypter.h b/src/wallet/crypter.h
index e95a841a4..b310b77b0 100644
--- a/src/wallet/crypter.h
+++ b/src/wallet/crypter.h
@@ -200,7 +200,9 @@ public:
             mi++;
         }
     }
-    virtual bool AddCryptedSpendingKey(const libzcash::PaymentAddress &address, const std::vector<unsigned char> &vchCryptedSecret);
+    virtual bool AddCryptedSpendingKey(const libzcash::PaymentAddress &address,
+                                       const libzcash::ViewingKey &vk,
+                                       const std::vector<unsigned char> &vchCryptedSecret);
     bool AddSpendingKey(const libzcash::SpendingKey &sk);
     bool HaveSpendingKey(const libzcash::PaymentAddress &address) const
     {