From e56f252a861a87c09f99f01c399d62846f56d025 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Wed, 8 Dec 2021 18:13:28 -0700
Subject: [PATCH] Add unified address tracking to KeyStore

---
 src/keystore.cpp      | 30 +++++++++++++++++++++++-------
 src/keystore.h        | 22 +++++++++++++++++++++-
 src/zcash/Address.cpp | 31 +++++++++++++++++++++++++++++++
 src/zcash/Address.hpp |  6 ++++++
 4 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/src/keystore.cpp b/src/keystore.cpp
index a4779a3dd..4230fd7cb 100644
--- a/src/keystore.cpp
+++ b/src/keystore.cpp
@@ -298,13 +298,6 @@ bool CBasicKeyStore::AddUnifiedFullViewingKey(
 {
     LOCK(cs_KeyStore);
 
-    auto tKey = ufvk.GetTransparentKey();
-    if (tKey.has_value()) {
-        // FIXME: this is probably not the right thing; we need to actually track
-        // addresses, not the id at the FVK level?
-        mapTKeyUnified.insert(std::make_pair(tKey.value().GetPubKey().GetID(), keyId));
-    }
-
     auto saplingKey = ufvk.GetSaplingKey();
     if (saplingKey.has_value()) {
         auto ivk = saplingKey.value().fvk.in_viewing_key();
@@ -315,3 +308,26 @@ bool CBasicKeyStore::AddUnifiedFullViewingKey(
 
     return true;
 }
+
+bool CBasicKeyStore::AddUnifiedAddress(
+        const libzcash::UFVKId& keyId,
+        const libzcash::UnifiedAddress& ua)
+{
+    LOCK(cs_KeyStore);
+
+    // It's only necessary to add p2pkh and p2sh components of
+    // the UA; all other lookups of the associated UFVK will be
+    // made via the protocol-specific viewing key that is used
+    // to trial-decrypt a transaction.
+
+    auto p2pkhReceiver = ua.GetP2PKHReceiver();
+    if (p2pkhReceiver.has_value()) {
+        mapP2PKHUnified.insert(std::make_pair(p2pkhReceiver.value(), keyId));
+    }
+
+    auto p2shReceiver = ua.GetP2SHReceiver();
+    if (p2shReceiver.has_value()) {
+        mapP2SHUnified.insert(std::make_pair(p2shReceiver.value(), keyId));
+    }
+}
+
diff --git a/src/keystore.h b/src/keystore.h
index cb1114665..dd731195a 100644
--- a/src/keystore.h
+++ b/src/keystore.h
@@ -107,6 +107,11 @@ public:
             const libzcash::UFVKId& keyId,
             const libzcash::ZcashdUnifiedFullViewingKey &ufvk
             ) = 0;
+
+    virtual bool AddUnifiedAddress(
+            const libzcash::UFVKId& keyId,
+            const libzcash::UnifiedAddress &ua
+            ) = 0;
 };
 
 typedef std::map<CKeyID, CKey> KeyMap;
@@ -128,6 +133,12 @@ typedef std::map<
 // Only maps from default addresses to ivk, may need to be reworked when adding diversified addresses.
 typedef std::map<libzcash::SaplingPaymentAddress, libzcash::SaplingIncomingViewingKey> SaplingIncomingViewingKeyMap;
 
+struct UnifiedAddressMetadata {
+    libzcash::UFVKId keyId;
+    libzcash::diversifier_index_t j;
+    std::vector<libzcash::ReceiverType> receiverTypes;
+};
+
 /** Basic key store, that keeps keys in an address->secret map */
 class CBasicKeyStore : public CKeyStore
 {
@@ -150,7 +161,8 @@ protected:
     SaplingIncomingViewingKeyMap mapSaplingIncomingViewingKeys;
 
     // Unified key support
-    std::map<CKeyID, libzcash::UFVKId> mapTKeyUnified;
+    std::map<CKeyID, libzcash::UFVKId> mapP2PKHUnified;
+    std::map<CScriptID, libzcash::UFVKId> mapP2SHUnified;
     std::map<libzcash::SaplingIncomingViewingKey, libzcash::UFVKId> mapSaplingKeyUnified;
     std::map<libzcash::UFVKId, libzcash::ZcashdUnifiedFullViewingKey> mapUnifiedFullViewingKeys;
 public:
@@ -329,6 +341,14 @@ public:
     virtual bool AddUnifiedFullViewingKey(
             const libzcash::UFVKId& keyId,
             const libzcash::ZcashdUnifiedFullViewingKey &ufvk);
+
+    /**
+     * Add the transparent component of the unified address, if any,
+     * to the keystore to make it possible to identify the
+     */
+    virtual bool AddUnifiedAddress(
+            const libzcash::UFVKId& keyId,
+            const libzcash::UnifiedAddress &ua);
 };
 
 typedef std::vector<unsigned char, secure_allocator<unsigned char> > CKeyingMaterial;
diff --git a/src/zcash/Address.cpp b/src/zcash/Address.cpp
index af2a1d873..ee6f30645 100644
--- a/src/zcash/Address.cpp
+++ b/src/zcash/Address.cpp
@@ -43,6 +43,37 @@ bool UnifiedAddress::AddReceiver(Receiver receiver) {
     return true;
 }
 
+std::optional<CKeyID> UnifiedAddress::GetP2PKHReceiver() const {
+    for (const auto& r : receivers) {
+        if (std::holds_alternative<CKeyID>(r)) {
+            return std::get<CKeyID>(r);
+        }
+    }
+
+    return std::nullopt;
+}
+
+std::optional<CScriptID> UnifiedAddress::GetP2SHReceiver() const {
+    for (const auto& r : receivers) {
+        if (std::holds_alternative<CScriptID>(r)) {
+            return std::get<CScriptID>(r);
+        }
+    }
+
+    return std::nullopt;
+}
+
+std::optional<SaplingPaymentAddress> UnifiedAddress::GetSaplingReceiver() const {
+    for (const auto& r : receivers) {
+        if (std::holds_alternative<SaplingPaymentAddress>(r)) {
+            return std::get<SaplingPaymentAddress>(r);
+        }
+    }
+
+    return std::nullopt;
+}
+
+
 std::pair<std::string, PaymentAddress> AddressInfoFromSpendingKey::operator()(const SproutSpendingKey &sk) const {
     return std::make_pair("sprout", sk.address());
 }
diff --git a/src/zcash/Address.hpp b/src/zcash/Address.hpp
index 99b9e4a43..88f7e2f88 100644
--- a/src/zcash/Address.hpp
+++ b/src/zcash/Address.hpp
@@ -123,6 +123,12 @@ public:
         return receivers.size();
     }
 
+    std::optional<CKeyID> GetP2PKHReceiver() const;
+
+    std::optional<CScriptID> GetP2SHReceiver() const;
+
+    std::optional<SaplingPaymentAddress> GetSaplingReceiver() const;
+
     friend inline bool operator==(const UnifiedAddress& a, const UnifiedAddress& b) {
         return a.receivers == b.receivers;
     }