btc: support Taproot policies

Cargo.toml

@@ -81,7 +81,7 @@ required-features = ["usb", "tokio/rt", "tokio/macros"]
 
 [[example]]
 name = "btc_miniscript"
-required-features = ["usb", "tokio/rt", "tokio/macros"]
+required-features = ["simulator", "tokio/rt", "tokio/macros"]
 
 [[example]]
 name = "eth"

Makefile

@@ -13,7 +13,7 @@ example-btc-psbt:
 example-btc-sign-msg:
 	cargo run --example btc_sign_msg --features=usb,tokio/rt,tokio/macros
 example-btc-miniscript:
-	cargo run --example btc_miniscript --features=usb,tokio/rt,tokio/macros
+	cargo run --example btc_miniscript --features=tokio/rt,tokio/macros,simulator
 example-eth:
 	cargo run --example eth --features=usb,tokio/rt,tokio/macros,rlp
 example-cardano:

src/btc.rs

@@ -248,13 +248,19 @@ pub enum PsbtError {
 enum OurKey {
     Segwit(bitcoin::secp256k1::PublicKey, Keypath),
     TaprootInternal(Keypath),
+    TaprootScript(
+        bitcoin::secp256k1::XOnlyPublicKey,
+        bitcoin::taproot::TapLeafHash,
+        Keypath,
+    ),
 }
 
 impl OurKey {
     fn keypath(&self) -> Keypath {
         match self {
             OurKey::Segwit(_, kp) => kp.clone(),
             OurKey::TaprootInternal(kp) => kp.clone(),
+            OurKey::TaprootScript(_, _, kp) => kp.clone(),
         }
     }
 }
@@ -328,19 +334,34 @@ fn find_our_key<T: PsbtOutputInfo>(
     our_root_fingerprint: &[u8],
     output_info: T,
 ) -> Result<OurKey, PsbtError> {
-    if let Some(tap_internal_key) = output_info.get_tap_internal_key() {
-        let (leaf_hashes, (fingerprint, derivation_path)) = output_info
-            .get_tap_key_origins()
-            .get(tap_internal_key)
-            .ok_or(PsbtError::KeyNotFound)?;
-        if !leaf_hashes.is_empty() {
-            return Err(PsbtError::UnsupportedTapScript);
-        }
+    for (xonly, (leaf_hashes, (fingerprint, derivation_path))) in
+        output_info.get_tap_key_origins().iter()
+    {
         if &fingerprint[..] == our_root_fingerprint {
             // TODO: check for fingerprint collision
-            return Ok(OurKey::TaprootInternal(derivation_path.into()));
+
+            if let Some(tap_internal_key) = output_info.get_tap_internal_key() {
+                if tap_internal_key == xonly {
+                    if !leaf_hashes.is_empty() {
+                        // TODO change err msg, we don't support the
+                        // same key as internal key and also in a leaf
+                        // script.
+                        return Err(PsbtError::UnsupportedTapScript);
+                    }
+                    return Ok(OurKey::TaprootInternal(derivation_path.into()));
+                }
+            }
+            if leaf_hashes.len() != 1 {
+                // TODO change err msg, per BIP-388 all pubkeys are
+                // unique, so it can't be in multiple leafs.
+                return Err(PsbtError::UnsupportedTapScript);
+            }
+            return Ok(OurKey::TaprootScript(
+                *xonly,
+                leaf_hashes[0],
+                derivation_path.into(),
+            ));
         }
-        return Err(PsbtError::KeyNotFound);
     }
     for (pubkey, (fingerprint, derivation_path)) in output_info.get_bip32_derivation().iter() {
         if &fingerprint[..] == our_root_fingerprint {
@@ -568,15 +589,26 @@ pub fn make_script_config_policy(policy: &str, keys: &[KeyOriginInfo]) -> pb::Bt
     }
 }
 
-fn is_taproot(script_config: &pb::BtcScriptConfigWithKeypath) -> bool {
-    matches!(script_config,
-        pb::BtcScriptConfigWithKeypath {
-            script_config:
-                Some(pb::BtcScriptConfig {
-                    config: Some(pb::btc_script_config::Config::SimpleType(simple_type)),
-                }),
-            ..
-        } if *simple_type == pb::btc_script_config::SimpleType::P2tr as i32)
+fn is_taproot_simple(script_config: &pb::BtcScriptConfigWithKeypath) -> bool {
+    matches!(
+        script_config.script_config.as_ref(),
+        Some(pb::BtcScriptConfig {
+            config: Some(pb::btc_script_config::Config::SimpleType(simple_type)),
+        }) if *simple_type == pb::btc_script_config::SimpleType::P2tr as i32
+    )
+}
+
+fn is_taproot_policy(script_config: &pb::BtcScriptConfigWithKeypath) -> bool {
+    matches!(
+        script_config.script_config.as_ref(),
+        Some(pb::BtcScriptConfig {
+            config: Some(pb::btc_script_config::Config::Policy(policy)),
+        })  if policy.policy.as_str().starts_with("tr("),
+    )
+}
+
+fn is_schnorr(script_config: &pb::BtcScriptConfigWithKeypath) -> bool {
+    is_taproot_simple(script_config) | is_taproot_policy(script_config)
 }
 
 impl<R: Runtime> PairedBitBox<R> {
@@ -673,7 +705,7 @@ impl<R: Runtime> PairedBitBox<R> {
         format_unit: pb::btc_sign_init_request::FormatUnit,
     ) -> Result<Vec<Vec<u8>>, Error> {
         self.validate_version(">=9.4.0")?; // anti-klepto since 9.4.0
-        if transaction.script_configs.iter().any(is_taproot) {
+        if transaction.script_configs.iter().any(is_taproot_simple) {
             self.validate_version(">=9.10.0")?; // taproot since 9.10.0
         }
 
@@ -700,7 +732,7 @@ impl<R: Runtime> PairedBitBox<R> {
                     let input_index: usize = next_response.index as _;
                     let tx_input: &TxInput = &transaction.inputs[input_index];
 
-                    let input_is_schnorr = is_taproot(
+                    let input_is_schnorr = is_schnorr(
                         &transaction.script_configs[tx_input.script_config_index as usize],
                     );
                     let perform_antiklepto = is_inputs_pass2 && !input_is_schnorr;
@@ -896,7 +928,12 @@ impl<R: Runtime> PairedBitBox<R> {
                     psbt_input.tap_key_sig = Some(
                         bitcoin::taproot::Signature::from_slice(signature)
                             .map_err(|_| Error::InvalidSignature)?,
-                    )
+                    );
+                }
+                OurKey::TaprootScript(xonly, leaf_hash, _) => {
+                    let sig = bitcoin::taproot::Signature::from_slice(signature)
+                        .map_err(|_| Error::InvalidSignature)?;
+                    psbt_input.tap_script_sigs.insert((xonly, leaf_hash), sig);
                 }
             }
         }