Add Sapling Add/Have/Get to keystore

src/keystore.cpp

@@ -93,6 +93,16 @@ bool CBasicKeyStore::AddSpendingKey(const libzcash::SproutSpendingKey &sk)
     return true;
 }
 
+//! Sapling 
+bool CBasicKeyStore::AddSaplingSpendingKey(const libzcash::SaplingSpendingKey &sk)
+{
+    LOCK(cs_SpendingKeyStore);
+    auto fvk = sk.full_viewing_key();
+    mapSaplingSpendingKeys[fvk] = sk;
+    //! TODO: Note decryptors for Sapling
+    return true;
+}
+
 bool CBasicKeyStore::AddViewingKey(const libzcash::SproutViewingKey &vk)
 {
     LOCK(cs_SpendingKeyStore);

src/keystore.h

@@ -55,6 +55,13 @@ public:
     virtual bool HaveSpendingKey(const libzcash::SproutPaymentAddress &address) const =0;
     virtual bool GetSpendingKey(const libzcash::SproutPaymentAddress &address, libzcash::SproutSpendingKey& skOut) const =0;
     virtual void GetPaymentAddresses(std::set<libzcash::SproutPaymentAddress> &setAddress) const =0;
+    
+    //! Add a Sapling spending key to the store.
+    virtual bool AddSaplingSpendingKey(const libzcash::SaplingSpendingKey &sk) =0;
+    
+    //! Check whether a Sapling spending key corresponding to a given Sapling viewing key is present in the store.
+    virtual bool HaveSaplingSpendingKey(const libzcash::SaplingFullViewingKey &fvk) const =0;
+    virtual bool GetSaplingSpendingKey(const libzcash::SaplingFullViewingKey &fvk, libzcash::SaplingSpendingKey& skOut) const =0;
 
     //! Support for viewing keys
     virtual bool AddViewingKey(const libzcash::SproutViewingKey &vk) =0;
@@ -70,6 +77,9 @@ typedef std::map<libzcash::SproutPaymentAddress, libzcash::SproutSpendingKey> Sp
 typedef std::map<libzcash::SproutPaymentAddress, libzcash::SproutViewingKey> ViewingKeyMap;
 typedef std::map<libzcash::SproutPaymentAddress, ZCNoteDecryption> NoteDecryptorMap;
 
+typedef std::map<libzcash::SaplingFullViewingKey, libzcash::SaplingSpendingKey> SaplingSpendingKeyMap;
+typedef std::map<libzcash::SaplingIncomingViewingKey, libzcash::SaplingFullViewingKey> SaplingFullViewingKeyMap;
+
 /** Basic key store, that keeps keys in an address->secret map */
 class CBasicKeyStore : public CKeyStore
 {
@@ -80,6 +90,8 @@ protected:
     SpendingKeyMap mapSpendingKeys;
     ViewingKeyMap mapViewingKeys;
     NoteDecryptorMap mapNoteDecryptors;
+    
+    SaplingSpendingKeyMap mapSaplingSpendingKeys;
 
 public:
     bool AddKeyPubKey(const CKey& key, const CPubKey &pubkey);
@@ -182,6 +194,32 @@ public:
             }
         }
     }
+    
+    //! Sapling 
+    bool AddSaplingSpendingKey(const libzcash::SaplingSpendingKey &sk);
+    bool HaveSaplingSpendingKey(const libzcash::SaplingFullViewingKey &fvk) const
+    {
+        bool result;
+        {
+            LOCK(cs_SpendingKeyStore);
+            result = (mapSaplingSpendingKeys.count(fvk) > 0);
+        }
+        return result;
+    }
+    bool GetSaplingSpendingKey(const libzcash::SaplingFullViewingKey &fvk, libzcash::SaplingSpendingKey &skOut) const
+    {
+        {
+            LOCK(cs_SpendingKeyStore);
+            
+            SaplingSpendingKeyMap::const_iterator mi = mapSaplingSpendingKeys.find(fvk);
+            if (mi != mapSaplingSpendingKeys.end())
+            {
+                skOut = mi->second;
+                return true;
+            }
+        }
+        return false;
+    }
 
     virtual bool AddViewingKey(const libzcash::SproutViewingKey &vk);
     virtual bool RemoveViewingKey(const libzcash::SproutViewingKey &vk);

src/wallet/gtest/test_wallet_zkeys.cpp

@@ -7,9 +7,42 @@
 
 #include <boost/filesystem.hpp>
 
+/**
+ * This test covers Sapling methods on CWallet
+ * GenerateNewSaplingZKey()
+ */
+TEST(wallet_zkeys_tests, store_and_load_sapling_zkeys) {
+    CWallet wallet;
+    
+    // wallet should be empty
+    // std::set<libzcash::SaplingPaymentAddress> addrs;
+    // wallet.GetSaplingPaymentAddresses(addrs);
+    // ASSERT_EQ(0, addrs.size());
+        
+    // wallet should have one key
+    auto saplingAddr = wallet.GenerateNewSaplingZKey();
+    // ASSERT_NE(boost::get<libzcash::SaplingPaymentAddress>(&address), nullptr);
+    // auto sapling_addr = boost::get<libzcash::SaplingPaymentAddress>(saplingAddr);
+    // wallet.GetSaplingPaymentAddresses(addrs);
+    // ASSERT_EQ(1, addrs.size());
+
+    auto sk = libzcash::SaplingSpendingKey::random();
+    auto full_viewing_key = sk.full_viewing_key();
+    ASSERT_TRUE(wallet.AddSaplingSpendingKey(sk));
+    
+    // verify wallet has spending key for the address
+    ASSERT_TRUE(wallet.HaveSaplingSpendingKey(full_viewing_key));
+    
+    // check key is the same
+    libzcash::SaplingSpendingKey keyOut;
+    wallet.GetSaplingSpendingKey(full_viewing_key, keyOut);
+    ASSERT_EQ(sk, keyOut);
+}
+
 /**
  * This test covers methods on CWallet
  * GenerateNewZKey()
+ * GenerateNewSaplingZKey()
  * AddZKey()
  * LoadZKey()
  * LoadZKeyMetadata()

src/wallet/wallet.cpp

@@ -100,8 +100,64 @@ libzcash::PaymentAddress CWallet::GenerateNewZKey()
     return addr;
 }
 
+//! TODO: Should be Sapling address format, SaplingPaymentAddress
+// Generate a new Sapling spending key and return its public payment address
+SaplingPaymentAddress CWallet::GenerateNewSaplingZKey()
+{
+    AssertLockHeld(cs_wallet); // mapZKeyMetadata
+    auto sk = SaplingSpendingKey::random();
+    auto fvk = sk.full_viewing_key();
+    auto addrOpt = sk.default_address();
+    if (addrOpt){
+        auto addr = addrOpt.value();
+        // Check for collision, even though it is unlikely to ever occur
+        if (CCryptoKeyStore::HaveSaplingSpendingKey(fvk))
+            throw std::runtime_error("CWallet::GenerateNewSaplingZKey(): Collision detected");
+
+        // Create new metadata
+        int64_t nCreationTime = GetTime();
+        mapSaplingZKeyMetadata[addr] = CKeyMetadata(nCreationTime);
+        
+        if (!AddSaplingZKey(sk)) {
+            throw std::runtime_error("CWallet::GenerateNewSaplingZKey(): AddSaplingZKey failed");
+        }
+        // return default sapling payment address.
+        return addr;
+    } else {
+        throw std::runtime_error("CWallet::GenerateNewSaplingZKey(): default_address() did not return address");
+    }
+}
+
+// Add spending key to keystore 
+// TODO: persist to disk
+bool CWallet::AddSaplingZKey(const libzcash::SaplingSpendingKey &sk)
+{
+    AssertLockHeld(cs_wallet); // mapSaplingZKeyMetadata
+    auto addr = sk.default_address();
+
+    if (!CCryptoKeyStore::AddSaplingSpendingKey(sk)) {
+        return false;
+    }
+
+    // // check if we need to remove from viewing keys
+    // if (HaveViewingKey(addr)) {
+    //     RemoveViewingKey(key.viewing_key());
+    // }
+
+    // if (!fFileBacked) {
+    //     return true;
+    // }
+
+    // if (!IsCrypted()) {
+    //     return CWalletDB(strWalletFile).WriteSaplingZKey(addr,
+    //                                               sk,
+    //                                               mapSaplingZKeyMetadata[addr]);
+    // }
+    return true;
+}
+
+
 // Add spending key to keystore and persist to disk
-// TODO: Add Sapling support
 bool CWallet::AddZKey(const libzcash::SproutSpendingKey &key)
 {
     AssertLockHeld(cs_wallet); // mapZKeyMetadata
@@ -176,7 +232,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)
@@ -516,7 +572,7 @@ bool CWallet::Verify(const string& walletFile, string& warningString, string& er
         } catch (const boost::filesystem::filesystem_error&) {
             // failure is ok (well, not really, but it's not worse than what we started with)
         }
-        
+
         // try again
         if (!bitdb.Open(GetDataDir())) {
             // if it still fails, it probably means we can't even create the database env
@@ -525,14 +581,14 @@ bool CWallet::Verify(const string& walletFile, string& warningString, string& er
             return true;
         }
     }
-    
+
     if (GetBoolArg("-salvagewallet", false))
     {
         // Recover readable keypairs:
         if (!CWalletDB::Recover(bitdb, walletFile, true))
             return false;
     }
-    
+
     if (boost::filesystem::exists(GetDataDir() / walletFile))
     {
         CDBEnv::VerifyResult r = bitdb.Verify(walletFile, CWalletDB::Recover);
@@ -546,7 +602,7 @@ bool CWallet::Verify(const string& walletFile, string& warningString, string& er
         if (r == CDBEnv::RECOVER_FAIL)
             errorString += _("wallet.dat corrupt, salvage failed");
     }
-    
+
     return true;
 }
 
@@ -2523,7 +2579,7 @@ bool CWallet::FundTransaction(CMutableTransaction& tx, CAmount &nFeeRet, int& nC
 
     CReserveKey reservekey(this);
     CWalletTx wtx;
-    
+
     if (!CreateTransaction(vecSend, wtx, reservekey, nFeeRet, nChangePosRet, strFailReason, &coinControl, false))
         return false;
 
@@ -2590,7 +2646,7 @@ bool CWallet::CreateTransaction(const vector<CRecipient>& vecSend, CWalletTx& wt
     if (!NetworkUpgradeActive(nextBlockHeight, Params().GetConsensus(), Consensus::UPGRADE_SAPLING)) {
         max_tx_size = MAX_TX_SIZE_BEFORE_SAPLING;
     }
-    
+
     // Discourage fee sniping.
     //
     // However because of a off-by-one-error in previous versions we need to
@@ -3062,7 +3118,7 @@ bool CWallet::SetDefaultKey(const CPubKey &vchPubKey)
 
 /**
  * Mark old keypool keys as used,
- * and generate all new keys 
+ * and generate all new keys
  */
 bool CWallet::NewKeyPool()
 {
@@ -3767,7 +3823,7 @@ void CWallet::GetFilteredNotes(
             if (ignoreUnspendable && !HaveSpendingKey(pa)) {
                 continue;
             }
-            
+
             // skip locked notes
             if (IsLockedNote(jsop)) {
                 continue;

src/wallet/wallet.h

@@ -786,6 +786,7 @@ public:
     std::set<int64_t> setKeyPool;
     std::map<CKeyID, CKeyMetadata> mapKeyMetadata;
     std::map<libzcash::SproutPaymentAddress, CKeyMetadata> mapZKeyMetadata;
+    std::map<libzcash::SaplingPaymentAddress, CKeyMetadata> mapSaplingZKeyMetadata;
 
     typedef std::map<unsigned int, CMasterKey> MasterKeyMap;
     MasterKeyMap mapMasterKeys;
@@ -977,6 +978,14 @@ public:
     bool RemoveViewingKey(const libzcash::SproutViewingKey &vk);
     //! Adds a viewing key to the store, without saving it to disk (used by LoadWallet)
     bool LoadViewingKey(const libzcash::SproutViewingKey &dest);
+    
+    /**
+      * Sapling ZKeys
+      */
+    //! Generates new Sapling key
+    libzcash::SaplingPaymentAddress GenerateNewSaplingZKey();
+    //! Adds Sapling spending key to the store, and saves it to disk
+    bool AddSaplingZKey(const libzcash::SaplingSpendingKey &key);
 
     /** 
      * Increment the next transaction order id