|
| 1 | +use std::collections::BTreeMap; |
| 2 | + |
| 3 | +use bitcoin::{ |
| 4 | + NetworkKind, PrivateKey, Psbt, PublicKey, Witness, |
| 5 | + bip32::{DerivationPath, Fingerprint, KeySource}, |
| 6 | + secp256k1::Secp256k1, |
| 7 | +}; |
| 8 | +use primitives::SignerError; |
| 9 | + |
| 10 | +pub fn sign_psbt(psbt_hex: &str, private_key: &[u8]) -> Result<String, SignerError> { |
| 11 | + let psbt_bytes = hex::decode(psbt_hex).map_err(|e| SignerError::invalid_input(format!("hex decode: {e}")))?; |
| 12 | + let mut psbt = Psbt::deserialize(&psbt_bytes).map_err(|e| SignerError::invalid_input(format!("psbt parse: {e}")))?; |
| 13 | + |
| 14 | + let secp = Secp256k1::new(); |
| 15 | + let key = PrivateKey::from_slice(private_key, NetworkKind::Main).map_err(|e| SignerError::invalid_input(format!("private key: {e}")))?; |
| 16 | + let pub_key = PublicKey::from_private_key(&secp, &key); |
| 17 | + |
| 18 | + prepare_inputs(&mut psbt, &pub_key); |
| 19 | + sign(&mut psbt, &pub_key, &key, &secp)?; |
| 20 | + finalize_inputs(&mut psbt, &pub_key)?; |
| 21 | + |
| 22 | + let tx = psbt.extract_tx_unchecked_fee_rate(); |
| 23 | + Ok(hex::encode(bitcoin::consensus::serialize(&tx))) |
| 24 | +} |
| 25 | + |
| 26 | +fn prepare_inputs(psbt: &mut Psbt, pub_key: &PublicKey) { |
| 27 | + let (x_only_key, _) = pub_key.inner.x_only_public_key(); |
| 28 | + let default_origin: KeySource = (Fingerprint::default(), DerivationPath::master()); |
| 29 | + |
| 30 | + for input in &mut psbt.inputs { |
| 31 | + let is_taproot = input.witness_utxo.as_ref().is_some_and(|utxo| utxo.script_pubkey.is_p2tr()); |
| 32 | + |
| 33 | + if is_taproot { |
| 34 | + input.tap_internal_key.get_or_insert(x_only_key); |
| 35 | + if input.tap_key_origins.is_empty() { |
| 36 | + input.tap_key_origins.insert(x_only_key, (vec![], default_origin.clone())); |
| 37 | + } |
| 38 | + } else if input.bip32_derivation.is_empty() { |
| 39 | + input.bip32_derivation.insert(pub_key.inner, default_origin.clone()); |
| 40 | + } |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +fn sign(psbt: &mut Psbt, pub_key: &PublicKey, key: &PrivateKey, secp: &Secp256k1<bitcoin::secp256k1::All>) -> Result<(), SignerError> { |
| 45 | + let keys = BTreeMap::from([(*pub_key, *key)]); |
| 46 | + psbt.sign(&keys, secp).map(|_| ()).map_err(|(_ok, errors)| { |
| 47 | + let messages: Vec<String> = errors.into_iter().map(|(idx, e)| format!("input {idx}: {e}")).collect(); |
| 48 | + SignerError::signing_error(messages.join(", ")) |
| 49 | + }) |
| 50 | +} |
| 51 | + |
| 52 | +fn finalize_inputs(psbt: &mut Psbt, pub_key: &PublicKey) -> Result<(), SignerError> { |
| 53 | + for (idx, input) in psbt.inputs.iter_mut().enumerate() { |
| 54 | + let script = &input |
| 55 | + .witness_utxo |
| 56 | + .as_ref() |
| 57 | + .ok_or_else(|| SignerError::signing_error(format!("missing witness_utxo for input {idx}")))? |
| 58 | + .script_pubkey; |
| 59 | + |
| 60 | + let witness = build_witness(input, pub_key, script, idx)?; |
| 61 | + input.final_script_witness = Some(witness); |
| 62 | + input.partial_sigs.clear(); |
| 63 | + input.sighash_type = None; |
| 64 | + input.redeem_script = None; |
| 65 | + input.witness_script = None; |
| 66 | + input.bip32_derivation.clear(); |
| 67 | + } |
| 68 | + Ok(()) |
| 69 | +} |
| 70 | + |
| 71 | +fn build_witness(input: &bitcoin::psbt::Input, pub_key: &PublicKey, script: &bitcoin::ScriptBuf, idx: usize) -> Result<Witness, SignerError> { |
| 72 | + if script.is_p2wpkh() { |
| 73 | + let sig = input |
| 74 | + .partial_sigs |
| 75 | + .get(pub_key) |
| 76 | + .ok_or_else(|| SignerError::signing_error(format!("missing signature for input {idx}")))?; |
| 77 | + let mut w = Witness::new(); |
| 78 | + w.push(sig.to_vec()); |
| 79 | + w.push(pub_key.to_bytes()); |
| 80 | + Ok(w) |
| 81 | + } else if script.is_p2tr() { |
| 82 | + let sig = input |
| 83 | + .tap_key_sig |
| 84 | + .ok_or_else(|| SignerError::signing_error(format!("missing taproot signature for input {idx}")))?; |
| 85 | + let mut w = Witness::new(); |
| 86 | + w.push(sig.to_vec()); |
| 87 | + Ok(w) |
| 88 | + } else { |
| 89 | + Err(SignerError::signing_error(format!("unsupported script type for input {idx}"))) |
| 90 | + } |
| 91 | +} |
| 92 | + |
| 93 | +#[cfg(test)] |
| 94 | +mod tests { |
| 95 | + use super::*; |
| 96 | + use crate::testkit::TEST_PRIVATE_KEY; |
| 97 | + use bitcoin::{Amount, OutPoint, ScriptBuf, Sequence, Transaction, TxIn, TxOut, Txid, hashes::Hash, secp256k1::Secp256k1}; |
| 98 | + |
| 99 | + fn build_test_psbt(script_pubkey: ScriptBuf) -> Psbt { |
| 100 | + let utxo = TxOut { |
| 101 | + value: Amount::from_sat(100_000), |
| 102 | + script_pubkey: script_pubkey.clone(), |
| 103 | + }; |
| 104 | + let tx = Transaction { |
| 105 | + version: bitcoin::transaction::Version(2), |
| 106 | + lock_time: bitcoin::absolute::LockTime::ZERO, |
| 107 | + input: vec![TxIn { |
| 108 | + previous_output: OutPoint::new(Txid::all_zeros(), 0), |
| 109 | + script_sig: ScriptBuf::new(), |
| 110 | + sequence: Sequence::ENABLE_RBF_NO_LOCKTIME, |
| 111 | + witness: Witness::new(), |
| 112 | + }], |
| 113 | + output: vec![TxOut { |
| 114 | + value: Amount::from_sat(90_000), |
| 115 | + script_pubkey, |
| 116 | + }], |
| 117 | + }; |
| 118 | + let mut psbt = Psbt::from_unsigned_tx(tx).unwrap(); |
| 119 | + psbt.inputs[0].witness_utxo = Some(utxo); |
| 120 | + psbt |
| 121 | + } |
| 122 | + |
| 123 | + fn sign_and_verify(psbt: Psbt) { |
| 124 | + let psbt_hex = hex::encode(psbt.serialize()); |
| 125 | + let result = sign_psbt(&psbt_hex, &TEST_PRIVATE_KEY).unwrap(); |
| 126 | + assert!(!result.is_empty()); |
| 127 | + |
| 128 | + let tx: Transaction = bitcoin::consensus::deserialize(&hex::decode(&result).unwrap()).unwrap(); |
| 129 | + assert_eq!(tx.input.len(), 1); |
| 130 | + assert!(!tx.input[0].witness.is_empty()); |
| 131 | + } |
| 132 | + |
| 133 | + #[test] |
| 134 | + fn test_sign_p2wpkh_psbt() { |
| 135 | + let secp = Secp256k1::new(); |
| 136 | + let key = PrivateKey::from_slice(&TEST_PRIVATE_KEY, NetworkKind::Main).unwrap(); |
| 137 | + let pub_key = PublicKey::from_private_key(&secp, &key); |
| 138 | + let script = ScriptBuf::new_p2wpkh(&pub_key.wpubkey_hash().unwrap()); |
| 139 | + |
| 140 | + sign_and_verify(build_test_psbt(script)); |
| 141 | + } |
| 142 | + |
| 143 | + #[test] |
| 144 | + fn test_sign_p2tr_psbt() { |
| 145 | + let secp = Secp256k1::new(); |
| 146 | + let key = PrivateKey::from_slice(&TEST_PRIVATE_KEY, NetworkKind::Main).unwrap(); |
| 147 | + let (x_only, _) = key.public_key(&secp).inner.x_only_public_key(); |
| 148 | + let script = ScriptBuf::new_p2tr(&secp, x_only, None); |
| 149 | + |
| 150 | + sign_and_verify(build_test_psbt(script)); |
| 151 | + } |
| 152 | + |
| 153 | + #[test] |
| 154 | + fn test_sign_psbt_invalid_hex() { |
| 155 | + assert!(sign_psbt("not_hex!", &TEST_PRIVATE_KEY).is_err()); |
| 156 | + } |
| 157 | + |
| 158 | + #[test] |
| 159 | + fn test_sign_psbt_invalid_psbt() { |
| 160 | + assert!(sign_psbt("deadbeef", &TEST_PRIVATE_KEY).is_err()); |
| 161 | + } |
| 162 | +} |
0 commit comments