Merge pull request #5569 from str4d/feature/wallet_orchard-ua_integra…

…tion

Integrate Orchard into Unified Address types

.cargo/config.offline

@@ -8,7 +8,7 @@ replace-with = "vendored-sources"
 
 [source."https://github.com/zcash/orchard.git"]
 git = "https://github.com/zcash/orchard.git"
-rev = "4dc1ae059a59ee911134cb3e731c7be627a71d4d"
+rev = "3b8d07f7b64b2329622089ac9698e4cce97e2f14"
 replace-with = "vendored-sources"
 
 [source."https://github.com/nuttycom/hdwallet.git"]

Cargo.toml

@@ -72,7 +72,7 @@ codegen-units = 1
 
 [patch.crates-io]
 hdwallet = { git = "https://github.com/nuttycom/hdwallet", rev = "576683b9f2865f1118c309017ff36e01f84420c9" }
-orchard = { git = "https://github.com/zcash/orchard.git", rev = "4dc1ae059a59ee911134cb3e731c7be627a71d4d" }
+orchard = { git = "https://github.com/zcash/orchard.git", rev = "3b8d07f7b64b2329622089ac9698e4cce97e2f14" }
 zcash_address = { git = "https://github.com/zcash/librustzcash.git", rev = "3d935a94e75786a67c3ea4992d7c372af203086f" }
 zcash_history = { git = "https://github.com/zcash/librustzcash.git", rev = "3d935a94e75786a67c3ea4992d7c372af203086f" }
 zcash_note_encryption = { git = "https://github.com/zcash/librustzcash.git", rev = "3d935a94e75786a67c3ea4992d7c372af203086f" }

src/gtest/test_keys.cpp

@@ -73,6 +73,12 @@ namespace libzcash {
     public:
         ReceiverToString() {}
 
+        std::string operator()(const OrchardRawAddress &zaddr) const {
+            CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
+            ss << zaddr;
+            return tfm::format("Orchard(%s)", HexStr(ss.begin(), ss.end()));
+        }
+
         std::string operator()(const SaplingPaymentAddress &zaddr) const {
             CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
             ss << zaddr;
@@ -126,8 +132,11 @@ TEST(Keys, EncodeAndDecodeUnifiedAddresses)
             // These were added to the UA in preference order by the Python test vectors.
             if (!test[3].isNull()) {
                 auto data = ParseHex(test[3].get_str());
-                libzcash::UnknownReceiver r(0x03, data);
-                ua.AddReceiver(r);
+                CDataStream ss(
+                    data,
+                    SER_NETWORK,
+                    PROTOCOL_VERSION);
+                ua.AddReceiver(libzcash::OrchardRawAddress::Read(ss));
             }
             if (!test[2].isNull()) {
                 auto data = ParseHex(test[2].get_str());
@@ -186,6 +195,16 @@ TEST(Keys, DeriveUnifiedFullViewingKeys)
         if (test.size() == 1) continue; // comment
 
         try {
+            // [
+            //     t_key_bytes,
+            //     sapling_fvk_bytes,
+            //     orchard_fvk_bytes,
+            //     unknown_fvk_typecode,
+            //     unknown_fvk_bytes,
+            //     unified_fvk,
+            //     root_seed,
+            //     account,
+            // ],
             auto seed_hex = test[6].get_str();
             auto seed_data = ParseHex(seed_hex);
             RawHDSeed raw_seed(seed_data.begin(), seed_data.end());
@@ -235,6 +254,24 @@ TEST(Keys, DeriveUnifiedFullViewingKeys)
                 auto key = libzcash::SaplingDiversifiableFullViewingKey::Read(ss);
                 EXPECT_EQ(key, saplingKey);
             }
+            if (!test[2].isNull()) {
+                auto expectedHex = test[2].get_str();
+
+                // Ensure that the serialized Orchard fvk matches the test data.
+                auto orchardKey = ufvk.GetOrchardKey().value();
+                CDataStream ssEncode(SER_NETWORK, PROTOCOL_VERSION);
+                ssEncode << orchardKey;
+                EXPECT_EQ(ssEncode.size(), 96);
+                auto skeyHex = HexStr(ssEncode.begin(), ssEncode.end());
+                EXPECT_EQ(expectedHex, skeyHex);
+
+                // Ensure that parsing the test data derives the correct dfvk
+                auto data = ParseHex(expectedHex);
+                ASSERT_EQ(data.size(), 96);
+                CDataStream ss(data, SER_NETWORK, PROTOCOL_VERSION);
+                auto key = libzcash::OrchardFullViewingKey::Read(ss);
+                EXPECT_EQ(key, orchardKey);
+            }
             // Enable the following after Orchard keys are supported.
             //{
             //    auto fvk_data = ParseHex(test[5].get_str());
@@ -284,14 +321,13 @@ TEST(Keys, EncodeAndDecodeUnifiedFullViewingKeys)
             auto key = libzcash::SaplingDiversifiableFullViewingKey::Read(ss);
             ASSERT_TRUE(builder.AddSaplingKey(key));
         }
-
-        // Orchard keys and unknown items are not yet supported; instead,
-        // we just test that we're able to parse the unified key string
-        // and that the constituent items match the elements; if no Sapling
-        // key is present then UFVK construction would fail because it might
-        // presume the UFVK to be transparent-only.
-        if (test[1].isNull())
-            continue;
+        if (!test[2].isNull()) {
+            auto data = ParseHex(test[2].get_str());
+            ASSERT_EQ(data.size(), 96);
+            CDataStream ss(data, SER_NETWORK, PROTOCOL_VERSION);
+            auto key = libzcash::OrchardFullViewingKey::Read(ss);
+            ASSERT_TRUE(builder.AddOrchardKey(key));
+        }
 
         auto built = builder.build();
         ASSERT_TRUE(built.has_value());
@@ -304,5 +340,6 @@ TEST(Keys, EncodeAndDecodeUnifiedFullViewingKeys)
 
         EXPECT_EQ(decoded.value().GetTransparentKey(), built.value().GetTransparentKey());
         EXPECT_EQ(decoded.value().GetSaplingKey(), built.value().GetSaplingKey());
+        EXPECT_EQ(decoded.value().GetOrchardKey(), built.value().GetOrchardKey());
     }
 }

src/gtest/test_keystore.cpp

@@ -557,17 +557,22 @@ TEST(KeystoreTests, StoreAndRetrieveUFVK) {
     EXPECT_EQ(keyStore.GetUnifiedFullViewingKey(ufvkid).value(), zufvk);
 
     auto addrPair = std::get<std::pair<UnifiedAddress, diversifier_index_t>>(zufvk.FindAddress(diversifier_index_t(0), {ReceiverType::Sapling}));
-
-
     EXPECT_TRUE(addrPair.first.GetSaplingReceiver().has_value());
     auto saplingReceiver = addrPair.first.GetSaplingReceiver().value();
-    auto ufvkmeta = keyStore.GetUFVKMetadataForReceiver(saplingReceiver);
-    EXPECT_FALSE(ufvkmeta.has_value());
 
+    // We detect this even though we haven't added the Sapling address, because
+    // we trial-decrypt diversifiers (which also means we learn the index).
+    auto ufvkmetaUnadded = keyStore.GetUFVKMetadataForReceiver(saplingReceiver);
+    EXPECT_TRUE(ufvkmetaUnadded.has_value());
+    EXPECT_EQ(ufvkmetaUnadded.value().first, ufvkid);
+    EXPECT_EQ(ufvkmetaUnadded.value().second.value(), addrPair.second);
+
+    // Adding the Sapling addr -> ivk map entry causes us to find the same UFVK,
+    // but as we're no longer trial-decrypting we don't learn the index.
     auto saplingIvk = zufvk.GetSaplingKey().value().ToIncomingViewingKey();
     keyStore.AddSaplingPaymentAddress(saplingIvk, saplingReceiver);
 
-    ufvkmeta = keyStore.GetUFVKMetadataForReceiver(saplingReceiver);
+    auto ufvkmeta = keyStore.GetUFVKMetadataForReceiver(saplingReceiver);
     EXPECT_TRUE(ufvkmeta.has_value());
     EXPECT_EQ(ufvkmeta.value().first, ufvkid);
     EXPECT_FALSE(ufvkmeta.value().second.has_value());

src/key_io.cpp

@@ -56,6 +56,7 @@ class DataLenForReceiver {
 public:
     DataLenForReceiver() {}
 
+    size_t operator()(const libzcash::OrchardRawAddress &zaddr) const { return 43; }
     size_t operator()(const libzcash::SaplingPaymentAddress &zaddr) const { return 43; }
     size_t operator()(const CScriptID &p2sh) const { return 20; }
     size_t operator()(const CKeyID &p2pkh) const { return 20; }
@@ -76,6 +77,13 @@ class CopyDataForReceiver {
 public:
     CopyDataForReceiver(unsigned char* data, size_t length) : data(data), length(length) {}
 
+    void operator()(const libzcash::OrchardRawAddress &zaddr) const {
+        CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
+        ss << zaddr;
+        assert(length == ss.size());
+        memcpy(data, ss.data(), ss.size());
+    }
+
     void operator()(const libzcash::SaplingPaymentAddress &zaddr) const {
         CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
         ss << zaddr;
@@ -433,6 +441,12 @@ std::optional<T1> DecodeAny(
     return std::nullopt;
 }
 
+static bool AddOrchardReceiver(void* ua, OrchardRawAddressPtr* ptr)
+{
+    return reinterpret_cast<libzcash::UnifiedAddress*>(ua)->AddReceiver(
+        libzcash::OrchardRawAddress::KeyIoOnlyFromReceiver(ptr));
+}
+
 /**
  * `raw` MUST be 43 bytes.
  */
@@ -492,6 +506,7 @@ std::optional<libzcash::PaymentAddress> KeyIO::DecodePaymentAddress(const std::s
         str.c_str(),
         keyConstants.NetworkIDString().c_str(),
         &ua,
+        AddOrchardReceiver,
         AddSaplingReceiver,
         AddP2SHReceiver,
         AddP2PKHReceiver,

src/keystore.cpp

@@ -252,6 +252,10 @@ bool CBasicKeyStore::GetSproutViewingKey(
     return false;
 }
 
+//
+// Sapling Keys
+//
+
 bool CBasicKeyStore::GetSaplingFullViewingKey(
     const libzcash::SaplingIncomingViewingKey &ivk,
     libzcash::SaplingExtendedFullViewingKey &extfvkOut) const
@@ -308,14 +312,26 @@ bool CBasicKeyStore::AddUnifiedFullViewingKey(
 {
     LOCK(cs_KeyStore);
 
+    auto ufvkId = ufvk.GetKeyID();
+
+    // Add the Orchard component of the UFVK to the wallet.
+    auto orchardKey = ufvk.GetOrchardKey();
+    if (orchardKey.has_value()) {
+        auto ivk = orchardKey.value().ToIncomingViewingKey();
+        mapOrchardKeyUnified.insert(std::make_pair(ivk, ufvkId));
+
+        auto ivkInternal = orchardKey.value().ToInternalIncomingViewingKey();
+        mapOrchardKeyUnified.insert(std::make_pair(ivkInternal, ufvkId));
+    }
+
     // Add the Sapling component of the UFVK to the wallet.
     auto saplingKey = ufvk.GetSaplingKey();
     if (saplingKey.has_value()) {
         auto ivk = saplingKey.value().ToIncomingViewingKey();
-        mapSaplingKeyUnified.insert(std::make_pair(ivk, ufvk.GetKeyID()));
+        mapSaplingKeyUnified.insert(std::make_pair(ivk, ufvkId));
 
         auto changeIvk = saplingKey.value().GetChangeIVK();
-        mapSaplingKeyUnified.insert(std::make_pair(changeIvk, ufvk.GetKeyID()));
+        mapSaplingKeyUnified.insert(std::make_pair(changeIvk, ufvkId));
     }
 
     // We can't reasonably add the transparent component here, because
@@ -325,7 +341,7 @@ bool CBasicKeyStore::AddUnifiedFullViewingKey(
     // transparent part of the address must be added to the keystore.
 
     // Add the UFVK by key identifier.
-    mapUnifiedFullViewingKeys.insert({ufvk.GetKeyID(), ufvk});
+    mapUnifiedFullViewingKeys.insert({ufvkId, ufvk});
 
     return true;
 }
@@ -386,6 +402,15 @@ std::optional<libzcash::UFVKId> CBasicKeyStore::GetUFVKIdForViewingKey(const lib
             }
         },
         [&](const libzcash::UnifiedFullViewingKey& ufvk) {
+            const auto orchardFvk = ufvk.GetOrchardKey();
+            if (orchardFvk.has_value()) {
+                const auto orchardIvk = orchardFvk.value().ToIncomingViewingKey();
+                const auto ufvkId = mapOrchardKeyUnified.find(orchardIvk);
+                if (ufvkId != mapOrchardKeyUnified.end()) {
+                    result = ufvkId->second;
+                    return;
+                }
+            }
             const auto saplingDfvk = ufvk.GetSaplingKey();
             if (saplingDfvk.has_value()) {
                 const auto saplingIvk = saplingDfvk.value().ToIncomingViewingKey();
@@ -399,19 +424,54 @@ std::optional<libzcash::UFVKId> CBasicKeyStore::GetUFVKIdForViewingKey(const lib
     return result;
 }
 
+std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>
+FindUFVKId::operator()(const libzcash::OrchardRawAddress& orchardAddr) const {
+    for (const auto& [k, v] : keystore.mapUnifiedFullViewingKeys) {
+        auto fvk = v.GetOrchardKey();
+        if (fvk.has_value()) {
+            auto d_idx = fvk.value().ToIncomingViewingKey().DecryptDiversifier(orchardAddr);
+            if (d_idx.has_value()) {
+                return std::make_pair(k, d_idx);
+            }
+        }
+    }
+    return std::nullopt;
+}
+
 std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>
 FindUFVKId::operator()(const libzcash::SaplingPaymentAddress& saplingAddr) const {
     const auto saplingIvk = keystore.mapSaplingIncomingViewingKeys.find(saplingAddr);
     if (saplingIvk != keystore.mapSaplingIncomingViewingKeys.end()) {
+        // We have either generated this as a receiver via `z_getaddressforaccount` or a
+        // legacy Sapling address via `z_getnewaddress`, or we have previously detected
+        // this via trial-decryption of a note.
         const auto ufvkId = keystore.mapSaplingKeyUnified.find(saplingIvk->second);
         if (ufvkId != keystore.mapSaplingKeyUnified.end()) {
             return std::make_pair(ufvkId->second, std::nullopt);
         } else {
+            // If we have the addr -> ivk map entry but not the ivk -> UFVK map entry,
+            // then this is definitely a legacy Sapling address.
             return std::nullopt;
         }
-    } else {
-        return std::nullopt;
     }
+
+    // We haven't generated this receiver via `z_getaddressforaccount` (or this is a
+    // recovery from a backed-up mnemonic which doesn't store receiver types selected by
+    // users). Trial-decrypt the diversifier of the Sapling address with every UFVK in the
+    // wallet, to check directly if it belongs to any of them.
+    for (const auto& [k, v] : keystore.mapUnifiedFullViewingKeys) {
+        auto dfvk = v.GetSaplingKey();
+        if (dfvk.has_value()) {
+            auto d_idx = dfvk.value().DecryptDiversifier(saplingAddr.d);
+            auto derived_addr = dfvk.value().Address(d_idx);
+            if (derived_addr.has_value() && derived_addr.value() == saplingAddr) {
+                return std::make_pair(k, d_idx);
+            }
+        }
+    }
+
+    // We definitely don't know of any UFVK linked to this Sapling address.
+    return std::nullopt;
 }
 std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>
 FindUFVKId::operator()(const CScriptID& scriptId) const {

src/keystore.h

@@ -185,6 +185,7 @@ class CBasicKeyStore : public CKeyStore
     std::map<CKeyID, std::pair<libzcash::UFVKId, libzcash::diversifier_index_t>> mapP2PKHUnified;
     std::map<CScriptID, std::pair<libzcash::UFVKId, libzcash::diversifier_index_t>> mapP2SHUnified;
     std::map<libzcash::SaplingIncomingViewingKey, libzcash::UFVKId> mapSaplingKeyUnified;
+    std::map<libzcash::OrchardIncomingViewingKey, libzcash::UFVKId> mapOrchardKeyUnified;
     std::map<libzcash::UFVKId, libzcash::ZcashdUnifiedFullViewingKey> mapUnifiedFullViewingKeys;
 
     friend class FindUFVKId;
@@ -398,6 +399,8 @@ class FindUFVKId {
 public:
     FindUFVKId(const CBasicKeyStore& keystore): keystore(keystore) {}
 
+    std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>
+        operator()(const libzcash::OrchardRawAddress& orchardAddr) const;
     std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>
         operator()(const libzcash::SaplingPaymentAddress& saplingAddr) const;
     std::optional<std::pair<libzcash::UFVKId, std::optional<libzcash::diversifier_index_t>>>

src/rust/include/rust/address.h

@@ -5,10 +5,13 @@
 #ifndef ZCASH_RUST_INCLUDE_RUST_ADDRESS_H
 #define ZCASH_RUST_INCLUDE_RUST_ADDRESS_H
 
+#include "rust/orchard/keys.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+typedef bool (*orchard_receiver_t)(void* ua, OrchardRawAddressPtr* addr);
 typedef bool (*raw_to_receiver_t)(void* ua, const unsigned char* raw);
 typedef bool (*unknown_receiver_t)(
     void* ua,
@@ -24,6 +27,7 @@ bool zcash_address_parse_unified(
     const char* str,
     const char* network,
     void* ua,
+    orchard_receiver_t orchard_cb,
     raw_to_receiver_t sapling_cb,
     raw_to_receiver_t p2sh_cb,
     raw_to_receiver_t p2pkh_cb,