Add support for spending keys to the encrypted wallet.

src/gtest/test_wallet_zkeys.cpp

@@ -138,3 +138,81 @@ TEST(wallet_zkeys_tests, write_zkey_direct_to_db) {
     ASSERT_EQ(m.nCreateTime, now);
 }
 
+
+
+/**
+ * This test covers methods on CWalletDB to load/save crypted z keys.
+ */
+TEST(wallet_zkeys_tests, write_cryptedzkey_direct_to_db) {
+    ECC_Start();
+
+    SelectParams(CBaseChainParams::TESTNET);
+
+    // Get temporary and unique path for file.
+    // Note: / operator to append paths
+    boost::filesystem::path temp = boost::filesystem::temp_directory_path() /
+            boost::filesystem::unique_path();
+    const std::string path = temp.native();
+
+    bool fFirstRun;
+    CWallet wallet(path);
+    ASSERT_EQ(DB_LOAD_OK, wallet.LoadWallet(fFirstRun));
+
+    // No default CPubKey set
+    ASSERT_TRUE(fFirstRun);
+
+    // wallet should be empty
+    std::set<libzcash::PaymentAddress> addrs;
+    wallet.GetPaymentAddresses(addrs);
+    ASSERT_EQ(0, addrs.size());
+
+    // Add random key to the wallet
+    auto paymentAddress = wallet.GenerateNewZKey();
+
+    // wallet should have one key
+    wallet.GetPaymentAddresses(addrs);
+    ASSERT_EQ(1, addrs.size());
+
+    // encrypt wallet
+    SecureString strWalletPass;
+    strWalletPass.reserve(100);
+    strWalletPass = "hello";
+    ASSERT_TRUE(wallet.EncryptWallet(strWalletPass));
+    
+    // adding a new key will fail as the wallet is locked
+    EXPECT_ANY_THROW(wallet.GenerateNewZKey());
+    
+    // unlock wallet and then add
+    wallet.Unlock(strWalletPass);
+    auto paymentAddress2 = wallet.GenerateNewZKey();
+
+    // Create a new wallet from the existing wallet path
+    CWallet wallet2(path);
+    ASSERT_EQ(DB_LOAD_OK, wallet2.LoadWallet(fFirstRun));
+
+    // Confirm it's not the same as the other wallet
+    ASSERT_TRUE(&wallet != &wallet2);
+    
+    // wallet should have two keys
+    wallet2.GetPaymentAddresses(addrs);
+    ASSERT_EQ(2, addrs.size());
+    
+    // check we have entries for our payment addresses
+    ASSERT_TRUE(addrs.count(paymentAddress.Get()));
+    ASSERT_TRUE(addrs.count(paymentAddress2.Get()));
+
+    // spending key is crypted, so we can't extract valid payment address
+    libzcash::SpendingKey keyOut;
+    wallet2.GetSpendingKey(paymentAddress.Get(), keyOut);
+    ASSERT_FALSE(paymentAddress.Get() == keyOut.address());
+    
+    // unlock wallet to get spending keys and verify payment addresses
+    wallet2.Unlock(strWalletPass);
+
+    wallet2.GetSpendingKey(paymentAddress.Get(), keyOut);
+    ASSERT_EQ(paymentAddress.Get(), keyOut.address());
+    
+    wallet2.GetSpendingKey(paymentAddress2.Get(), keyOut);
+    ASSERT_EQ(paymentAddress2.Get(), keyOut.address());
+}
+

src/test/rpc_wallet_tests.cpp

@@ -1041,4 +1041,58 @@ BOOST_AUTO_TEST_CASE(rpc_z_sendmany_internals)
 }
 
 
+/*
+ * This test covers storing encrypted zkeys in the wallet.
+ */
+BOOST_AUTO_TEST_CASE(rpc_wallet_encrypted_wallet_zkeys)
+{
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+    Value retValue;
+    int n = 100;
+
+    // wallet should currently be empty
+    std::set<libzcash::PaymentAddress> addrs;
+    pwalletMain->GetPaymentAddresses(addrs);
+    BOOST_CHECK(addrs.size()==0);
+
+    // create keys
+    for (int i = 0; i < n; i++) {
+        CallRPC("z_getnewaddress");
+    }
+
+    // Verify we can list the keys imported
+    BOOST_CHECK_NO_THROW(retValue = CallRPC("z_listaddresses"));
+    Array arr = retValue.get_array();
+    BOOST_CHECK(arr.size() == n);
+
+    // Encrypt the wallet (we can''t call RPC encryptwallet as that shuts down node)
+    SecureString strWalletPass;
+    strWalletPass.reserve(100);
+    strWalletPass = "hello";
+    BOOST_CHECK(pwalletMain->EncryptWallet(strWalletPass));
+    
+    // Verify we can still list the keys imported
+    BOOST_CHECK_NO_THROW(retValue = CallRPC("z_listaddresses"));
+    arr = retValue.get_array();
+    BOOST_CHECK(arr.size() == n);
+    
+    // Try to add a new key, but we can't as the wallet is locked
+    BOOST_CHECK_THROW(CallRPC("z_getnewaddress"), runtime_error);
+    
+    // We can't call RPC walletpassphrase as that invokes RPCRunLater which breaks tests.
+    // So we manually unlock.
+    BOOST_CHECK(pwalletMain->Unlock(strWalletPass));
+    
+    // Now add a key
+    BOOST_CHECK_NO_THROW(CallRPC("z_getnewaddress"));
+    
+    // Verify the key has been added
+    BOOST_CHECK_NO_THROW(retValue = CallRPC("z_listaddresses"));
+    arr = retValue.get_array();
+    BOOST_CHECK(arr.size() == n+1);    
+
+    // We can't simulate over RPC the wallet closing and being reloaded
+    // but there are tests for this in gtest.
+}
+
 BOOST_AUTO_TEST_SUITE_END()

src/wallet/rpcwallet.cpp

@@ -2796,6 +2796,8 @@ Value z_getnewaddress(const Array& params, bool fHelp)
 
     LOCK2(cs_main, pwalletMain->cs_wallet);
 
+    EnsureWalletIsUnlocked();
+
     CZCPaymentAddress pubaddr = pwalletMain->GenerateNewZKey();
     std::string result = pubaddr.ToString();
     return result;

src/wallet/wallet.cpp

@@ -18,6 +18,7 @@
 #include "util.h"
 #include "utilmoneystr.h"
 #include "zcash/Note.hpp"
+#include "crypter.h"
 
 #include <assert.h>
 
@@ -169,6 +170,7 @@ bool CWallet::AddKeyPubKey(const CKey& secret, const CPubKey &pubkey)
 bool CWallet::AddCryptedKey(const CPubKey &vchPubKey,
                             const vector<unsigned char> &vchCryptedSecret)
 {
+    
     if (!CCryptoKeyStore::AddCryptedKey(vchPubKey, vchCryptedSecret))
         return false;
     if (!fFileBacked)
@@ -187,6 +189,28 @@ bool CWallet::AddCryptedKey(const CPubKey &vchPubKey,
     return false;
 }
 
+
+bool CWallet::AddCryptedSpendingKey(const libzcash::PaymentAddress &address,
+                                    const std::vector<unsigned char> &vchCryptedSecret)
+{
+    if (!CCryptoKeyStore::AddCryptedSpendingKey(address, vchCryptedSecret))
+        return false;
+    if (!fFileBacked)
+        return true;
+    {
+        LOCK(cs_wallet);
+        if (pwalletdbEncryption)
+            return pwalletdbEncryption->WriteCryptedZKey(address,
+                                                         vchCryptedSecret,
+                                                         mapZKeyMetadata[address]);
+        else
+            return CWalletDB(strWalletFile).WriteCryptedZKey(address,
+                                                             vchCryptedSecret,
+                                                             mapZKeyMetadata[address]);
+    }
+    return false;
+}
+
 bool CWallet::LoadKeyMetadata(const CPubKey &pubkey, const CKeyMetadata &meta)
 {
     AssertLockHeld(cs_wallet); // mapKeyMetadata
@@ -209,6 +233,11 @@ bool CWallet::LoadCryptedKey(const CPubKey &vchPubKey, const std::vector<unsigne
     return CCryptoKeyStore::AddCryptedKey(vchPubKey, vchCryptedSecret);
 }
 
+bool CWallet::LoadCryptedZKey(const libzcash::PaymentAddress &addr, const std::vector<unsigned char> &vchCryptedSecret)
+{
+    return CCryptoKeyStore::AddCryptedSpendingKey(addr, vchCryptedSecret);
+}
+
 bool CWallet::LoadZKey(const libzcash::SpendingKey &key)
 {
     return CCryptoKeyStore::AddSpendingKey(key);

src/wallet/wallet.h

@@ -752,7 +752,10 @@ public:
     bool LoadZKey(const libzcash::SpendingKey &key);
     //! Load spending key metadata (used by LoadWallet)
     bool LoadZKeyMetadata(const libzcash::PaymentAddress &addr, const CKeyMetadata &meta);
-
+    //! Adds an encrypted spending key to the store, without saving it to disk (used by LoadWallet)
+    bool LoadCryptedZKey(const libzcash::PaymentAddress &addr, const std::vector<unsigned char> &vchCryptedSecret);
+    //! 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 std::vector<unsigned char> &vchCryptedSecret);
 
     /** 
      * Increment the next transaction order id

src/wallet/walletdb.cpp

@@ -104,6 +104,25 @@ bool CWalletDB::WriteCryptedKey(const CPubKey& vchPubKey,
     return true;
 }
 
+bool CWalletDB::WriteCryptedZKey(const libzcash::PaymentAddress & addr,
+                                 const std::vector<unsigned char>& vchCryptedSecret,
+                                 const CKeyMetadata &keyMeta)
+{
+    const bool fEraseUnencryptedKey = true;
+    nWalletDBUpdated++;
+
+    if (!Write(std::make_pair(std::string("zkeymeta"), addr), keyMeta))
+        return false;
+
+    if (!Write(std::make_pair(std::string("czkey"), addr), vchCryptedSecret, false))
+        return false;
+    if (fEraseUnencryptedKey)
+    {
+        Erase(std::make_pair(std::string("zkey"), addr));
+    }
+    return true;
+}
+
 bool CWalletDB::WriteMasterKey(unsigned int nID, const CMasterKey& kMasterKey)
 {
     nWalletDBUpdated++;
@@ -348,6 +367,7 @@ public:
     unsigned int nCKeys;
     unsigned int nKeyMeta;
     unsigned int nZKeys;
+    unsigned int nCZKeys;
     unsigned int nZKeyMeta;
     bool fIsEncrypted;
     bool fAnyUnordered;
@@ -355,7 +375,7 @@ public:
     vector<uint256> vWalletUpgrade;
 
     CWalletScanState() {
-        nKeys = nCKeys = nKeyMeta = nZKeys = nZKeyMeta = 0;
+        nKeys = nCKeys = nKeyMeta = nZKeys = nCZKeys = nZKeyMeta = 0;
         fIsEncrypted = false;
         fAnyUnordered = false;
         nFileVersion = 0;
@@ -557,6 +577,21 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
             }
             wss.fIsEncrypted = true;
         }
+        else if (strType == "czkey")
+        {
+            libzcash::PaymentAddress addr;
+            ssKey >> addr;
+            vector<unsigned char> vchCryptedSecret;
+            ssValue >> vchCryptedSecret;
+            wss.nCKeys++;
+
+            if (!pwallet->LoadCryptedZKey(addr, vchCryptedSecret))
+            {
+                strErr = "Error reading wallet database: LoadCryptedZKey failed";
+                return false;
+            }
+            wss.fIsEncrypted = true;
+        }
         else if (strType == "keymeta")
         {
             CPubKey vchPubKey;
@@ -651,7 +686,7 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
 static bool IsKeyType(string strType)
 {
     return (strType== "key" || strType == "wkey" ||
-            strType == "zkey" ||
+            strType == "zkey" || strType == "czkey" ||
             strType == "mkey" || strType == "ckey");
 }
 
@@ -736,9 +771,8 @@ DBErrors CWalletDB::LoadWallet(CWallet* pwallet)
     LogPrintf("Keys: %u plaintext, %u encrypted, %u w/ metadata, %u total\n",
            wss.nKeys, wss.nCKeys, wss.nKeyMeta, wss.nKeys + wss.nCKeys);
 
-    // TODO: Keep track of encrypted ZKeys
+    LogPrintf("ZKeys: %u plaintext, %u encrypted, %u w/metadata, %u total\n",
-    LogPrintf("ZKeys: %u plaintext, -- encrypted, %u w/metadata, %u total\n",
+           wss.nZKeys, wss.nCZKeys, wss.nZKeyMeta, wss.nZKeys + wss.nCZKeys);
-           wss.nZKeys, wss.nZKeyMeta, wss.nZKeys + 0);
 
     // nTimeFirstKey is only reliable if all keys have metadata
     if ((wss.nKeys + wss.nCKeys) != wss.nKeyMeta)

src/wallet/walletdb.h

@@ -135,6 +135,9 @@ public:
 
     /// Write spending key to wallet database, where key is payment address and value is spending key.
     bool WriteZKey(const libzcash::PaymentAddress& addr, const libzcash::SpendingKey& key, const CKeyMetadata &keyMeta);
+    bool WriteCryptedZKey(const libzcash::PaymentAddress & addr,
+                          const std::vector<unsigned char>& vchCryptedSecret,
+                          const CKeyMetadata &keyMeta);
 
 private:
     CWalletDB(const CWalletDB&);