IsUsedDestination should count any known single-key address

instagibbs · instagibbs · commit 09502452bbbe · 2020-01-03T17:20:46.000-05:00
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
@@ -2927,7 +2927,7 @@ static UniValue listunspent(const JSONRPCRequest& request)
         CTxDestination address;
         const CScript& scriptPubKey = out.tx->tx->vout[out.i].scriptPubKey;
         bool fValidAddress = ExtractDestination(scriptPubKey, address);
-        bool reused = avoid_reuse && pwallet->IsUsedDestination(address);
+        bool reused = avoid_reuse && pwallet->IsUsedDestination(out.tx->GetHash(), out.i);
 
         if (destinations.size() && (!fValidAddress || !destinations.count(address)))
             continue;
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
@@ -719,17 +719,33 @@ void CWallet::SetUsedDestinationState(WalletBatch& batch, const uint256& hash, u
     }
 }
 
-bool CWallet::IsUsedDestination(const CTxDestination& dst) const
-{
-    LOCK(cs_wallet);
-    return IsMine(dst) && GetDestData(dst, "used", nullptr);
-}
-
 bool CWallet::IsUsedDestination(const uint256& hash, unsigned int n) const
 {
+    AssertLockHeld(cs_wallet);
     CTxDestination dst;
     const CWalletTx* srctx = GetWalletTx(hash);
-    return srctx && ExtractDestination(srctx->tx->vout[n].scriptPubKey, dst) && IsUsedDestination(dst);
+    if (srctx) {
+        assert(srctx->tx->vout.size() > n);
+        LegacyScriptPubKeyMan* spk_man = GetLegacyScriptPubKeyMan();
+        // When descriptor wallets arrive, these additional checks are
+        // likely superfluous and can be optimized out
+        assert(spk_man != nullptr);
+        for (const auto& keyid : GetAffectedKeys(srctx->tx->vout[n].scriptPubKey, *spk_man)) {
+            WitnessV0KeyHash wpkh_dest(keyid);
+            if (GetDestData(wpkh_dest, "used", nullptr)) {
+                return true;
+            }
+            ScriptHash sh_wpkh_dest(wpkh_dest);
+            if (GetDestData(sh_wpkh_dest, "used", nullptr)) {
+                return true;
+            }
+            PKHash pkh_dest(keyid);
+            if (GetDestData(pkh_dest, "used", nullptr)) {
+                return true;
+            }
+        }
+    }
+    return false;
 }
 
 bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFlushOnClose)
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -810,9 +810,8 @@ class CWallet final : public WalletStorage, private interfaces::Chain::Notificat
 
     bool IsSpent(const uint256& hash, unsigned int n) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
-    // Whether this or any UTXO with the same CTxDestination has been spent.
-    bool IsUsedDestination(const CTxDestination& dst) const;
-    bool IsUsedDestination(const uint256& hash, unsigned int n) const;
+    // Whether this or any known UTXO with the same single key has been spent.
+    bool IsUsedDestination(const uint256& hash, unsigned int n) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     void SetUsedDestinationState(WalletBatch& batch, const uint256& hash, unsigned int n, bool used) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     std::vector<OutputGroup> GroupOutputs(const std::vector<COutput>& outputs, bool single_coin) const;
diff --git a/test/functional/wallet_avoidreuse.py b/test/functional/wallet_avoidreuse.py
@@ -86,7 +86,12 @@ def run_test(self):
         reset_balance(self.nodes[1], self.nodes[0].getnewaddress())
         self.test_fund_send_fund_senddirty()
         reset_balance(self.nodes[1], self.nodes[0].getnewaddress())
-        self.test_fund_send_fund_send()
+        self.test_fund_send_fund_send("legacy")
+        reset_balance(self.nodes[1], self.nodes[0].getnewaddress())
+        self.test_fund_send_fund_send("p2sh-segwit")
+        reset_balance(self.nodes[1], self.nodes[0].getnewaddress())
+        self.test_fund_send_fund_send("bech32")
+
 
     def test_persistence(self):
         '''Test that wallet files persist the avoid_reuse flag.'''
@@ -182,7 +187,7 @@ def test_fund_send_fund_senddirty(self):
         assert_approx(self.nodes[1].getbalance(), 5, 0.001)
         assert_approx(self.nodes[1].getbalance(avoid_reuse=False), 5, 0.001)
 
-    def test_fund_send_fund_send(self):
+    def test_fund_send_fund_send(self, second_addr_type):
         '''
         Test the simple case where [1] generates a new address A, then
         [0] sends 10 BTC to A.
@@ -193,7 +198,7 @@ def test_fund_send_fund_send(self):
         '''
         self.log.info("Test fund send fund send")
 
-        fundaddr = self.nodes[1].getnewaddress()
+        fundaddr = self.nodes[1].getnewaddress(label="", address_type="legacy")
         retaddr = self.nodes[0].getnewaddress()
 
         self.nodes[0].sendtoaddress(fundaddr, 10)
@@ -214,7 +219,19 @@ def test_fund_send_fund_send(self):
         # getbalances should show no used, 5 btc trusted
         assert_balances(self.nodes[1], mine={"used": 0, "trusted": 5})
 
-        self.nodes[0].sendtoaddress(fundaddr, 10)
+        # For the second send, we transmute it to a related single-key address
+        # to make sure it's also detected as re-use
+        fund_spk = self.nodes[0].getaddressinfo(fundaddr)["scriptPubKey"]
+        fund_decoded = self.nodes[0].decodescript(fund_spk)
+        if second_addr_type == "p2sh-segwit":
+            new_fundaddr = fund_decoded["segwit"]["p2sh-segwit"]
+        elif second_addr_type == "bech32":
+            new_fundaddr = fund_decoded["segwit"]["addresses"][0]
+        else:
+            new_fundaddr = fundaddr
+            assert_equal(second_addr_type, "legacy")
+
+        self.nodes[0].sendtoaddress(new_fundaddr, 10)
         self.nodes[0].generate(1)
         self.sync_all()
 
