Implement viewing key storage in the wallet

src/wallet/gtest/test_wallet_zkeys.cpp

@@ -66,6 +66,53 @@ TEST(wallet_zkeys_tests, store_and_load_zkeys) {
     ASSERT_EQ(m.nCreateTime, now);
 }
 
+/**
+ * This test covers methods on CWallet
+ * AddViewingKey()
+ * RemoveViewingKey()
+ * LoadViewingKey()
+ */
+TEST(wallet_zkeys_tests, StoreAndLoadViewingKeys) {
+    SelectParams(CBaseChainParams::MAIN);
+
+    CWallet wallet;
+
+    // wallet should be empty
+    std::set<libzcash::PaymentAddress> addrs;
+    wallet.GetPaymentAddresses(addrs);
+    ASSERT_EQ(0, addrs.size());
+
+    // manually add new viewing key to wallet
+    auto sk = libzcash::SpendingKey::random();
+    auto vk = sk.viewing_key();
+    ASSERT_TRUE(wallet.AddViewingKey(vk));
+
+    // verify wallet did add it
+    auto addr = sk.address();
+    ASSERT_TRUE(wallet.HaveViewingKey(addr));
+    // and that we don't have the corresponding spending key
+    ASSERT_FALSE(wallet.HaveSpendingKey(addr));
+
+    // verify viewing key stored correctly
+    libzcash::ViewingKey vkOut;
+    wallet.GetViewingKey(addr, vkOut);
+    ASSERT_EQ(vk, vkOut);
+
+    // Load a second viewing key into the wallet
+    auto sk2 = libzcash::SpendingKey::random();
+    ASSERT_TRUE(wallet.LoadViewingKey(sk2.viewing_key()));
+
+    // verify wallet did add it
+    auto addr2 = sk2.address();
+    ASSERT_TRUE(wallet.HaveViewingKey(addr2));
+    ASSERT_FALSE(wallet.HaveSpendingKey(addr2));
+
+    // Remove the first viewing key
+    ASSERT_TRUE(wallet.RemoveViewingKey(vk));
+    ASSERT_FALSE(wallet.HaveViewingKey(addr));
+    ASSERT_TRUE(wallet.HaveViewingKey(addr2));
+}
+
 /**
  * This test covers methods on CWalletDB
  * WriteZKey()
@@ -138,6 +185,50 @@ TEST(wallet_zkeys_tests, write_zkey_direct_to_db) {
     ASSERT_EQ(m.nCreateTime, now);
 }
 
+/**
+ * This test covers methods on CWalletDB
+ * WriteViewingKey()
+ */
+TEST(wallet_zkeys_tests, WriteViewingKeyDirectToDB) {
+    SelectParams(CBaseChainParams::TESTNET);
+
+    // Get temporary and unique path for file.
+    // Note: / operator to append paths
+    boost::filesystem::path pathTemp = boost::filesystem::temp_directory_path() / boost::filesystem::unique_path();
+    boost::filesystem::create_directories(pathTemp);
+    mapArgs["-datadir"] = pathTemp.string();
+
+    bool fFirstRun;
+    CWallet wallet("wallet-vkey.dat");
+    ASSERT_EQ(DB_LOAD_OK, wallet.LoadWallet(fFirstRun));
+
+    // No default CPubKey set
+    ASSERT_TRUE(fFirstRun);
+
+    // create random viewing key and add it to database directly, bypassing wallet
+    auto sk = libzcash::SpendingKey::random();
+    auto vk = sk.viewing_key();
+    auto addr = sk.address();
+    int64_t now = GetTime();
+    CKeyMetadata meta(now);
+    CWalletDB db("wallet-vkey.dat");
+    db.WriteViewingKey(vk);
+
+    // wallet should not be aware of viewing key
+    ASSERT_FALSE(wallet.HaveViewingKey(addr));
+
+    // load the wallet again
+    ASSERT_EQ(DB_LOAD_OK, wallet.LoadWallet(fFirstRun));
+
+    // wallet can now see the viewing key
+    ASSERT_TRUE(wallet.HaveViewingKey(addr));
+
+    // check key is the same
+    libzcash::ViewingKey vkOut;
+    wallet.GetViewingKey(addr, vkOut);
+    ASSERT_EQ(vk, vkOut);
+}
+
 
 
 /**

src/wallet/wallet.cpp

@@ -107,6 +107,10 @@ bool CWallet::AddZKey(const libzcash::SpendingKey &key)
     if (!CCryptoKeyStore::AddSpendingKey(key))
         return false;
 
+    // check if we need to remove from viewing keys
+    if (HaveViewingKey(addr))
+        RemoveViewingKey(key.viewing_key());
+
     if (!fFileBacked)
         return true;
 
@@ -246,6 +250,33 @@ bool CWallet::LoadZKey(const libzcash::SpendingKey &key)
     return CCryptoKeyStore::AddSpendingKey(key);
 }
 
+bool CWallet::AddViewingKey(const libzcash::ViewingKey &vk)
+{
+    if (!CCryptoKeyStore::AddViewingKey(vk))
+        return false;
+    nTimeFirstKey = 1; // No birthday information for viewing keys.
+    if (!fFileBacked)
+        return true;
+    return CWalletDB(strWalletFile).WriteViewingKey(vk);
+}
+
+bool CWallet::RemoveViewingKey(const libzcash::ViewingKey &vk)
+{
+    AssertLockHeld(cs_wallet);
+    if (!CCryptoKeyStore::RemoveViewingKey(vk))
+        return false;
+    if (fFileBacked)
+        if (!CWalletDB(strWalletFile).EraseViewingKey(vk))
+            return false;
+
+    return true;
+}
+
+bool CWallet::LoadViewingKey(const libzcash::ViewingKey &vk)
+{
+    return CCryptoKeyStore::AddViewingKey(vk);
+}
+
 bool CWallet::AddCScript(const CScript& redeemScript)
 {
     if (!CCryptoKeyStore::AddCScript(redeemScript))

src/wallet/wallet.h

@@ -954,6 +954,12 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     //! Adds an encrypted spending key to the store, and saves it to disk (virtual method, declared in crypter.h)
     bool AddCryptedSpendingKey(const libzcash::PaymentAddress &address, const libzcash::ReceivingKey &rk, const std::vector<unsigned char> &vchCryptedSecret);
 
+    //! Adds a viewing key to the store, and saves it to disk.
+    bool AddViewingKey(const libzcash::ViewingKey &vk);
+    bool RemoveViewingKey(const libzcash::ViewingKey &vk);
+    //! Adds a viewing key to the store, without saving it to disk (used by LoadWallet)
+    bool LoadViewingKey(const libzcash::ViewingKey &dest);
+
     /** 
      * Increment the next transaction order id
      * @return next transaction order id

src/wallet/walletdb.cpp

@@ -142,6 +142,18 @@ bool CWalletDB::WriteZKey(const libzcash::PaymentAddress& addr, const libzcash::
     return Write(std::make_pair(std::string("zkey"), addr), key, false);
 }
 
+bool CWalletDB::WriteViewingKey(const libzcash::ViewingKey &vk)
+{
+    nWalletDBUpdated++;
+    return Write(std::make_pair(std::string("vkey"), vk), '1');
+}
+
+bool CWalletDB::EraseViewingKey(const libzcash::ViewingKey &vk)
+{
+    nWalletDBUpdated++;
+    return Erase(std::make_pair(std::string("vkey"), vk));
+}
+
 bool CWalletDB::WriteCScript(const uint160& hash, const CScript& redeemScript)
 {
     nWalletDBUpdated++;
@@ -471,6 +483,19 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
             // so set the wallet birthday to the beginning of time.
             pwallet->nTimeFirstKey = 1;
         }
+        else if (strType == "vkey")
+        {
+            libzcash::ViewingKey vk;
+            ssKey >> vk;
+            char fYes;
+            ssValue >> fYes;
+            if (fYes == '1')
+                pwallet->LoadViewingKey(vk);
+
+            // Viewing keys have no birthday information for now,
+            // so set the wallet birthday to the beginning of time.
+            pwallet->nTimeFirstKey = 1;
+        }
         else if (strType == "zkey")
         {
             libzcash::PaymentAddress addr;
@@ -694,6 +719,7 @@ static bool IsKeyType(string strType)
 {
     return (strType== "key" || strType == "wkey" ||
             strType == "zkey" || strType == "czkey" ||
+            strType == "vkey" ||
             strType == "mkey" || strType == "ckey");
 }
 

src/wallet/walletdb.h

@@ -140,6 +140,9 @@ class CWalletDB : public CDB
                           const std::vector<unsigned char>& vchCryptedSecret,
                           const CKeyMetadata &keyMeta);
 
+    bool WriteViewingKey(const libzcash::ViewingKey &vk);
+    bool EraseViewingKey(const libzcash::ViewingKey &vk);
+
 private:
     CWalletDB(const CWalletDB&);
     void operator=(const CWalletDB&);