From 290b2e4e3972787e70336a5059af475f1bc633ec Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Fri, 24 Sep 2021 15:20:13 -0600
Subject: [PATCH] Use the default UA-based Sapling address for the
 saplingmigration tool.

---
 src/test/arith_uint256_tests.cpp                  |  8 ++++++++
 src/uint256.h                                     | 13 +++++--------
 src/wallet/asyncrpcoperation_saplingmigration.cpp |  9 +++++----
 src/wallet/wallet.cpp                             | 11 +++++++++--
 src/zcash/address/zip32.h                         | 10 ++++++++++
 5 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/test/arith_uint256_tests.cpp b/src/test/arith_uint256_tests.cpp
index a1b7e2625..47563fdd8 100644
--- a/src/test/arith_uint256_tests.cpp
+++ b/src/test/arith_uint256_tests.cpp
@@ -564,4 +564,12 @@ BOOST_AUTO_TEST_CASE( getmaxcoverage ) // some more tests just to get 100% cover
     CHECKBITWISEOPERATOR(R1,~R2,&)
 }
 
+BOOST_AUTO_TEST_CASE( blob88_increment )
+{
+    blob88 bzero(0);
+    blob88 bone(1);
+    BOOST_CHECK(bzero.increment());
+    BOOST_CHECK(bzero == bone);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/uint256.h b/src/uint256.h
index 76a7917ce..68000108d 100644
--- a/src/uint256.h
+++ b/src/uint256.h
@@ -120,18 +120,15 @@ public:
     }
     explicit blob88(const std::vector<unsigned char>& vch) : base_blob<88>(vch) {}
 
-    std::optional<blob88> increment() const {
-        blob88 result = *this;
-
+    bool increment() {
         for (int i = 0; i < 11; i++) {
-            result.data[i] += 1;
-            if (result.data[i] != 0) {
-                // no overflow
-                return result;
+            this->data[i] += 1;
+            if (this->data[i] != 0) {
+                return true; // no overflow
             }
         }
 
-        return std::nullopt;
+        return false; //overflow
     }
 
     // treat as little-endian for numeric comparison
diff --git a/src/wallet/asyncrpcoperation_saplingmigration.cpp b/src/wallet/asyncrpcoperation_saplingmigration.cpp
index 5bfdb759a..24438e6a2 100644
--- a/src/wallet/asyncrpcoperation_saplingmigration.cpp
+++ b/src/wallet/asyncrpcoperation_saplingmigration.cpp
@@ -192,6 +192,7 @@ CAmount AsyncRPCOperation_saplingmigration::chooseAmount(const CAmount& availabl
 }
 
 // Unless otherwise specified, the migration destination address is the address for Sapling account 0
+// at the smallest diversifier index which produces a valid diversified address.
 libzcash::SaplingPaymentAddress AsyncRPCOperation_saplingmigration::getMigrationDestAddress(const HDSeed& seed) {
     KeyIO keyIO(Params());
     if (mapArgs.count("-migrationdestaddress")) {
@@ -202,12 +203,12 @@ libzcash::SaplingPaymentAddress AsyncRPCOperation_saplingmigration::getMigration
         return *saplingAddress;
     }
 
-    // TODO: use UVK-based derivation here instead.
     auto usk = pwalletMain->GetUnifiedSpendingKeyForAccount(0);
     assert(usk.has_value()); // mnemonic seeds are currently always generated to have valid USKs at account 0
-    auto xsk = usk.value().GetSaplingExtendedSpendingKey();
-    if (xsk.has_value()) {
-        return xsk.value().ToXFVK().DefaultAddress();
+    auto ua = usk.value().ToFullViewingKey().FindAddress(libzcash::diversifier_index_t(0));
+    auto addr = ua.first.GetSaplingPaymentAddress();
+    if (addr.has_value()) {
+        return addr.value();
     } else {
         // This error will only occur if Sapling address generation has been disbled for USKs from this
         // wallet.
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index bdf98e4da..7d475dcca 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -394,8 +394,15 @@ UnifiedSpendingKey CWallet::GenerateNewUnifiedSpendingKey() {
 }
 
 std::optional<libzcash::UnifiedSpendingKey> CWallet::GetUnifiedSpendingKeyForAccount(uint32_t accountId) {
-    //TODO
-    return std::nullopt;
+    auto seed = GetMnemonicSeed();
+    assert(seed.has_value());
+    // TODO: is there any reason to cache and not re-derive this every time?
+    auto usk = UnifiedSpendingKey::Derive(seed.value(), BIP44CoinType(), accountId);
+    if (usk.has_value()) {
+        return usk.value().first;
+    } else {
+        return std::nullopt;
+    }
 }
 
 void CWallet::LoadKeyMetadata(const CPubKey &pubkey, const CKeyMetadata &meta)
diff --git a/src/zcash/address/zip32.h b/src/zcash/address/zip32.h
index 46eff708d..3aef9d960 100644
--- a/src/zcash/address/zip32.h
+++ b/src/zcash/address/zip32.h
@@ -321,6 +321,16 @@ public:
     }
 
     std::optional<ZcashdUnifiedAddress> Address(diversifier_index_t j) const;
+
+    std::pair<ZcashdUnifiedAddress, diversifier_index_t> FindAddress(diversifier_index_t j) const {
+        auto addr = Address(j);
+        while (!addr.has_value()) {
+            if (!j.increment())
+                throw std::runtime_error(std::string(__func__) + ": diversifier index overflow.");;
+            addr = Address(j);
+        }
+        return std::make_pair(addr.value(), j);
+    }
 };
 
 class UnifiedSpendingKey {