Merge bitcoin/bitcoin#36025: psbt: avoid duplicate taproot leaf script keys when merging

1cb4163 psbt: avoid duplicate taproot leaf script keys when merging (Shuvam Pandey)

Pull request description:

  Follow-up to #35665, which fixed the same combiner defect for `PSBT_GLOBAL_XPUB`. thomasbuilds
  and winterrdog asked for this one as its own PR when I reported it there.

  `m_tap_scripts` maps a leaf script to a set of control blocks, but is serialized as one record
  per control block, keyed by the control block (`SerializeToVector(s, PSBT_IN_TAP_LEAF_SCRIPT,
  std::span{control_block})`). `PSBTInput::Merge` unions it by the map key, so two PSBTs that map
  the same control block to different leaf scripts merge into an input that serializes the `0x15`
  key twice. Duplicate keys make a PSBT invalid, so it is the same `combinepsbt` then
  `decodepsbt` failure as the xpub case, at the input level. Present since #22558 (v24.0).

  Both decode on their own, and differ only in the leaf script the control block maps to, `OP_1`
  against `OP_1 OP_1`:

  ```
  $ A=cHNidP8BADwCAAAAAaqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAIhXAUJKbdMGgSVS3i0tgNel6XgeKWg8o7JbVR7/ums6AOsACUcAAAA==
  $ B=cHNidP8BADwCAAAAAaqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAIhXAUJKbdMGgSVS3i0tgNel6XgeKWg8o7JbVR7/ums6AOsADUVHAAAA=
  $ bitcoin-cli -regtest decodepsbt "$(bitcoin-cli -regtest combinepsbt "[\"$A\",\"$B\"]")"
  error code: -22
  error message:
  TX decode failed Duplicate Key, input key "15c050929b74c1a04954b78b4b6035e97a5e078a5a0f28ec96d547bfee9ace803ac0" already provided: unspecified iostream_category error
  ```

  winterrdog reproduced it on the #35665 thread with another pair.

  Merge the records rather than the map entries, keeping the leaf script already there. BIP 174
  lets the combiner "pick arbitrarily when conflicts occur", and unknown and proprietary records
  already resolve that way. Refusing to combine is the BIP's other option, but that would fail
  `combinepsbt` on input it accepts today.

  Merging by map key drops records as well. `std::map::insert` leaves existing keys alone, so
  when both PSBTs carry the same leaf script with different control blocks, the incoming set was
  dropped. Those keys do not conflict, so merging per record keeps them.

  The control blocks already present are collected once per merge rather than searched for per
  incoming record, which would be quadratic in the size of the two PSBTs `combinepsbt` takes from
  the caller.

  Since this is the second field with this shape I checked the rest. `m_xpubs` (#35665) and
  `m_tap_scripts` are the only two whose record key comes from the value, so two map entries can
  serialize the same key. `partial_sigs` is keyed by `CKeyID` and serialized under the pubkey,
  but the pubkey determines the `CKeyID`, so those records stay distinct. The others key the
  record by the map key, `m_proprietary` included, and `PSBTOutput` has no such field.

  The test fails on master on both counts, and covers the merges that do not conflict as well.

  I found this with a local assertion in the psbt fuzz target that a combined PSBT must
  roundtrip. That assertion can go in a follow-up.

  Tested:

  ```
  ./build/bin/test_bitcoin --run_test=psbt_tests
  ./build/bin/test_bitcoin --run_test=psbt_wallet_tests
  ./build/test/functional/test_runner.py rpc_psbt.py rpc_rawtransaction.py wallet_taproot.py wallet_signer.py feature_taproot.py wallet_basic.py
  ```

ACKs for top commit:
  achow101:
    ACK 1cb4163
  winterrdog:
    re-ACK 1cb4163

Tree-SHA512: 2599beffe701ba9b672e8dcc3d43844f3853ceeb3d86fc53798a428d3288aa8edd2e42f338b32a1046fa558d6cc5cfa5d55b3a0aaf960278b0f3002158381623

Author: achow101

src/psbt.cpp

@@ -14,6 +14,7 @@
 #include <util/strencodings.h>
 
 #include <algorithm>
+#include <set>
 
 using common::PSBTError;
 
@@ -436,7 +437,17 @@ bool PSBTInput::Merge(const PSBTInput& input)
     m_proprietary.insert(input.m_proprietary.begin(), input.m_proprietary.end());
     unknown.insert(input.unknown.begin(), input.unknown.end());
     m_tap_script_sigs.insert(input.m_tap_script_sigs.begin(), input.m_tap_script_sigs.end());
-    m_tap_scripts.insert(input.m_tap_scripts.begin(), input.m_tap_scripts.end());
+    // Merge by control block, the serialized key (BIP 371), to avoid duplicate keys. Keep the
+    // leaf script already present; BIP 174 lets the Combiner pick arbitrarily on conflict.
+    std::set<std::vector<unsigned char>> seen_control_blocks;
+    for (const auto& [_, control_blocks] : m_tap_scripts) {
+        seen_control_blocks.insert(control_blocks.begin(), control_blocks.end());
+    }
+    for (const auto& [leaf, control_blocks] : input.m_tap_scripts) {
+        for (const auto& control_block : control_blocks) {
+            if (seen_control_blocks.insert(control_block).second) m_tap_scripts[leaf].insert(control_block);
+        }
+    }
     m_tap_bip32_paths.insert(input.m_tap_bip32_paths.begin(), input.m_tap_bip32_paths.end());
 
     if (redeem_script.empty() && !input.redeem_script.empty()) redeem_script = input.redeem_script;

test/functional/rpc_psbt.py

@@ -43,14 +43,15 @@
     PSBT_IN_PROPRIETARY,
     PSBT_IN_TAP_BIP32_DERIVATION,
     PSBT_IN_TAP_INTERNAL_KEY,
+    PSBT_IN_TAP_LEAF_SCRIPT,
     PSBT_IN_WITNESS_UTXO,
     PSBT_IN_FINAL_SCRIPTWITNESS,
     PSBT_OUT_MUSIG2_PARTICIPANT_PUBKEYS,
     PSBT_OUT_PROPRIETARY,
     PSBT_OUT_TAP_TREE,
     PSBT_OUT_SCRIPT,
 )
-from test_framework.script import CScript, OP_TRUE, SIGHASH_ALL, SIGHASH_ANYONECANPAY, hash160
+from test_framework.script import CScript, LEAF_VERSION_TAPSCRIPT, OP_TRUE, SIGHASH_ALL, SIGHASH_ANYONECANPAY, hash160
 from test_framework.script_util import MIN_STANDARD_TX_NONWITNESS_SIZE, output_key_to_p2tr_script
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
@@ -409,6 +410,66 @@ def psbt_with_origin(fingerprint):
         decoded = self.nodes[0].decodepsbt(combined)
         assert_equal(decoded["global_xpubs"], [{"xpub": xpub, "master_fingerprint": "00000000", "path": "m"}])
 
+    def test_combinepsbt_tap_leaf_script_conflict(self):
+        self.log.info("Test that combining PSBTs with conflicting leaf scripts for the same control block keeps a single record")
+
+        tx = CTransaction()
+        tx.vin = [CTxIn(outpoint=COutPoint(hash=int('aa' * 32, 16), n=0), scriptSig=b"")]
+        tx.vout = [CTxOut(nValue=0, scriptPubKey=b"")]
+
+        def psbt_with_leaf_scripts(*records):
+            return PSBT(
+                g=PSBTMap({PSBT_GLOBAL_UNSIGNED_TX: tx.serialize()}),
+                i=[PSBTMap({
+                    bytes([PSBT_IN_TAP_LEAF_SCRIPT]) + control_block: bytes(leaf_script) + bytes([LEAF_VERSION_TAPSCRIPT])
+                    for leaf_script, control_block in records
+                })],
+                o=[PSBTMap({})],
+            ).to_base64()
+
+        def combined_tap_scripts(psbts):
+            return self.nodes[0].decodepsbt(self.nodes[0].combinepsbt(psbts))["inputs"][0]["taproot_scripts"]
+
+        def tap_script(leaf_script, control_blocks):
+            return {"script": leaf_script.hex(), "leaf_ver": LEAF_VERSION_TAPSCRIPT, "control_blocks": [cb.hex() for cb in control_blocks]}
+
+        control_block = bytes([LEAF_VERSION_TAPSCRIPT]) + bytes.fromhex(H_POINT)
+        control_block_with_path = control_block + bytes(32)
+        control_block_with_longer_path = control_block + bytes(64)
+        leaf_script_a = CScript([OP_TRUE])
+        leaf_script_b = CScript([OP_TRUE, OP_TRUE])
+        leaf_script_c = CScript([OP_TRUE, OP_TRUE, OP_TRUE])
+
+        psbt_a = psbt_with_leaf_scripts((leaf_script_a, control_block))
+        psbt_b = psbt_with_leaf_scripts((leaf_script_b, control_block))
+        psbt_c = psbt_with_leaf_scripts((leaf_script_c, control_block))
+
+        # The same control block under two leaf scripts would serialize as duplicate keys
+        assert_equal(combined_tap_scripts([psbt_a, psbt_b]), [tap_script(leaf_script_a, [control_block])])
+        # Reversed, so the leaf script kept is decided by the argument order and not by its content
+        assert_equal(combined_tap_scripts([psbt_b, psbt_a]), [tap_script(leaf_script_b, [control_block])])
+        # A third PSBT conflicting with what the first merge kept is dropped the same way
+        assert_equal(combined_tap_scripts([psbt_a, psbt_b, psbt_c]), [tap_script(leaf_script_a, [control_block])])
+        # Combining a PSBT with itself leaves it untouched
+        assert_equal(self.nodes[0].combinepsbt([psbt_a, psbt_a]), psbt_a)
+
+        # Records that do not conflict are all kept, whether or not they share a leaf script
+        psbt_same_leaf = psbt_with_leaf_scripts((leaf_script_a, control_block_with_path))
+        assert_equal(combined_tap_scripts([psbt_a, psbt_same_leaf]), [tap_script(leaf_script_a, [control_block, control_block_with_path])])
+        psbt_other_leaf = psbt_with_leaf_scripts((leaf_script_b, control_block_with_path))
+        assert_equal(combined_tap_scripts([psbt_a, psbt_other_leaf]), [
+            tap_script(leaf_script_a, [control_block]),
+            tap_script(leaf_script_b, [control_block_with_path]),
+        ])
+
+        # Only the conflicting control block of an incoming leaf script is dropped, not all of them
+        psbt_leaf_a_two_blocks = psbt_with_leaf_scripts((leaf_script_a, control_block), (leaf_script_a, control_block_with_path))
+        psbt_leaf_b_two_blocks = psbt_with_leaf_scripts((leaf_script_b, control_block_with_path), (leaf_script_b, control_block_with_longer_path))
+        assert_equal(combined_tap_scripts([psbt_leaf_a_two_blocks, psbt_leaf_b_two_blocks]), [
+            tap_script(leaf_script_a, [control_block, control_block_with_path]),
+            tap_script(leaf_script_b, [control_block_with_longer_path]),
+        ])
+
     def test_sighash_mismatch(self):
         self.log.info("Test sighash type mismatches")
         self.nodes[0].createwallet("sighash_mismatch")
@@ -1470,6 +1531,7 @@ def global_xpub_key(extended_pubkey):
 
         self.test_combinepsbt_preserves_proprietary_fields()
         self.test_combinepsbt_global_xpub_origin_conflict()
+        self.test_combinepsbt_tap_leaf_script_conflict()
 
         self.log.info("Test that combining PSBTs with different transactions fails")
         tx = CTransaction()