diff --git a/src/sign.rs b/src/sign.rs index 78df32e..05f4ea4 100644 --- a/src/sign.rs +++ b/src/sign.rs @@ -2,8 +2,8 @@ //! according to the BIP-322 standard. use crate::{ - BIP322, Error, MessageProof, MessageVerificationResult, SignatureFormat, derive_tx_params, - to_sign, to_spend, validate_witness, verify_psbt_proof, verify_signed_proof, + BIP322, Error, MessageProof, MessageVerificationResult, SignatureFormat, configure_p2sh_input, + derive_tx_params, to_sign, to_spend, validate_witness, verify_psbt_proof, verify_signed_proof, }; use alloc::{string::ToString, vec::Vec}; @@ -190,19 +190,23 @@ fn configure_psbt_inputs( Some(PsbtSighashType::from(EcdsaSighashType::All)) }; - if input_spk.is_p2tr() || input_spk.is_p2wpkh() || input_spk.is_p2wsh() { + let descriptor = wallet.public_descriptor(keychain); + let derived_descriptor = descriptor + .at_derivation_index(derivation_index) + .map_err(|e| Error::InvalidFormat(e.to_string()))?; + + if input_spk.is_p2tr() || input_spk.is_p2wpkh() { + psbt_input.witness_utxo = Some(txout); + } else if input_spk.is_p2wsh() { psbt_input.witness_utxo = Some(txout); - if input_spk.is_p2wsh() { - let desc = wallet.public_descriptor(keychain); - let derived = desc - .at_derivation_index(derivation_index) - .map_err(|e| Error::InvalidFormat(e.to_string()))?; - let script = derived - .explicit_script() - .map_err(|e| Error::InvalidFormat(e.to_string()))?; - psbt_input.witness_script = Some(script); - } + let script = derived_descriptor + .explicit_script() + .map_err(|e| Error::InvalidFormat(e.to_string()))?; + psbt_input.witness_script = Some(script); + } else if script_pubkey.is_p2sh() { + psbt_input.witness_utxo = Some(txout); + configure_p2sh_input(psbt_input, &derived_descriptor)?; } else if input_spk.is_p2pkh() { // P2PKH requires full transaction as non-witness UTXO if i == 0 { @@ -250,6 +254,13 @@ fn encode_signature( "Legacy format is verify-only. Use Simple or Full for P2PKH addresses".to_string(), )), SignatureFormat::Simple => { + if script_pubkey.is_p2sh() { + return Err(Error::InvalidFormat( + "Simple format is not supported for P2SH addresses. Use Full format." + .to_string(), + )); + } + let witness = psbt.inputs[0] .final_script_witness .as_ref() @@ -294,7 +305,7 @@ mod tests { } #[test] - fn test_p2pkh_format() { + fn test_simple_format_p2pkh() { const EXTERNAL_DESC: &str = "pkh(tprv8ZgxMBicQKsPfGXKjYNsw4gayjfBsq6FHxvNZ8LSBdz4DSTeBPd7cjvVQXTdMH9NJBVwNrNKLDr58dcrf4YmWLYBs4KogJhSgUELXuo1JwH/44'/1'/0'/0/*)"; const INTERNAL_DESC: &str = "pkh(tprv8ZgxMBicQKsPfGXKjYNsw4gayjfBsq6FHxvNZ8LSBdz4DSTeBPd7cjvVQXTdMH9NJBVwNrNKLDr58dcrf4YmWLYBs4KogJhSgUELXuo1JwH/44'/1'/0'/1/*)"; @@ -554,6 +565,43 @@ mod tests { assert_ne!(verify.proven_amount.unwrap(), Amount::from_sat(0)) } + #[test] + fn test_simple_format_p2sh_p2wpkh() { + const EXTERNAL_DESC: &str = "sh(wpkh(tprv8ZgxMBicQKsPdy6LMhUtFHAgpocR8GC6QmwMSFpZs7h6Eziw3SpThFfczTDh5rW2krkqffa11UpX3XkeTTB2FvzZKWXqPY54Y6Rq4AQ5R8L/49'/1'/0'/0/*))"; + const INTERNAL_DESC: &str = "sh(wpkh(tprv8ZgxMBicQKsPdy6LMhUtFHAgpocR8GC6QmwMSFpZs7h6Eziw3SpThFfczTDh5rW2krkqffa11UpX3XkeTTB2FvzZKWXqPY54Y6Rq4AQ5R8L/49'/1'/0'/1/*))"; + + let (mut wallet, _) = get_funded_wallet(EXTERNAL_DESC, INTERNAL_DESC); + let address = wallet.peek_address(KeychainKind::External, 0).address; + + let sign = wallet + .sign_message("HELLO WORLD", SignatureFormat::Simple, &address, None) + .unwrap(); + + let verify = wallet + .verify_message(&sign, "HELLO WORLD", &address) + .unwrap(); + + assert!(verify.valid); + } + + #[test] + fn test_full_format_p2sh_p2wsh() { + let (mut wallet, _) = get_funded_wallet_single( + "sh(wsh(pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW)))", + ); + let address = wallet.peek_address(KeychainKind::External, 0).address; + + let sign = wallet + .sign_message("HELLO WORLD", SignatureFormat::Full, &address, None) + .unwrap(); + + let verify = wallet + .verify_message(&sign, "HELLO WORLD", &address) + .unwrap(); + + assert!(verify.valid); + } + #[test] fn test_wrong_message_fails_verification() { const EXTERNAL_DESC: &str = "wpkh(tprv8ZgxMBicQKsPdy6LMhUtFHAgpocR8GC6QmwMSFpZs7h6Eziw3SpThFfczTDh5rW2krkqffa11UpX3XkeTTB2FvzZKWXqPY54Y6Rq4AQ5R8L/84'/1'/0'/0/*)"; diff --git a/src/utils.rs b/src/utils.rs index 3616d29..2969e42 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -6,7 +6,7 @@ use bdk_wallet::{ Wallet, keys::ScriptContext, miniscript::{ - Descriptor, Miniscript, Terminal, + DefiniteDescriptorKey, Descriptor, Miniscript, Terminal, ToPublicKey as _, descriptor::{ShInner, WshInner}, }, }; @@ -117,9 +117,15 @@ pub fn validate_witness(witness: &Witness, script_pubkey: &ScriptBuf) -> Result< "P2TR requires at least 1 witness element".to_string(), )); } + } else if script_pubkey.is_p2sh() { + if witness.is_empty() { + return Err(Error::InvalidFormat( + "P2SH-wrapped SegWit requires witness data".to_string(), + )); + } } else { return Err(Error::InvalidFormat( - "Simple format only supports segwit script types (P2WPKH, P2WSH, P2TR)".to_string(), + "Unsupported script type for Simple format".to_string(), )); } @@ -325,6 +331,66 @@ fn find_timelocks( (max_csv, max_cltv) } +/// Extracts the redeem script from a P2SH script_sig. +pub fn extract_redeem_script(script_sig: &ScriptBuf) -> Result { + let mut instructions = script_sig.instructions(); + + let redeem_script = match instructions.next() { + Some(Ok(Instruction::PushBytes(bytes))) => ScriptBuf::from_bytes(bytes.as_bytes().to_vec()), + _ => { + return Err(Error::InvalidFormat( + "P2SH scriptSig must start with a push".to_string(), + )); + } + }; + + if instructions.next().is_some() { + return Err(Error::InvalidFormat( + "P2SH-wrapped SegWit scriptSig must be a single push".to_string(), + )); + } + + Ok(redeem_script) +} + +/// Configures PSBT input fields for P2SH-wrapped SegWit descriptors. +pub fn configure_p2sh_input( + psbt_input: &mut bitcoin::psbt::Input, + derived: &Descriptor, +) -> Result<(), Error> { + let Descriptor::Sh(sh) = derived else { + return Err(Error::InvalidFormat( + "Expected Sh descriptor for P2SH scriptPubKey".to_string(), + )); + }; + + match sh.as_inner() { + ShInner::Wpkh(wpkh) => { + let pk = wpkh.as_inner().to_public_key(); + let wpkh_hash = pk + .wpubkey_hash() + .map_err(|e| Error::InvalidPublicKey(e.to_string()))?; + psbt_input.redeem_script = Some(ScriptBuf::new_p2wpkh(&wpkh_hash)); + } + ShInner::Wsh(wsh) => { + let witness_script = match wsh.as_inner() { + WshInner::Ms(ms) => ms.encode(), + WshInner::SortedMulti(sm) => sm.encode(), + }; + let wsh_hash = witness_script.wscript_hash(); + psbt_input.redeem_script = Some(ScriptBuf::new_p2wsh(&wsh_hash)); + psbt_input.witness_script = Some(witness_script); + } + _ => { + return Err(Error::UnsupportedScriptType( + "Only P2SH-P2WPKH and P2SH-P2WSH are supported".to_string(), + )); + } + } + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/verify.rs b/src/verify.rs index c3c0fad..45f840a 100644 --- a/src/verify.rs +++ b/src/verify.rs @@ -5,7 +5,7 @@ use crate::{ Error, MessageVerificationResult, SignatureFormat, detect_signature_format, extract_pubkeys, - to_sign, to_spend, validate_to_sign, validate_witness, + extract_redeem_script, to_sign, to_spend, validate_to_sign, validate_witness, }; use alloc::{string::ToString, vec::Vec}; use bdk_wallet::Wallet; @@ -206,13 +206,15 @@ fn verify_message( if wp.version() != WitnessVersion::V0 { return Err(Error::UnsupportedSegwitVersion("v0".to_string())); } - verify_p2wsh(to_sign, &prevout, address, 0, secp)? + verify_p2wsh(to_sign, &prevout, &script_pubkey, 0, secp)? } else if script_pubkey.is_p2tr() { let wp = address.witness_program().ok_or(Error::NotSegwitAddress)?; if wp.version() != WitnessVersion::V1 { return Err(Error::UnsupportedSegwitVersion("v1".to_string())); } verify_p2tr(to_sign, &prevout, 0, wallet, &to_spend, secp)? + } else if script_pubkey.is_p2sh() { + verify_p2sh(to_sign, &prevout, &script_pubkey, 0, secp)? } else { return Ok(false); }; @@ -267,7 +269,11 @@ fn verify_proof_of_funds( return Ok(false); } } else if script_pubkey.is_p2wsh() { - if !verify_p2wsh(to_sign, &utxo.txout, address, i, secp)? { + if !verify_p2wsh(to_sign, &utxo.txout, &script_pubkey, i, secp)? { + return Ok(false); + } + } else if script_pubkey.is_p2sh() { + if !verify_p2sh(to_sign, &utxo.txout, &script_pubkey, i, secp)? { return Ok(false); } } else { @@ -498,11 +504,11 @@ fn verify_p2wpkh( fn verify_p2wsh( to_sign: &Transaction, prevout: &TxOut, - address: &Address, + script_pubkey: &ScriptBuf, input_index: usize, secp: &Secp256k1, ) -> Result { - let script_pubkey = address.script_pubkey(); + // let script_pubkey = address.script_pubkey(); let witness = &to_sign.input[input_index].witness; if witness.len() < 2 { @@ -520,7 +526,7 @@ fn verify_p2wsh( let script_hash = witness_script.wscript_hash(); let expected_script_pubkey = ScriptBuf::new_p2wsh(&script_hash); - if script_pubkey != expected_script_pubkey { + if *script_pubkey != expected_script_pubkey { return Err(Error::InvalidSignature( "Witness script hash doesn't match address".to_string(), )); @@ -683,3 +689,51 @@ fn verify_p2tr( Ok(secp.verify_schnorr(&signature, msg, &pub_key).is_ok()) } + +/// Verifies P2SH-wrapped SegWit signatures (P2SH-P2WPKH and P2SH-P2WSH). +/// +/// Extracts the redeem script from the input's script_sig, determines +/// the inner SegWit type, and delegates to the appropriate verifier. +fn verify_p2sh( + to_sign: &Transaction, + prevout: &TxOut, + script_pubkey: &ScriptBuf, + input_index: usize, + secp: &Secp256k1, +) -> Result { + let input = &to_sign.input[input_index]; + + // P2SH-wrapped SegWit: the script_sig contains a single push of the redeem script. + // The redeem script IS the inner SegWit scriptPubKey (e.g., OP_0 <20-byte-hash>). + let redeem_script = extract_redeem_script(&input.script_sig)?; + + // Verify the redeem script hashes to the P2SH address + let script_hash = redeem_script.script_hash(); + let expected_script_pubkey = ScriptBuf::new_p2sh(&script_hash); + + if *script_pubkey != expected_script_pubkey { + return Err(Error::InvalidSignature( + "Redeem script hash doesn't match P2SH address".to_string(), + )); + } + + // Create a synthetic prevout with the inner SegWit scriptPubKey + // so the existing verifiers compute the correct sighash. + let inner_prevout = TxOut { + value: prevout.value, + script_pubkey: redeem_script.clone(), + }; + + if redeem_script.is_p2wpkh() { + verify_p2wpkh(to_sign, &inner_prevout, input_index, secp) + } else if redeem_script.is_p2wsh() { + // For P2SH-P2WSH, we need a synthetic address wrapping the inner P2WSH + // so verify_p2wsh can check the witness script hash. + // We construct a temporary address from the inner script. + verify_p2wsh(to_sign, &inner_prevout, &redeem_script, input_index, secp) + } else { + Err(Error::UnsupportedScriptType( + "Only P2SH-P2WPKH and P2SH-P2WSH are supported.".to_string(), + )) + } +}