From 4715b31c769ffabdb1083b0f41691b0222fb11ae Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Mon, 9 Jul 2018 20:12:37 +0100
Subject: [PATCH] Implement CKeyStore::GetSaplingPaymentAddresses()

---
 src/keystore.h                         | 14 ++++++++++++++
 src/wallet/gtest/test_wallet_zkeys.cpp | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/keystore.h b/src/keystore.h
index 76772dbf1..c856005c1 100644
--- a/src/keystore.h
+++ b/src/keystore.h
@@ -75,6 +75,7 @@ public:
     virtual bool GetSaplingIncomingViewingKey(
         const libzcash::SaplingPaymentAddress &addr, 
         libzcash::SaplingIncomingViewingKey& ivkOut) const =0;
+    virtual void GetSaplingPaymentAddresses(std::set<libzcash::SaplingPaymentAddress> &setAddress) const =0;
 
     //! Support for viewing keys
     virtual bool AddViewingKey(const libzcash::SproutViewingKey &vk) =0;
@@ -251,6 +252,19 @@ public:
     virtual bool GetSaplingIncomingViewingKey(
         const libzcash::SaplingPaymentAddress &addr, 
         libzcash::SaplingIncomingViewingKey& ivkOut) const;
+    void GetSaplingPaymentAddresses(std::set<libzcash::SaplingPaymentAddress> &setAddress) const
+    {
+        setAddress.clear();
+        {
+            LOCK(cs_SpendingKeyStore);
+            auto mi = mapSaplingIncomingViewingKeys.begin();
+            while (mi != mapSaplingIncomingViewingKeys.end())
+            {
+                setAddress.insert((*mi).first);
+                mi++;
+            }
+        }
+    }
 
     virtual bool AddViewingKey(const libzcash::SproutViewingKey &vk);
     virtual bool RemoveViewingKey(const libzcash::SproutViewingKey &vk);
diff --git a/src/wallet/gtest/test_wallet_zkeys.cpp b/src/wallet/gtest/test_wallet_zkeys.cpp
index 317e4ff54..53963c482 100644
--- a/src/wallet/gtest/test_wallet_zkeys.cpp
+++ b/src/wallet/gtest/test_wallet_zkeys.cpp
@@ -16,7 +16,15 @@ 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 address = wallet.GenerateNewSaplingZKey();
+    wallet.GetSaplingPaymentAddresses(addrs);
+    ASSERT_EQ(1, addrs.size());
     
     // verify wallet has incoming viewing key for the address
     ASSERT_TRUE(wallet.HaveSaplingIncomingViewingKey(address));
@@ -28,6 +36,16 @@ TEST(wallet_zkeys_tests, store_and_load_sapling_zkeys) {
     // verify wallet did add it
     auto fvk = sk.full_viewing_key();
     ASSERT_TRUE(wallet.HaveSaplingSpendingKey(fvk));
+
+    // verify spending key stored correctly
+    libzcash::SaplingSpendingKey keyOut;
+    wallet.GetSaplingSpendingKey(fvk, keyOut);
+    ASSERT_EQ(sk, keyOut);
+
+    // verify there are two keys
+    wallet.GetSaplingPaymentAddresses(addrs);
+    ASSERT_EQ(2, addrs.size());
+    ASSERT_EQ(1, addrs.count(address));
 }
 
 /**