From 49aaee04232d6ac522ad5bedde79e1e618ead077 Mon Sep 17 00:00:00 2001
From: Marko Bencun <marko@shiftcrypto.ch>
Date: Sun, 11 Aug 2024 15:20:28 +0200
Subject: [PATCH] btc: add multisig support

---
 Cargo.lock                      |   2 +-
 Cargo.toml                      |   2 +-
 NPM_VERSION                     |   2 +-
 sandbox/src/Bitcoin.tsx         |  63 +++++++++++++-
 scripts/build-protos.rs         |   4 +
 src/btc.rs                      |  52 ++++++++++++
 src/shiftcrypto.bitbox02.rs     |   4 +
 src/wasm/types.rs               |   9 +-
 tests/subtests/test_btc_psbt.rs | 142 ++++++++++++++++++++++++++++++++
 9 files changed, 274 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 62d1195..ffa9522 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -162,7 +162,7 @@ checksum = "d965446196e3b7decd44aa7ee49e31d630118f90ef12f97900f262eb915c951d"
 
 [[package]]
 name = "bitbox-api"
-version = "0.4.1"
+version = "0.5.0"
 dependencies = [
  "async-trait",
  "base32",
diff --git a/Cargo.toml b/Cargo.toml
index 92febca..90aa336 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "bitbox-api"
 authors = ["Marko Bencun <benma@bitbox.swiss>"]
-version = "0.4.1"
+version = "0.5.0"
 homepage = "https://bitbox.swiss/"
 repository = "https://github.com/BitBoxSwiss/bitbox-api-rs/"
 readme = "README-rust.md"
diff --git a/NPM_VERSION b/NPM_VERSION
index 79a2734..09a3acf 100644
--- a/NPM_VERSION
+++ b/NPM_VERSION
@@ -1 +1 @@
-0.5.0
\ No newline at end of file
+0.6.0
\ No newline at end of file
diff --git a/sandbox/src/Bitcoin.tsx b/sandbox/src/Bitcoin.tsx
index 091bf36..8c7cee2 100644
--- a/sandbox/src/Bitcoin.tsx
+++ b/sandbox/src/Bitcoin.tsx
@@ -195,7 +195,7 @@ function BtcSignPSBT({ bb02 }: Props) {
   }
 
   return (
-    <div> 
+    <div>
       <h4>Sign PSBT</h4>
       <form className="verticalForm" onSubmit={submitForm}>
         <label>
@@ -295,7 +295,7 @@ function BtcSignMessage({ bb02 }: Props) {
         <button type='submit' disabled={running}>Sign message</button>
         {result ? (
           <div className="resultContainer">
-            <label>Result: 
+            <label>Result:
             {
               <textarea
                 rows={32}
@@ -313,6 +313,62 @@ function BtcSignMessage({ bb02 }: Props) {
 
 }
 
+function BtcMultisigAddress({ bb02 }: Props) {
+  const [running, setRunning] = useState(false);
+  const [result, setResult] = useState('');
+  const [err, setErr] = useState<bitbox.Error>();
+
+  const coin = 'tbtc';
+  const keypath = "m/48'/1'/0'/2'";
+  const someXPub = "tpubDFgycCkexSxkdZfeyaasDHityE97kiYM1BeCNoivDHvydGugKtoNobt4vEX6YSHNPy2cqmWQHKjKxciJuocepsGPGxcDZVmiMBnxgA1JKQk";
+
+  const submitForm = async (e: FormEvent) => {
+    e.preventDefault();
+    setRunning(true);
+    setResult('');
+    setErr(undefined);
+    try {
+      const ourXPub = await bb02.btcXpub(coin, keypath, 'tpub', false);
+      const scriptConfig: bitbox.BtcScriptConfig = {
+        multisig: {
+          threshold: 1,
+          xpubs: [ourXPub, someXPub],
+          ourXpubIndex: 0,
+          scriptType: 'p2wsh',
+        },
+      };
+      const is_script_config_registered = await bb02.btcIsScriptConfigRegistered(coin, scriptConfig, keypath);
+      if (!is_script_config_registered) {
+        await bb02.btcRegisterScriptConfig(coin, scriptConfig, keypath, 'autoXpubTpub', undefined);
+      }
+      const address = await bb02.btcAddress(coin, keypath + "/0/10", scriptConfig, true);
+      setResult(address);
+    } catch (err) {
+      setErr(bitbox.ensureError(err));
+    } finally {
+      setRunning(false);
+    }
+  }
+
+  return (
+    <div>
+      <h4>Multisig</h4>
+      <form className="verticalForm" onSubmit={submitForm}>
+        Address for a multisig wallet using the BitBox02 xpub at
+        <pre>{keypath}</pre>
+        <p>and some other arbitrary xpub: <code>{someXPub}</code></p>
+        <button type='submit' disabled={running}>Multisig address</button>
+        {result ? (
+            <div className="resultContainer">
+              <label>Result: <b><code>{result}</code></b></label>
+            </div>
+        ) : null }
+        <ShowError err={err} />
+      </form>
+    </div>
+  );
+}
+
 function BtcMiniscriptAddress({ bb02 }: Props) {
   const [running, setRunning] = useState(false);
   const [result, setResult] = useState('');
@@ -391,6 +447,9 @@ export function Bitcoin({ bb02 } : Props) {
       <div className="action">
         <BtcSignMessage bb02={bb02} />
       </div>
+      <div className="action">
+        <BtcMultisigAddress bb02={bb02} />
+      </div>
       <div className="action">
         <BtcMiniscriptAddress bb02={bb02} />
       </div>
diff --git a/scripts/build-protos.rs b/scripts/build-protos.rs
index 2a4c951..8547470 100644
--- a/scripts/build-protos.rs
+++ b/scripts/build-protos.rs
@@ -43,6 +43,10 @@ fn add_serde_attrs(c: &mut prost_build::Config) {
             "shiftcrypto.bitbox02.BTCScriptConfig.config.simple_type",
             "serde(deserialize_with = \"crate::btc::serde_deserialize_simple_type\")",
         ),
+        (
+            "shiftcrypto.bitbox02.BTCScriptConfig.config.multisig",
+            "serde(deserialize_with = \"crate::btc::serde_deserialize_multisig\")",
+        ),
         (
             "shiftcrypto.bitbox02.RootFingerprintResponse.fingerprint",
             "serde(deserialize_with = \"hex::serde::deserialize\")",
diff --git a/src/btc.rs b/src/btc.rs
index 5ea10be..4c63c36 100644
--- a/src/btc.rs
+++ b/src/btc.rs
@@ -25,6 +25,39 @@ where
     Ok(pb::btc_script_config::SimpleType::deserialize(deserializer)?.into())
 }
 
+#[cfg(feature = "wasm")]
+pub(crate) fn serde_deserialize_multisig<'de, D>(
+    deserializer: D,
+) -> Result<pb::btc_script_config::Multisig, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    use serde::Deserialize;
+    use std::str::FromStr;
+
+    #[derive(serde::Deserialize)]
+    #[serde(rename_all = "camelCase")]
+    struct Multisig {
+        threshold: u32,
+        xpubs: Vec<String>,
+        our_xpub_index: u32,
+        script_type: pb::btc_script_config::multisig::ScriptType,
+    }
+    let ms = Multisig::deserialize(deserializer)?;
+    let xpubs = ms
+        .xpubs
+        .iter()
+        .map(|s| Xpub::from_str(s.as_str()))
+        .collect::<Result<Vec<Xpub>, _>>()
+        .map_err(serde::de::Error::custom)?;
+    Ok(pb::btc_script_config::Multisig {
+        threshold: ms.threshold,
+        xpubs: xpubs.iter().map(convert_xpub).collect(),
+        our_xpub_index: ms.our_xpub_index,
+        script_type: ms.script_type.into(),
+    })
+}
+
 #[cfg(feature = "wasm")]
 #[derive(serde::Deserialize)]
 pub(crate) struct SerdeScriptConfig(pb::btc_script_config::Config);
@@ -500,6 +533,25 @@ impl From<KeyOriginInfo> for pb::KeyOriginInfo {
     }
 }
 
+/// Create a multi-sig script config.
+pub fn make_script_config_multisig(
+    threshold: u32,
+    xpubs: &[bitcoin::bip32::Xpub],
+    our_xpub_index: u32,
+    script_type: pb::btc_script_config::multisig::ScriptType,
+) -> pb::BtcScriptConfig {
+    pb::BtcScriptConfig {
+        config: Some(pb::btc_script_config::Config::Multisig(
+            pb::btc_script_config::Multisig {
+                threshold,
+                xpubs: xpubs.iter().map(convert_xpub).collect(),
+                our_xpub_index,
+                script_type: script_type as _,
+            },
+        )),
+    }
+}
+
 /// Create a wallet policy script config according to the wallet policies BIP:
 /// <https://github.com/bitcoin/bips/pull/1389>
 ///
diff --git a/src/shiftcrypto.bitbox02.rs b/src/shiftcrypto.bitbox02.rs
index 50a49c1..bd320fd 100644
--- a/src/shiftcrypto.bitbox02.rs
+++ b/src/shiftcrypto.bitbox02.rs
@@ -411,6 +411,10 @@ pub mod btc_script_config {
         )]
         SimpleType(i32),
         #[prost(message, tag = "2")]
+        #[cfg_attr(
+            feature = "wasm",
+            serde(deserialize_with = "crate::btc::serde_deserialize_multisig")
+        )]
         Multisig(Multisig),
         #[prost(message, tag = "3")]
         Policy(Policy),
diff --git a/src/wasm/types.rs b/src/wasm/types.rs
index 7c9d8ad..5cdb0dc 100644
--- a/src/wasm/types.rs
+++ b/src/wasm/types.rs
@@ -32,8 +32,15 @@ type KeyOriginInfo = {
   xpub: XPub;
 };
 type BtcRegisterXPubType = 'autoElectrum' | 'autoXpubTpub';
+type BtcMultisigScriptType = 'p2wsh' | 'p2wshP2sh';
+type BtcMultisig = {
+  threshold: number;
+  xpubs: XPub[];
+  ourXpubIndex: number;
+  scriptType: BtcMultisigScriptType;
+};
 type BtcPolicy = { policy: string; keys: KeyOriginInfo[] };
-type BtcScriptConfig = { simpleType: BtcSimpleType; } | { policy: BtcPolicy };
+type BtcScriptConfig = { simpleType: BtcSimpleType; } | { multisig: BtcMultisig } | { policy: BtcPolicy };
 type BtcScriptConfigWithKeypath = {
   scriptConfig: BtcScriptConfig;
   keypath: Keypath;
diff --git a/tests/subtests/test_btc_psbt.rs b/tests/subtests/test_btc_psbt.rs
index 4bb3ca0..4e1f199 100644
--- a/tests/subtests/test_btc_psbt.rs
+++ b/tests/subtests/test_btc_psbt.rs
@@ -55,6 +55,7 @@ fn verify_transaction(psbt: Psbt) {
 pub async fn test(bitbox: &PairedBitBox) {
     test_taproot_key_spend(bitbox).await;
     test_mixed_spend(bitbox).await;
+    test_policy_multisig_p2wsh(bitbox).await;
     test_policy_wsh(bitbox).await;
 }
 
@@ -325,6 +326,147 @@ async fn test_mixed_spend(bitbox: &PairedBitBox) {
     verify_transaction(psbt);
 }
 
+async fn test_policy_multisig_p2wsh(bitbox: &PairedBitBox) {
+    let secp = secp256k1::Secp256k1::new();
+
+    let coin = pb::BtcCoin::Tbtc;
+
+    let our_root_fingerprint = super::simulator_xprv().fingerprint(&secp);
+
+    let threshold: u32 = 1;
+    let keypath_account: DerivationPath = "m/48'/1'/0'/2'".parse().unwrap();
+
+    let our_xpub: Xpub = super::simulator_xpub_at(&secp, &keypath_account);
+    let some_xpub: Xpub = "tpubDFgycCkexSxkdZfeyaasDHityE97kiYM1BeCNoivDHvydGugKtoNobt4vEX6YSHNPy2cqmWQHKjKxciJuocepsGPGxcDZVmiMBnxgA1JKQk".parse().unwrap();
+
+    // We use the miniscript library to build a multipath descriptor including key origin so we can
+    // easily derive the receive/change descriptor, pubkey scripts, populate the PSBT input key
+    // infos and convert the sigs to final witnesses.
+
+    let multi_descriptor: miniscript::Descriptor<miniscript::DescriptorPublicKey> = format!(
+        "wsh(sortedmulti({},[{}/48'/1'/0'/2']{}/<0;1>/*,{}/<0;1>/*))",
+        threshold, our_root_fingerprint, &our_xpub, &some_xpub
+    )
+    .parse::<miniscript::Descriptor<miniscript::DescriptorPublicKey>>()
+    .unwrap();
+    assert!(multi_descriptor.sanity_check().is_ok());
+
+    let [descriptor_receive, descriptor_change] = multi_descriptor
+        .into_single_descriptors()
+        .unwrap()
+        .try_into()
+        .unwrap();
+    // Derive /0/0 (first receive) and /1/0 (first change) descriptors.
+    let input_descriptor = descriptor_receive.at_derivation_index(0).unwrap();
+    let change_descriptor = descriptor_change.at_derivation_index(0).unwrap();
+    let multisig_config = bitbox_api::btc::make_script_config_multisig(
+        threshold,
+        &[our_xpub, some_xpub],
+        0,
+        pb::btc_script_config::multisig::ScriptType::P2wsh,
+    );
+
+    // Register policy if not already registered. This must be done before any receive address is
+    // created or any transaction is signed.
+    let is_registered = bitbox
+        .btc_is_script_config_registered(coin, &multisig_config, None)
+        .await
+        .unwrap();
+
+    if !is_registered {
+        bitbox
+            .btc_register_script_config(
+                coin,
+                &multisig_config,
+                Some(&(&keypath_account).into()),
+                pb::btc_register_script_config_request::XPubType::AutoXpubTpub,
+                Some("test wsh multisig"),
+            )
+            .await
+            .unwrap();
+    }
+
+    // A previous tx which creates some UTXOs we can reference later.
+    let prev_tx = Transaction {
+        version: transaction::Version::TWO,
+        lock_time: bitcoin::absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: "3131313131313131313131313131313131313131313131313131313131313131:0"
+                .parse()
+                .unwrap(),
+            script_sig: ScriptBuf::new(),
+            sequence: Sequence(0xFFFFFFFF),
+            witness: Witness::default(),
+        }],
+        output: vec![TxOut {
+            value: Amount::from_sat(100_000_000),
+            script_pubkey: input_descriptor.script_pubkey(),
+        }],
+    };
+
+    let tx = Transaction {
+        version: transaction::Version::TWO,
+        lock_time: bitcoin::absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint {
+                txid: prev_tx.compute_txid(),
+                vout: 0,
+            },
+            script_sig: ScriptBuf::new(),
+            sequence: Sequence(0xFFFFFFFF),
+            witness: Witness::default(),
+        }],
+        output: vec![
+            TxOut {
+                value: Amount::from_sat(70_000_000),
+                script_pubkey: change_descriptor.script_pubkey(),
+            },
+            TxOut {
+                value: Amount::from_sat(20_000_000),
+                script_pubkey: ScriptBuf::new_p2tr(
+                    &secp,
+                    // random private key:
+                    // 9dbb534622a6100a39b73dece43c6d4db14b9a612eb46a6c64c2bb849e283ce8
+                    "e4adbb12c3426ec71ebb10688d8ae69d531ca822a2b790acee216a7f1b95b576"
+                        .parse()
+                        .unwrap(),
+                    None,
+                ),
+            },
+        ],
+    };
+
+    let mut psbt = Psbt::from_unsigned_tx(tx).unwrap();
+
+    // Add input and change infos.
+    psbt.inputs[0].non_witness_utxo = Some(prev_tx.clone());
+    // These add the input/output bip32_derivation entries / key infos.
+    psbt.update_input_with_descriptor(0, &input_descriptor)
+        .unwrap();
+    psbt.update_output_with_descriptor(0, &change_descriptor)
+        .unwrap();
+
+    // Sign.
+    bitbox
+        .btc_sign_psbt(
+            pb::BtcCoin::Tbtc,
+            &mut psbt,
+            Some(pb::BtcScriptConfigWithKeypath {
+                script_config: Some(multisig_config),
+                keypath: keypath_account.to_u32_vec(),
+            }),
+            pb::btc_sign_init_request::FormatUnit::Default,
+        )
+        .await
+        .unwrap();
+
+    // Finalize, add witnesses.
+    psbt.finalize_mut(&secp).unwrap();
+
+    // Verify the signed tx, including that all sigs/witnesses are correct.
+    verify_transaction(psbt);
+}
+
 async fn test_policy_wsh(bitbox: &PairedBitBox) {
     let secp = secp256k1::Secp256k1::new();