wallet_musig.py raw

   1  #!/usr/bin/env python3
   2  # Copyright (c) 2024-present The Bitcoin Core developers
   3  # Distributed under the MIT software license, see the accompanying
   4  # file COPYING or http://www.opensource.org/licenses/mit-license.php.
   5  
   6  import re
   7  
   8  from test_framework.descriptors import descsum_create
   9  from test_framework.key import H_POINT
  10  from test_framework.script import hash160
  11  from test_framework.test_framework import BitcoinTestFramework
  12  from test_framework.util import (
  13      assert_equal,
  14      assert_greater_than,
  15      assert_not_equal,
  16  )
  17  
  18  PRIVKEY_RE = re.compile(r"^tr\((.+?)/.+\)#.{8}$")
  19  PUBKEY_RE = re.compile(r"^tr\((\[.+?\].+?)/.+\)#.{8}$")
  20  ORIGIN_PATH_RE = re.compile(r"^\[\w{8}(/.*)\].*$")
  21  MULTIPATH_TWO_RE = re.compile(r"<(\d+);(\d+)>")
  22  MUSIG_RE = re.compile(r"musig\((.*?)\)")
  23  PLACEHOLDER_RE = re.compile(r"\$\d")
  24  
  25  class WalletMuSigTest(BitcoinTestFramework):
  26      wallet_num = 0
  27      def set_test_params(self):
  28          self.num_nodes = 1
  29  
  30      def skip_test_if_missing_module(self):
  31          self.skip_if_no_wallet()
  32  
  33      # Create wallets and extract keys
  34      def create_wallets_and_keys_from_pattern(self, pat):
  35          wallets = []
  36          keys = []
  37  
  38          for musig in MUSIG_RE.findall(pat):
  39              for placeholder in PLACEHOLDER_RE.findall(musig):
  40                  wallet_index = int(placeholder[1:])
  41                  if wallet_index < len(wallets):
  42                      continue
  43  
  44                  wallet_name = f"musig_{self.wallet_num}"
  45                  self.wallet_num += 1
  46                  self.nodes[0].createwallet(wallet_name)
  47                  wallet = self.nodes[0].get_wallet_rpc(wallet_name)
  48                  wallets.append(wallet)
  49  
  50                  for priv_desc in wallet.listdescriptors(True)["descriptors"]:
  51                      desc = priv_desc["desc"]
  52                      if not desc.startswith("tr("):
  53                          continue
  54                      privkey = PRIVKEY_RE.search(desc).group(1)
  55                      break
  56                  for pub_desc in wallet.listdescriptors()["descriptors"]:
  57                      desc = pub_desc["desc"]
  58                      if not desc.startswith("tr("):
  59                          continue
  60                      pubkey = PUBKEY_RE.search(desc).group(1)
  61                      # Since the pubkey is derived from the private key that we have, we need
  62                      # to extract and insert the origin path from the pubkey as well.
  63                      privkey += ORIGIN_PATH_RE.search(pubkey).group(1)
  64                      break
  65                  keys.append((privkey, pubkey))
  66  
  67          return wallets, keys
  68  
  69      # Construct and import each wallet's musig descriptor that
  70      # contains the private key from that wallet and pubkeys of the others
  71      def construct_and_import_musig_descriptor_in_wallets(self, pat, wallets, keys, only_one_musig_wallet=False):
  72          for i, wallet in enumerate(wallets):
  73              if only_one_musig_wallet and i > 0:
  74                  continue
  75              desc = pat
  76              for j, (priv, pub) in enumerate(keys):
  77                  if j == i:
  78                      desc = desc.replace(f"${i}", priv)
  79                  else:
  80                      desc = desc.replace(f"${j}", pub)
  81  
  82              import_descs = [{
  83                  "desc": descsum_create(desc),
  84                  "active": True,
  85                  "timestamp": "now",
  86              }]
  87  
  88              res = wallet.importdescriptors(import_descs)
  89              for r in res:
  90                  assert_equal(r["success"], True)
  91  
  92      def setup_musig_scenario(self, pat):
  93          wallets, keys = self.create_wallets_and_keys_from_pattern(pat)
  94          self.construct_and_import_musig_descriptor_in_wallets(pat, wallets, keys, only_one_musig_wallet=False)
  95  
  96          # Fund address
  97          addr = wallets[0].getnewaddress(address_type="bech32m")
  98          for wallet in wallets[1:]:
  99              assert_equal(addr, wallet.getnewaddress(address_type="bech32m"))
 100  
 101          self.def_wallet.sendtoaddress(addr, 10)
 102          self.generate(self.nodes[0], 1)
 103  
 104          # Create PSBT
 105          utxo = wallets[0].listunspent()[0]
 106          psbt = wallets[0].walletcreatefundedpsbt(
 107              outputs=[{self.def_wallet.getnewaddress(): 5}],
 108              inputs=[utxo],
 109              change_type="bech32m",
 110              changePosition=1
 111          )["psbt"]
 112  
 113          return wallets, psbt
 114  
 115      def assert_musig_signer_data(self, first, second, different_field):
 116          assert_equal(first["participant_pubkey"], second["participant_pubkey"])
 117          assert_equal(first["aggregate_pubkey"], second["aggregate_pubkey"])
 118          if "leaf_hash" in first:
 119              assert_equal(first["leaf_hash"], second["leaf_hash"])
 120          else:
 121              assert "leaf_hash" not in second
 122  
 123          assert_not_equal(first[different_field], second[different_field])
 124  
 125      def assert_musig_aggregate_in_script(self, signer_data, pattern, psbtin):
 126          pubkey = signer_data["aggregate_pubkey"][2:]
 127          if "pkh" in pattern or "pk_h" in pattern:
 128              pubkey = hash160(bytes.fromhex(pubkey)).hex()
 129          if pubkey in psbtin["witness_utxo"]["scriptPubKey"]["hex"]:
 130              return
 131          elif "taproot_scripts" in psbtin:
 132              for leaf_scripts in psbtin["taproot_scripts"]:
 133                  if pubkey in leaf_scripts["script"]:
 134                      break
 135              else:
 136                  assert False, "Aggregate pubkey not seen as output key, or in any scripts"
 137          else:
 138              assert False, "Aggregate pubkey not seen as output key or internal key"
 139  
 140      def test_failure_case_1(self, comment, pat):
 141          self.log.info(f"Testing {comment}")
 142          wallets, psbt = self.setup_musig_scenario(pat)
 143  
 144          # Only 2 out of 3 participants provide nonces
 145          nonce_psbts = []
 146          for i in range(2):
 147              proc = wallets[i].walletprocesspsbt(psbt=psbt)
 148              nonce_psbts.append(proc["psbt"])
 149  
 150          comb_nonce_psbt = self.nodes[0].combinepsbt(nonce_psbts)
 151  
 152          # Attempt to create partial sigs. This should not complete due to the
 153          # missing nonce.
 154          for wallet in wallets[:2]:
 155              proc = wallet.walletprocesspsbt(psbt=comb_nonce_psbt)
 156              assert_equal(proc["complete"], False)
 157              # No partial sigs are created
 158              dec = self.nodes[0].decodepsbt(proc["psbt"])
 159              # There are still only two nonces
 160              assert_equal(len(dec["inputs"][0].get("musig2_pubnonces", [])), 2)
 161  
 162      def test_failure_case_2(self, comment, pat):
 163          self.log.info(f"Testing {comment}")
 164          wallets, psbt = self.setup_musig_scenario(pat)
 165          nonce_psbts = [w.walletprocesspsbt(psbt=psbt)["psbt"] for w in wallets]
 166          comb_nonce_psbt = self.nodes[0].combinepsbt(nonce_psbts)
 167  
 168          # Only 2 out of 3 provide partial sigs
 169          psig_psbts = []
 170          for i in range(2):
 171              proc = wallets[i].walletprocesspsbt(psbt=comb_nonce_psbt)
 172              psig_psbts.append(proc["psbt"])
 173  
 174          comb_psig_psbt = self.nodes[0].combinepsbt(psig_psbts)
 175  
 176          # Finalization fails due to missing partial sig
 177          finalized = self.nodes[0].finalizepsbt(comb_psig_psbt)
 178          assert_equal(finalized["complete"], False)
 179  
 180          # Still only two partial sigs in combined PSBT
 181          dec = self.nodes[0].decodepsbt(comb_psig_psbt)
 182          assert_equal(len(dec["inputs"][0]["musig2_partial_sigs"]), 2)
 183  
 184      def test_failure_case_3(self, comment, pat):
 185          self.log.info(f"Testing {comment}")
 186          wallets, psbt = self.setup_musig_scenario(pat)
 187          nonce_psbts = [w.walletprocesspsbt(psbt=psbt)["psbt"] for w in wallets]
 188          comb_nonce_psbt = self.nodes[0].combinepsbt(nonce_psbts)
 189  
 190          finalized = self.nodes[0].finalizepsbt(comb_nonce_psbt)
 191          assert_equal(finalized["complete"], False)
 192  
 193          dec = self.nodes[0].decodepsbt(comb_nonce_psbt)
 194          assert "musig2_pubnonces" in dec["inputs"][0]
 195          assert "musig2_partial_sigs" not in dec["inputs"][0]
 196  
 197      def test_success_case(self, comment, pattern, sighash_type=None, scriptpath=False, nosign_wallets=None, only_one_musig_wallet=False):
 198          self.log.info(f"Testing {comment}")
 199          has_internal = MULTIPATH_TWO_RE.search(pattern) is not None
 200  
 201          pat = pattern.replace("$H", H_POINT)
 202          wallets, keys = self.create_wallets_and_keys_from_pattern(pat)
 203          self.construct_and_import_musig_descriptor_in_wallets(pat, wallets, keys, only_one_musig_wallet)
 204  
 205          expected_pubnonces = 0
 206          expected_partial_sigs = 0
 207          for musig in MUSIG_RE.findall(pat):
 208              musig_partial_sigs = 0
 209              for placeholder in PLACEHOLDER_RE.findall(musig):
 210                  wallet_index = int(placeholder[1:])
 211                  if nosign_wallets is None or wallet_index not in nosign_wallets:
 212                      expected_pubnonces += 1
 213                  else:
 214                      musig_partial_sigs = None
 215                  if musig_partial_sigs is not None:
 216                      musig_partial_sigs += 1
 217                  if wallet_index < len(wallets):
 218                      continue
 219              if musig_partial_sigs is not None:
 220                  expected_partial_sigs += musig_partial_sigs
 221  
 222          # Check that the wallets agree on the same musig address
 223          addr = None
 224          change_addr = None
 225          for i, wallet in enumerate(wallets):
 226              if only_one_musig_wallet and i > 0:
 227                  continue
 228              if addr is None:
 229                  addr = wallet.getnewaddress(address_type="bech32m")
 230              else:
 231                  assert_equal(addr, wallet.getnewaddress(address_type="bech32m"))
 232              if has_internal:
 233                  if change_addr is None:
 234                      change_addr = wallet.getrawchangeaddress(address_type="bech32m")
 235                  else:
 236                      assert_equal(change_addr, wallet.getrawchangeaddress(address_type="bech32m"))
 237  
 238          # Fund that address
 239          self.def_wallet.sendtoaddress(addr, 10)
 240          self.generate(self.nodes[0], 1)
 241  
 242          # Spend that UTXO
 243          utxo = None
 244          for i, wallet in enumerate(wallets):
 245              if only_one_musig_wallet and i > 0:
 246                  continue
 247              if utxo is None:
 248                  utxo = wallet.listunspent()[0]
 249              else:
 250                  assert_equal(utxo, wallet.listunspent()[0])
 251          psbt = wallets[0].walletcreatefundedpsbt(outputs=[{self.def_wallet.getnewaddress(): 5}], inputs=[utxo], change_type="bech32m", changePosition=1, locktime=self.nodes[0].getblockcount())["psbt"]
 252  
 253          dec_psbt = self.nodes[0].decodepsbt(psbt)
 254          assert_equal(len(dec_psbt["inputs"]), 1)
 255          assert_equal(len(dec_psbt["inputs"][0]["musig2_participant_pubkeys"]), pattern.count("musig("))
 256          if has_internal:
 257              assert_equal(len(dec_psbt["outputs"][1]["musig2_participant_pubkeys"]), pattern.count("musig("))
 258  
 259          # Check all participant pubkeys in the input and change output
 260          psbt_maps = [dec_psbt["inputs"][0]]
 261          if has_internal:
 262              psbt_maps.append(dec_psbt["outputs"][1])
 263          for psbt_map in psbt_maps:
 264              part_pks = set()
 265              for agg in psbt_map["musig2_participant_pubkeys"]:
 266                  for part_pub in agg["participant_pubkeys"]:
 267                      part_pks.add(part_pub[2:])
 268              # Check that there are as many participants as we expected
 269              assert_equal(len(part_pks), len(keys))
 270              # Check that each participant has a derivation path
 271              for deriv_path in psbt_map["taproot_bip32_derivs"]:
 272                  if deriv_path["pubkey"] in part_pks:
 273                      part_pks.remove(deriv_path["pubkey"])
 274              assert_equal(len(part_pks), 0)
 275  
 276          # Run 2 signing sessions simultaneously to verify no nonce reuse
 277          # Add pubnonces
 278          nonce_psbts = []
 279          nonce_psbts2 = []
 280          for i, wallet in enumerate(wallets):
 281              if nosign_wallets and i in nosign_wallets:
 282                  continue
 283              for psbt_list in [nonce_psbts, nonce_psbts2]:
 284                  proc = wallet.walletprocesspsbt(psbt=psbt, sighashtype=sighash_type)
 285                  assert_equal(proc["complete"], False)
 286                  psbt_list.append(proc["psbt"])
 287  
 288          comb_nonce_psbt = self.nodes[0].combinepsbt(nonce_psbts)
 289          comb_nonce_psbt2 = self.nodes[0].combinepsbt(nonce_psbts2)
 290  
 291          dec_psbt = self.nodes[0].decodepsbt(comb_nonce_psbt)
 292          dec_psbt2 = self.nodes[0].decodepsbt(comb_nonce_psbt2)
 293          assert_equal(len(dec_psbt["inputs"][0]["musig2_pubnonces"]), len(dec_psbt2["inputs"][0]["musig2_pubnonces"]), expected_pubnonces)
 294          for pn, pn2 in zip(dec_psbt["inputs"][0]["musig2_pubnonces"], dec_psbt2["inputs"][0]["musig2_pubnonces"]):
 295              self.assert_musig_signer_data(pn, pn2, "pubnonce")
 296              self.assert_musig_aggregate_in_script(pn, pattern, dec_psbt["inputs"][0])
 297  
 298          # Add partial sigs
 299          psig_psbts = []
 300          psig_psbts2 = []
 301          for i, wallet in enumerate(wallets):
 302              if nosign_wallets and i in nosign_wallets:
 303                  continue
 304              for psbt, psbt_list in [(comb_nonce_psbt, psig_psbts), (comb_nonce_psbt2, psig_psbts2)]:
 305                  proc = wallet.walletprocesspsbt(psbt=psbt, sighashtype=sighash_type)
 306                  assert_equal(proc["complete"], False)
 307                  psbt_list.append(proc["psbt"])
 308  
 309          comb_psig_psbt = self.nodes[0].combinepsbt(psig_psbts)
 310          comb_psig_psbt2 = self.nodes[0].combinepsbt(psig_psbts2)
 311  
 312          dec_psbt = self.nodes[0].decodepsbt(comb_psig_psbt)
 313          dec_psbt2 = self.nodes[0].decodepsbt(comb_psig_psbt2)
 314          assert_equal(len(dec_psbt["inputs"][0]["musig2_partial_sigs"]), len(dec_psbt2["inputs"][0]["musig2_partial_sigs"]), expected_partial_sigs)
 315          for ps, ps2 in zip(dec_psbt["inputs"][0]["musig2_partial_sigs"], dec_psbt2["inputs"][0]["musig2_partial_sigs"]):
 316              self.assert_musig_signer_data(ps, ps2, "partial_sig")
 317              self.assert_musig_aggregate_in_script(ps, pattern, dec_psbt["inputs"][0])
 318  
 319          # Non-participant aggregates partial sigs and send
 320          finalized = self.nodes[0].finalizepsbt(psbt=comb_psig_psbt, extract=False)
 321          finalized2 = self.nodes[0].finalizepsbt(psbt=comb_psig_psbt2, extract=False)
 322          assert_equal(finalized["complete"], finalized2["complete"], True)
 323          witness = self.nodes[0].decodepsbt(finalized["psbt"])["inputs"][0]["final_scriptwitness"]
 324          assert_not_equal(witness, self.nodes[0].decodepsbt(finalized2["psbt"])["inputs"][0]["final_scriptwitness"])
 325          if scriptpath:
 326              assert_greater_than(len(witness), 1)
 327          else:
 328              assert_equal(len(witness), 1)
 329          finalized = self.nodes[0].finalizepsbt(comb_psig_psbt)
 330          assert "hex" in finalized
 331          self.nodes[0].sendrawtransaction(finalized["hex"])
 332  
 333      def run_test(self):
 334          self.def_wallet = self.nodes[0].get_wallet_rpc(self.default_wallet_name)
 335  
 336          self.test_success_case("rawtr(musig(keys/*))", "rawtr(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*))")
 337          self.test_success_case("rawtr(musig(keys/*)) with ALL|ANYONECANPAY", "rawtr(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*))", "ALL|ANYONECANPAY")
 338          self.test_success_case("tr(musig(keys/*)) no multipath", "tr(musig($0/0/*,$1/1/*,$2/2/*))")
 339          self.test_success_case("tr(musig(keys/*)) 2 index multipath", "tr(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*))")
 340          self.test_success_case("tr(musig(keys/*)) 3 index multipath", "tr(musig($0/<0;1;2>/*,$1/<1;2;3>/*,$2/<2;3;4>/*))")
 341          self.test_success_case("rawtr(musig/*)", "rawtr(musig($0,$1,$2)/<0;1>/*)")
 342          self.test_success_case("tr(musig/*)", "tr(musig($0,$1,$2)/<0;1>/*)")
 343          self.test_success_case("rawtr(musig(keys/*)) without all wallets importing", "rawtr(musig($0/<0;1>/*,$1/<0;1>/*,$2/<0;1>/*))", only_one_musig_wallet=True)
 344          self.test_success_case("tr(musig(keys/*)) without all wallets importing", "tr(musig($0/<0;1>/*,$1/<0;1>/*,$2/<0;1>/*))", only_one_musig_wallet=True)
 345          self.test_success_case("tr(H, pk(musig(keys/*)))", "tr($H,pk(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*)))", scriptpath=True)
 346          self.test_success_case("tr(H,pk(musig/*))", "tr($H,pk(musig($0,$1,$2)/<0;1>/*))", scriptpath=True)
 347          self.test_success_case("tr(H,{pk(musig/*), pk(musig/*)})", "tr($H,{pk(musig($0,$1,$2)/<0;1>/*),pk(musig($3,$4,$5)/0/*)})", scriptpath=True)
 348          self.test_success_case("tr(H,{pk(musig/*), pk(same keys different musig/*)})", "tr($H,{pk(musig($0,$1,$2)/<0;1>/*),pk(musig($1,$2)/0/*)})", scriptpath=True)
 349          self.test_success_case("tr(musig/*,{pk(partial keys diff musig-1/*),pk(partial keys diff musig-2/*)})}", "tr(musig($0,$1,$2)/<3;4>/*,{pk(musig($0,$1)/<5;6>/*),pk(musig($1,$2)/7/*)})")
 350          self.test_success_case("tr(musig/*,{pk(partial keys diff musig-1/*),pk(partial keys diff musig-2/*)})} script-path", "tr(musig($0,$1,$2)/<3;4>/*,{pk(musig($0,$1)/<5;6>/*),pk(musig($1,$2)/7/*)})", scriptpath=True, nosign_wallets=[0])
 351          self.test_success_case("tr(H,and(pk(musig/*),after(1)))", "tr($H,and_v(v:pk(musig($0,$1,$2)/<0;1>/*),after(1)))", scriptpath=True)
 352          self.test_success_case("tr(H,and(pk_k(musig/*),after(1)))", "tr($H,and_v(vc:pk_k(musig($0,$1,$2)/<0;1>/*),after(1)))", scriptpath=True)
 353          self.test_success_case("tr(H,and(pkh(musig/*),after(1)))", "tr($H,and_v(v:pkh(musig($0,$1,$2)/<0;1>/*),after(1)))", scriptpath=True)
 354          self.test_success_case("tr(H,and(pk_h(musig/*),after(1)))", "tr($H,and_v(vc:pk_h(musig($0,$1,$2)/<0;1>/*),after(1)))", scriptpath=True)
 355          self.test_success_case("tr(H,{and(pk(musig/*),after(1)),and(pk(musig/*),after(1))})", "tr($H,{and_v(v:pk(musig($0,$2)/0/*),after(1)),and_v(v:pk(musig($1,$2)/0/*),after(1))})", scriptpath=True)
 356  
 357          self.test_failure_case_1("missing participant nonce", "tr(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*))")
 358          self.test_failure_case_2("insufficient partial signatures", "rawtr(musig($0/<0;1>/*,$1/<1;2>/*,$2/<2;3>/*))")
 359          self.test_failure_case_3("finalize without partial sigs", "rawtr(musig($0/<0;1>/*,$1/<1;2>/*))")
 360  
 361  if __name__ == '__main__':
 362      WalletMuSigTest(__file__).main()
 363