Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 114 additions & 19 deletions src/wallet/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,55 @@ impl fmt::Display for AddressInfo {
/// A `CanonicalTx` managed by a `Wallet`.
pub type WalletTx<'a> = CanonicalTx<'a, Arc<Transaction>, ConfirmationBlockTime>;

/// The finalization status for a single PSBT input.
#[derive(Debug, PartialEq)]
pub enum FinalizeInputOutcome {
/// The input was already finalized before this call.
AlreadyFinalized,
/// The input was successfully finalized during this call.
Finalized,
/// The wallet could not derive a descriptor for the input.
MissingDescriptor,
/// The wallet found the descriptor but could not construct the input satisfaction.
CouldNotSatisfy(miniscript::Error),
}

impl FinalizeInputOutcome {
/// Whether the input is finalized after this call.
pub fn is_finalized(&self) -> bool {
matches!(self, Self::AlreadyFinalized | Self::Finalized)
}
}

/// Holds per-input PSBT finalization outcomes.
#[derive(Debug, PartialEq)]
pub struct FinalizedInputs {
outcomes: BTreeMap<usize, FinalizeInputOutcome>,
}

impl FinalizedInputs {
fn new(outcomes: BTreeMap<usize, FinalizeInputOutcome>) -> Self {
Self { outcomes }
}

/// Whether all inputs are finalized after this call.
pub fn is_finalized(&self) -> bool {
self.outcomes
.values()
.all(FinalizeInputOutcome::is_finalized)
}

/// Borrow the per-input finalization outcomes.
pub fn outcomes(&self) -> &BTreeMap<usize, FinalizeInputOutcome> {
&self.outcomes
}

/// Consume the collection and return the per-input finalization outcomes.
pub fn into_outcomes(self) -> BTreeMap<usize, FinalizeInputOutcome> {
self.outcomes
}
}

impl Wallet {
/// Build a new single descriptor [`Wallet`].
///
Expand Down Expand Up @@ -1866,23 +1915,55 @@ impl Wallet {
}
})
.collect::<HashMap<Txid, u32>>();
let current_height = sign_options
.assume_height
.unwrap_or_else(|| self.chain.tip().height());

Ok(self
.try_finalize_psbt_with(psbt, |_, input| {
Some((
current_height,
confirmation_heights
.get(&input.previous_output.txid)
.copied(),
))
})?
.is_finalized())
}

/// Finalize a PSBT and return per-input finalization results. Use this method when you need to
/// inspect why a specific input could not be finalized.
///
/// The method should only return `Err` when the PSBT is malformed, for example if its inputs
/// are out of bounds.
pub fn try_finalize_psbt(
&self,
psbt: &mut Psbt,
) -> Result<FinalizedInputs, IndexOutOfBoundsError> {
self.try_finalize_psbt_with(psbt, |_, _| None)
}

let mut finished = true;
fn try_finalize_psbt_with<F>(
&self,
psbt: &mut Psbt,
mut wallet_timelocks: F,
) -> Result<FinalizedInputs, IndexOutOfBoundsError>
where
F: FnMut(usize, &bitcoin::TxIn) -> Option<(u32, Option<u32>)>,
{
let tx = &psbt.unsigned_tx;

let mut outcomes = BTreeMap::new();

for (n, input) in tx.input.iter().enumerate() {
let psbt_input = &psbt
.inputs
.get(n)
.ok_or(IndexOutOfBoundsError::new(n, psbt.inputs.len()))?;
if psbt_input.final_script_sig.is_some() || psbt_input.final_script_witness.is_some() {
outcomes.insert(n, FinalizeInputOutcome::AlreadyFinalized);
continue;
}
let confirmation_height = confirmation_heights
.get(&input.previous_output.txid)
.copied();
let current_height = sign_options
.assume_height
.unwrap_or_else(|| self.chain.tip().height());

// - Try to derive the descriptor by looking at the txout. If it's in our database, we
// know exactly which `keychain` to use, and which derivation index it is.
Expand All @@ -1902,14 +1983,22 @@ impl Wallet {
match desc {
Some(desc) => {
let mut tmp_input = bitcoin::TxIn::default();
match desc.satisfy(
&mut tmp_input,
(
PsbtInputSatisfier::new(psbt, n),
After::new(Some(current_height), false),
Older::new(Some(current_height), confirmation_height, false),
),
) {
let satisfy_result = if let Some((current_height, confirmation_height)) =
wallet_timelocks(n, input)
{
desc.satisfy(
&mut tmp_input,
(
PsbtInputSatisfier::new(psbt, n),
After::new(Some(current_height), false),
Older::new(Some(current_height), confirmation_height, false),
),
)
} else {
desc.satisfy(&mut tmp_input, PsbtInputSatisfier::new(psbt, n))
};

match satisfy_result {
Ok(_) => {
let length = psbt.inputs.len();
// Set the UTXO fields, final script_sig and witness
Expand All @@ -1927,23 +2016,29 @@ impl Wallet {
if !tmp_input.witness.is_empty() {
psbt_input.final_script_witness = Some(tmp_input.witness);
}
outcomes.insert(n, FinalizeInputOutcome::Finalized);
}
Err(err) => {
outcomes.insert(n, FinalizeInputOutcome::CouldNotSatisfy(err));
}
Err(_) => finished = false,
}
}
None => finished = false,
None => {
outcomes.insert(n, FinalizeInputOutcome::MissingDescriptor);
}
}
}

// Clear derivation paths from outputs.
if finished {
let finalized = FinalizedInputs::new(outcomes);
if finalized.is_finalized() {
for output in &mut psbt.outputs {
output.bip32_derivation.clear();
output.tap_key_origins.clear();
}
}

Ok(finished)
Ok(finalized)
}

/// Return the secp256k1 context used for all signing operations.
Expand Down
184 changes: 183 additions & 1 deletion tests/wallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@ use bdk_wallet::psbt::PsbtUtils;
use bdk_wallet::signer::{SignOptions, SignerError};
use bdk_wallet::test_utils::*;
use bdk_wallet::KeychainKind;
use bdk_wallet::{AddressInfo, Balance, PersistedWallet, Update, Wallet, WalletTx};
use bdk_wallet::{
AddressInfo, Balance, FinalizeInputOutcome, IndexOutOfBoundsError, PersistedWallet, Update,
Wallet, WalletTx,
};
use bitcoin::constants::COINBASE_MATURITY;
use bitcoin::hashes::Hash;
use bitcoin::script::PushBytesBuf;
Expand Down Expand Up @@ -1560,6 +1563,185 @@ fn test_try_finalize_sign_option() {
}
}

#[test]
fn test_try_finalize_psbt_outcomes() {
{
let (mut wallet, _) = get_funded_wallet_single(get_test_wpkh());
let addr = wallet.next_unused_address(KeychainKind::External);
let mut builder = wallet.build_tx();
builder.drain_to(addr.script_pubkey()).drain_wallet();
let mut psbt = builder.finish().unwrap();

let is_final = wallet
.sign(
&mut psbt,
SignOptions {
try_finalize: false,
..Default::default()
},
)
.unwrap();
assert!(!is_final);
assert!(
psbt.outputs
.iter()
.any(|output| !output.bip32_derivation.is_empty()),
"expected wallet-owned outputs to retain derivation data before finalization"
);

let finalized = wallet.try_finalize_psbt(&mut psbt).unwrap();

assert!(finalized.is_finalized());
assert_matches!(
finalized.outcomes().get(&0),
Some(FinalizeInputOutcome::Finalized)
);
assert!(
psbt.inputs[0].final_script_sig.is_some()
|| psbt.inputs[0].final_script_witness.is_some()
);
assert!(psbt
.outputs
.iter()
.all(|output| output.bip32_derivation.is_empty()));

let finalized = wallet.try_finalize_psbt(&mut psbt).unwrap();

assert!(finalized.is_finalized());
assert_matches!(
finalized.outcomes().get(&0),
Some(FinalizeInputOutcome::AlreadyFinalized)
);
}

{
let (mut wallet, _) = get_funded_wallet_single(get_test_wpkh());
let addr = wallet.next_unused_address(KeychainKind::External);
let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
let mut psbt = builder.finish().unwrap();

let dud_input = bitcoin::psbt::Input {
witness_utxo: Some(TxOut {
value: Amount::from_sat(100_000),
script_pubkey: miniscript::Descriptor::<bitcoin::PublicKey>::from_str(
"wpkh(025476c2e83188368da1ff3e292e7acafcdb3566bb0ad253f62fc70f07aeee6357)",
)
.unwrap()
.script_pubkey(),
}),
..Default::default()
};

psbt.inputs.push(dud_input);
psbt.unsigned_tx.input.push(bitcoin::TxIn::default());

let finalized = wallet.try_finalize_psbt(&mut psbt).unwrap();

assert!(!finalized.is_finalized());
assert_matches!(
finalized.outcomes().get(&0),
Some(FinalizeInputOutcome::CouldNotSatisfy(
bdk_wallet::miniscript::Error::MissingSig(_)
))
);
assert_matches!(
finalized.outcomes().get(&1),
Some(FinalizeInputOutcome::MissingDescriptor)
);
}
}

#[test]
fn test_try_finalize_psbt_returns_index_out_of_bounds_for_malformed_psbt() {
let (mut wallet, _) = get_funded_wallet_single(get_test_wpkh());
let addr = wallet.next_unused_address(KeychainKind::External);
let mut builder = wallet.build_tx();
builder.drain_to(addr.script_pubkey()).drain_wallet();
let mut psbt = builder.finish().unwrap();

psbt.inputs.clear();

let err = wallet.try_finalize_psbt(&mut psbt).unwrap_err();

assert_eq!(err, IndexOutOfBoundsError::new(0, 0));
}

#[test]
fn test_try_finalize_psbt_uses_psbt_timelocks() {
{
let (mut wallet, _) = get_funded_wallet_single(get_test_single_sig_cltv());
let addr = wallet.next_unused_address(KeychainKind::External);
let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
let mut psbt = builder.finish().unwrap();

wallet
.sign(
&mut psbt,
SignOptions {
try_finalize: false,
..Default::default()
},
)
.unwrap();

let mut valid_psbt = psbt.clone();
assert!(wallet
.try_finalize_psbt(&mut valid_psbt)
.unwrap()
.is_finalized());

psbt.unsigned_tx.lock_time = absolute::LockTime::from_height(0).unwrap();

let finalized = wallet.try_finalize_psbt(&mut psbt).unwrap();

assert!(!finalized.is_finalized());
assert_matches!(
finalized.outcomes().get(&0),
Some(FinalizeInputOutcome::CouldNotSatisfy(_))
);
assert!(psbt.inputs[0].final_script_sig.is_none());
assert!(psbt.inputs[0].final_script_witness.is_none());
}

{
let (mut wallet, _) = get_funded_wallet_single(get_test_single_sig_csv());
let addr = wallet.next_unused_address(KeychainKind::External);
let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), Amount::from_sat(25_000));
let mut psbt = builder.finish().unwrap();

wallet
.sign(
&mut psbt,
SignOptions {
try_finalize: false,
..Default::default()
},
)
.unwrap();

let mut valid_psbt = psbt.clone();
assert!(wallet
.try_finalize_psbt(&mut valid_psbt)
.unwrap()
.is_finalized());

psbt.unsigned_tx.input[0].sequence = Sequence::MAX;

let finalized = wallet.try_finalize_psbt(&mut psbt).unwrap();

assert!(!finalized.is_finalized());
assert_matches!(
finalized.outcomes().get(&0),
Some(FinalizeInputOutcome::CouldNotSatisfy(_))
);
assert!(psbt.inputs[0].final_script_sig.is_none());
assert!(psbt.inputs[0].final_script_witness.is_none());
}
}

#[test]
fn test_taproot_try_finalize_sign_option() {
let (mut wallet, _) = get_funded_wallet_single(get_test_tr_with_taptree());
Expand Down