diff --git a/wallet/core/src/account/pskb.rs b/wallet/core/src/account/pskb.rs index 6d8dabb49..d8a7c018c 100644 --- a/wallet/core/src/account/pskb.rs +++ b/wallet/core/src/account/pskb.rs @@ -153,7 +153,6 @@ pub async fn bundle_from_pskt_generator(generator: PSKTGenerator) -> Result, @@ -163,64 +162,78 @@ pub async fn pskb_signer_for_address( key_fingerprint: KeyFingerprint, ) -> Result { let mut signed_bundle = Bundle::new(); - let reused_values = SigHashReusedValuesUnsync::new(); - // If set, sign-for address is used for signing. - // Else, all addresses from inputs are. - let addresses: Vec
= match sign_for_address { - Some(signer) => vec![signer.clone()], - None => bundle + // If sign_for_address is provided, we'll use it for all signatures + // Otherwise, collect addresses per PSKT + let addresses_per_pskt: Vec> = if sign_for_address.is_some() { + // Create a vec of single-address vecs + bundle.iter().map(|_| vec![sign_for_address.unwrap().clone()]).collect() + } else { + // Collect addresses for each PSKT separately + bundle .iter() - .flat_map(|inner| { - inner.inputs + .map(|inner| { + inner + .inputs .iter() - .filter_map(|input| input.utxo_entry.as_ref()) // Filter out None and get a reference to UtxoEntry if it exists + .filter_map(|input| input.utxo_entry.as_ref()) .filter_map(|utxo_entry| { extract_script_pub_key_address(&utxo_entry.script_public_key.clone(), network_id.into()).ok() }) - .collect::>() + .collect() }) - .collect(), + .collect() }; - // Prepare the signer. - signer.ingest(addresses.as_ref())?; + // Prepare the signer with all unique addresses + let all_addresses: Vec
= addresses_per_pskt.iter().flat_map(|addresses| addresses.iter().cloned()).collect(); + signer.ingest(all_addresses.as_slice())?; - for pskt_inner in bundle.iter().cloned() { + // Process each PSKT in the bundle + for (pskt_idx, pskt_inner) in bundle.iter().cloned().enumerate() { let pskt: PSKT = PSKT::from(pskt_inner); + let current_addresses = &addresses_per_pskt[pskt_idx]; + + // Create new reused values for each PSKT + let reused_values = SigHashReusedValuesUnsync::new(); - let sign = |signer_pskt: PSKT| { + let sign = |signer_pskt: PSKT| -> Result, Error> { signer_pskt .pass_signature_sync(|tx, sighash| -> Result, String> { tx.tx .inputs .iter() .enumerate() - .map(|(idx, _input)| { - let hash = calc_schnorr_signature_hash(&tx.as_verifiable(), idx, sighash[idx], &reused_values); - let msg = secp256k1::Message::from_digest_slice(hash.as_bytes().as_slice()).unwrap(); - - // When address represents a locked UTXO, no private key is available. - // Instead, use the account receive address' private key. - let address: &Address = match sign_for_address { - Some(address) => address, - None => addresses.get(idx).expect("Input indexed address"), + .map(|(input_idx, _input)| { + let hash = calc_schnorr_signature_hash(&tx.as_verifiable(), input_idx, sighash[input_idx], &reused_values); + let msg = secp256k1::Message::from_digest_slice(hash.as_bytes().as_slice()).map_err(|e| e.to_string())?; + + // Get the appropriate address for this input + let address = if let Some(sign_addr) = sign_for_address { + sign_addr + } else { + current_addresses.get(input_idx).ok_or_else(|| format!("No address found for input {}", input_idx))? }; - let public_key = signer.public_key(address).expect("Public key for input indexed address"); + let public_key = signer.public_key(address).map_err(|e| format!("Failed to get public key: {}", e))?; + + let signature = signer.sign_schnorr(address, msg).map_err(|e| format!("Failed to sign: {}", e))?; Ok(SignInputOk { - signature: Signature::Schnorr(signer.sign_schnorr(address, msg).unwrap()), + signature: Signature::Schnorr(signature), pub_key: public_key, key_source: Some(KeySource { key_fingerprint, derivation_path: derivation_path.clone() }), }) }) .collect() }) - .unwrap() + .map_err(Error::from) }; - signed_bundle.add_pskt(sign(pskt.clone())); + + let signed_pskt = sign(pskt)?; + signed_bundle.add_pskt(signed_pskt); } + Ok(signed_bundle) }