Fix weight estimation for P2TR inputs

Estimation can only occur if the P2TR input is a key spend, since
we can't reliably predict the size of potential unlocking scripts

Author: shinghim

payjoin/src/core/psbt/mod.rs

@@ -9,7 +9,7 @@ use std::fmt;
 use bitcoin::address::FromScriptError;
 use bitcoin::psbt::Psbt;
 use bitcoin::transaction::InputWeightPrediction;
-use bitcoin::{bip32, psbt, Address, AddressType, Network, TxIn, TxOut, Weight};
+use bitcoin::{bip32, psbt, Address, AddressType, Network, TapSighashType, TxIn, TxOut, Weight};
 
 #[derive(Debug, PartialEq)]
 pub(crate) enum InconsistentPsbt {
@@ -98,6 +98,9 @@ impl PsbtExt for Psbt {
 // https://github.com/bitcoin/bips/blob/master/bip-0141.mediawiki#p2wpkh-nested-in-bip16-p2sh
 const NESTED_P2WPKH_MAX: InputWeightPrediction = InputWeightPrediction::from_slice(23, &[72, 33]);
 
+const SCHNORR_SIG_DEFAULT_SIZE: usize = 64;
+const SCHNORR_SIG_WITH_SIGHASH_SIZE: usize = 65;
+
 #[derive(Clone, Debug)]
 pub(crate) struct InternalInputPair<'a> {
     pub txin: &'a TxIn,
@@ -207,7 +210,38 @@ impl InternalInputPair<'_> {
             }
             P2wpkh => Ok(InputWeightPrediction::P2WPKH_MAX),
             P2wsh => Err(InputWeightError::NotSupported),
-            P2tr => Ok(InputWeightPrediction::P2TR_KEY_DEFAULT_SIGHASH),
+            P2tr => {
+                // For unfinalized psbts, handle via tap_key_sig
+                if let Some(sig) = self.psbtin.tap_key_sig {
+                    if !self.psbtin.tap_scripts.is_empty()
+                        || !self.psbtin.tap_script_sigs.is_empty()
+                        || self.psbtin.tap_merkle_root.is_some()
+                    {
+                        return Err(InputWeightError::NotSupported);
+                    }
+                    match sig.sighash_type {
+                        TapSighashType::Default =>
+                            Ok(InputWeightPrediction::P2TR_KEY_DEFAULT_SIGHASH),
+                        _ => Ok(InputWeightPrediction::P2TR_KEY_NON_DEFAULT_SIGHASH),
+                    }
+                }
+                // For finalized psbts, check final_script_witness
+                else if let Some(witness) = &self.psbtin.final_script_witness {
+                    if witness.len() == 1 {
+                        match witness[0].len() {
+                            SCHNORR_SIG_DEFAULT_SIZE =>
+                                Ok(InputWeightPrediction::P2TR_KEY_DEFAULT_SIGHASH),
+                            SCHNORR_SIG_WITH_SIGHASH_SIZE =>
+                                Ok(InputWeightPrediction::P2TR_KEY_NON_DEFAULT_SIGHASH),
+                            _ => Err(InputWeightError::NotSupported),
+                        }
+                    } else {
+                        Err(InputWeightError::NotSupported)
+                    }
+                } else {
+                    Err(InputWeightError::NotSupported)
+                }
+            }
             _ => Err(AddressTypeError::UnknownAddressType.into()),
         }?;
 
@@ -376,3 +410,132 @@ impl std::error::Error for InputWeightError {
 impl From<AddressTypeError> for InputWeightError {
     fn from(value: AddressTypeError) -> Self { Self::AddressType(value) }
 }
+
+#[cfg(test)]
+mod tests {
+    use bitcoin::key::Secp256k1;
+    use bitcoin::taproot::{ControlBlock, LeafVersion};
+    use bitcoin::{psbt, secp256k1, taproot, PublicKey, ScriptBuf, TapNodeHash, XOnlyPublicKey};
+
+    use super::*;
+    use crate::core::psbt::InternalInputPair;
+    use crate::receive::InputPair;
+
+    /// Lengths of txid, index and sequence: (32, 4, 4)
+    const TXID_INDEX_SEQUENCE_WEIGHT: Weight = Weight::from_non_witness_data_size(32 + 4 + 4);
+
+    #[test]
+    fn expected_weight_for_p2tr() {
+        let pubkey_string = "0347ff3dacd07a1f43805ec6808e801505a6e18245178609972a68afbc2777ff2b";
+        let pubkey = pubkey_string.parse::<PublicKey>().expect("valid pubkey");
+        let xonly_pubkey = XOnlyPublicKey::from(pubkey.inner);
+        let p2tr_utxo = TxOut {
+            value: Default::default(),
+            script_pubkey: ScriptBuf::new_p2tr(&Secp256k1::new(), xonly_pubkey, None),
+        };
+        let default_sighash_pair = InputPair {
+            txin: Default::default(),
+            psbtin: psbt::Input {
+                tap_key_sig: Some(
+                    taproot::Signature::from_slice(
+                        &[0; secp256k1::constants::SCHNORR_SIGNATURE_SIZE],
+                    )
+                    .unwrap(),
+                ),
+                witness_utxo: Some(p2tr_utxo.clone()),
+                ..Default::default()
+            },
+        };
+        assert_eq!(
+            InternalInputPair::from(&default_sighash_pair).expected_input_weight().unwrap(),
+            InputWeightPrediction::P2TR_KEY_DEFAULT_SIGHASH.weight() + TXID_INDEX_SEQUENCE_WEIGHT
+        );
+
+        // Add a sighash byte
+        let mut sig_bytes = [0; secp256k1::constants::SCHNORR_SIGNATURE_SIZE + 1];
+        sig_bytes[sig_bytes.len() - 1] = 1;
+        let non_default_sighash_pair = InputPair {
+            txin: Default::default(),
+            psbtin: psbt::Input {
+                tap_key_sig: Some(taproot::Signature::from_slice(&sig_bytes).unwrap()),
+                witness_utxo: Some(p2tr_utxo),
+                ..Default::default()
+            },
+        };
+        assert_eq!(
+            InternalInputPair::from(&non_default_sighash_pair).expected_input_weight().unwrap(),
+            InputWeightPrediction::P2TR_KEY_NON_DEFAULT_SIGHASH.weight()
+                + TXID_INDEX_SEQUENCE_WEIGHT
+        );
+    }
+
+    #[test]
+    fn not_supported_p2tr_expected_weights() {
+        let pubkey_string = "0347ff3dacd07a1f43805ec6808e801505a6e18245178609972a68afbc2777ff2b";
+        let pubkey = pubkey_string.parse::<PublicKey>().expect("valid pubkey");
+        let xonly_pubkey = XOnlyPublicKey::from(pubkey.inner);
+        let p2tr_script = ScriptBuf::new_p2tr(&Secp256k1::new(), xonly_pubkey.clone(), None);
+        let p2tr_utxo = TxOut { value: Default::default(), script_pubkey: p2tr_script.clone() };
+
+        let mut tap_scripts = BTreeMap::new();
+        let leaf_version: u8 = 0xC0;
+        let mut control_block_vec = Vec::with_capacity(33);
+        control_block_vec.push(leaf_version);
+        control_block_vec.extend_from_slice(&xonly_pubkey.serialize());
+        let control_block = ControlBlock::decode(control_block_vec.as_slice()).unwrap();
+        tap_scripts
+            .insert(control_block.clone(), (p2tr_script.clone(), control_block.leaf_version));
+
+        let pair_with_tapscripts = InputPair {
+            txin: Default::default(),
+            psbtin: psbt::Input {
+                tap_scripts,
+                witness_utxo: Some(p2tr_utxo.clone()),
+                ..Default::default()
+            },
+        };
+        assert_eq!(
+            InternalInputPair::from(&pair_with_tapscripts).expected_input_weight().err().unwrap(),
+            InputWeightError::NotSupported
+        );
+
+        let mut tap_script_sigs = BTreeMap::new();
+        tap_script_sigs.insert(
+            (xonly_pubkey.clone(), p2tr_script.tapscript_leaf_hash()),
+            taproot::Signature::from_slice(&[0; secp256k1::constants::SCHNORR_SIGNATURE_SIZE])
+                .unwrap(),
+        );
+        let pair_with_tap_script_sigs = InputPair {
+            txin: Default::default(),
+            psbtin: psbt::Input {
+                tap_script_sigs,
+                witness_utxo: Some(p2tr_utxo.clone()),
+                ..Default::default()
+            },
+        };
+        assert_eq!(
+            InternalInputPair::from(&pair_with_tap_script_sigs)
+                .expected_input_weight()
+                .err()
+                .unwrap(),
+            InputWeightError::NotSupported
+        );
+
+        let tap_merkle_root = TapNodeHash::from_script(&p2tr_script, LeafVersion::TapScript);
+        let pair_with_tap_merkle_root = InputPair {
+            txin: Default::default(),
+            psbtin: psbt::Input {
+                tap_merkle_root: Some(tap_merkle_root),
+                witness_utxo: Some(p2tr_utxo.clone()),
+                ..Default::default()
+            },
+        };
+        assert_eq!(
+            InternalInputPair::from(&pair_with_tap_merkle_root)
+                .expected_input_weight()
+                .err()
+                .unwrap(),
+            InputWeightError::NotSupported
+        );
+    }
+}