Skip to content

Commit

Permalink
Allow signing only specific leaf hashes
Browse files Browse the repository at this point in the history 
Fixes #616
  • Loading branch information
@danielabrozzoni
danielabrozzoni committed Jul 13, 2022
1 parent dd51380 commit b8459a8
Show file tree
Hide file tree
Showing 3 changed files with 214 additions and 23 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [Unreleased]
- Add the ability to specify which leaves to sign in a taproot transaction through `TapLeavesOptions` in `SignOptions`
- New MSRV set to `1.56.1`
- Fee sniping discouraging through nLockTime - if the user specifies a `current_height`, we use that as a nlocktime, otherwise we use the last sync height (or 0 if we never synced)
- Fix hang when `ElectrumBlockchainConfig::stop_gap` is zero.
Expand Down
142 changes: 141 additions & 1 deletion src/wallet/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1071,7 +1071,7 @@ where
.iter()
.chain(self.change_signers.signers().iter())
{
signer.sign_transaction(psbt, &self.secp)?;
signer.sign_transaction(psbt, &sign_options.tap_leaves_options, &self.secp)?;
}

// attempt to finalize
Expand Down Expand Up @@ -1917,6 +1917,10 @@ pub(crate) mod test {
"tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
}

pub(crate) fn get_test_tr_with_taptree_both_priv() -> &'static str {
"tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh),pk(cNaQCDwmmh4dS9LzCgVtyy1e1xjCJ21GUDHe9K98nzb689JvinGV)})"
}

pub(crate) fn get_test_tr_repeated_key() -> &'static str {
"tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(100)),and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(200))})"
}
Expand Down Expand Up @@ -4753,6 +4757,142 @@ pub(crate) mod test {
test_spend_from_wallet(wallet);
}

#[test]
fn test_taproot_script_spend_sign_all_leaves() {
use crate::signer::TapLeavesOptions;
let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
let addr = wallet.get_address(AddressIndex::New).unwrap();

let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), 25_000);
let (mut psbt, _) = builder.finish().unwrap();

assert!(
wallet
.sign(
&mut psbt,
SignOptions {
tap_leaves_options: TapLeavesOptions::All,
..Default::default()
},
)
.unwrap(),
"Unable to finalize tx"
);

assert!(psbt
.inputs
.iter()
.all(|i| i.tap_script_sigs.len() == i.tap_scripts.len()));
}

#[test]
fn test_taproot_script_spend_sign_include_some_leaves() {
use crate::signer::TapLeavesOptions;
use crate::wallet::taproot::TapLeafHash;

let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
let addr = wallet.get_address(AddressIndex::New).unwrap();

let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), 25_000);
let (mut psbt, _) = builder.finish().unwrap();
let mut script_leaves: Vec<_> = psbt.inputs[0]
.tap_scripts
.clone()
.values()
.map(|(script, version)| TapLeafHash::from_script(script, *version))
.collect();
let included_script_leaves = vec![script_leaves.pop().unwrap()];
let excluded_script_leaves = script_leaves;

assert!(
wallet
.sign(
&mut psbt,
SignOptions {
tap_leaves_options: TapLeavesOptions::Include(
included_script_leaves.clone()
),
..Default::default()
},
)
.unwrap(),
"Unable to finalize tx"
);

assert!(psbt.inputs[0]
.tap_script_sigs
.iter()
.all(|s| included_script_leaves.contains(&s.0 .1)
&& !excluded_script_leaves.contains(&s.0 .1)));
}

#[test]
fn test_taproot_script_spend_sign_exclude_some_leaves() {
use crate::signer::TapLeavesOptions;
use crate::wallet::taproot::TapLeafHash;

let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
let addr = wallet.get_address(AddressIndex::New).unwrap();

let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), 25_000);
let (mut psbt, _) = builder.finish().unwrap();
let mut script_leaves: Vec<_> = psbt.inputs[0]
.tap_scripts
.clone()
.values()
.map(|(script, version)| TapLeafHash::from_script(script, *version))
.collect();
let included_script_leaves = vec![script_leaves.pop().unwrap()];
let excluded_script_leaves = script_leaves;

assert!(
wallet
.sign(
&mut psbt,
SignOptions {
tap_leaves_options: TapLeavesOptions::Exclude(
excluded_script_leaves.clone()
),
..Default::default()
},
)
.unwrap(),
"Unable to finalize tx"
);

assert!(psbt.inputs[0]
.tap_script_sigs
.iter()
.all(|s| included_script_leaves.contains(&s.0 .1)
&& !excluded_script_leaves.contains(&s.0 .1)));
}

#[test]
fn test_taproot_script_spend_sign_no_leaves() {
use crate::signer::TapLeavesOptions;
let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
let addr = wallet.get_address(AddressIndex::New).unwrap();

let mut builder = wallet.build_tx();
builder.add_recipient(addr.script_pubkey(), 25_000);
let (mut psbt, _) = builder.finish().unwrap();

wallet
.sign(
&mut psbt,
SignOptions {
tap_leaves_options: TapLeavesOptions::None,
..Default::default()
},
)
.unwrap();

assert!(psbt.inputs.iter().all(|i| i.tap_script_sigs.is_empty()));
}

#[test]
fn test_taproot_sign_derive_index_from_psbt() {
let (wallet, _, _) = get_funded_wallet(get_test_tr_single_sig_xprv());
Expand Down
94 changes: 72 additions & 22 deletions src/wallet/signer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
//! &self,
//! psbt: &mut psbt::PartiallySignedTransaction,
//! input_index: usize,
//! _tap_leaves_options: &TapLeavesOptions,
//! _secp: &Secp256k1<All>,
//! ) -> Result<(), SignerError> {
//! self.device.hsm_sign_input(psbt, input_index)?;
Expand Down Expand Up @@ -241,6 +242,7 @@ pub trait InputSigner: SignerCommon {
&self,
psbt: &mut psbt::PartiallySignedTransaction,
input_index: usize,
tap_leaves_options: &TapLeavesOptions,
secp: &SecpCtx,
) -> Result<(), SignerError>;
}
Expand All @@ -254,6 +256,7 @@ pub trait TransactionSigner: SignerCommon {
fn sign_transaction(
&self,
psbt: &mut psbt::PartiallySignedTransaction,
tap_leaves_options: &TapLeavesOptions,
secp: &SecpCtx,
) -> Result<(), SignerError>;
}
Expand All @@ -262,10 +265,11 @@ impl<T: InputSigner> TransactionSigner for T {
fn sign_transaction(
&self,
psbt: &mut psbt::PartiallySignedTransaction,
tap_leaves_options: &TapLeavesOptions,
secp: &SecpCtx,
) -> Result<(), SignerError> {
for input_index in 0..psbt.inputs.len() {
self.sign_input(psbt, input_index, secp)?;
self.sign_input(psbt, input_index, tap_leaves_options, secp)?;
}

Ok(())
Expand All @@ -287,6 +291,7 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
&self,
psbt: &mut psbt::PartiallySignedTransaction,
input_index: usize,
tap_leaves_options: &TapLeavesOptions,
secp: &SecpCtx,
) -> Result<(), SignerError> {
if input_index >= psbt.inputs.len() {
Expand Down Expand Up @@ -346,7 +351,12 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
inner: derived_key.private_key,
};

SignerWrapper::new(priv_key, self.ctx).sign_input(psbt, input_index, secp)
SignerWrapper::new(priv_key, self.ctx).sign_input(
psbt,
input_index,
tap_leaves_options,
secp,
)
}
}
}
Expand All @@ -369,6 +379,7 @@ impl InputSigner for SignerWrapper<PrivateKey> {
&self,
psbt: &mut psbt::PartiallySignedTransaction,
input_index: usize,
tap_leaves_options: &TapLeavesOptions,
secp: &SecpCtx,
) -> Result<(), SignerError> {
if input_index >= psbt.inputs.len() || input_index >= psbt.unsigned_tx.input.len() {
Expand Down Expand Up @@ -401,26 +412,36 @@ impl InputSigner for SignerWrapper<PrivateKey> {
if let Some((leaf_hashes, _)) =
psbt.inputs[input_index].tap_key_origins.get(&x_only_pubkey)
{
let leaf_hashes = leaf_hashes
.iter()
.filter(|lh| {
!psbt.inputs[input_index]
.tap_script_sigs
.contains_key(&(x_only_pubkey, **lh))
})
.cloned()
.collect::<Vec<_>>();
for lh in leaf_hashes {
let (hash, hash_ty) = Tap::sighash(psbt, input_index, Some(lh))?;
sign_psbt_schnorr(
&self.inner,
x_only_pubkey,
Some(lh),
&mut psbt.inputs[input_index],
hash,
hash_ty,
secp,
);
if *tap_leaves_options != TapLeavesOptions::None {
let leaf_hashes = leaf_hashes
.iter()
.filter(|lh| {
// Removing the leaves we shouldn't sign for
let should_sign = match tap_leaves_options {
TapLeavesOptions::Include(v) => v.contains(lh),
TapLeavesOptions::Exclude(v) => !v.contains(lh),
_ => true,
};
// Filtering out the leaves without our key
should_sign
&& !psbt.inputs[input_index]
.tap_script_sigs
.contains_key(&(x_only_pubkey, **lh))
})
.cloned()
.collect::<Vec<_>>();
for lh in leaf_hashes {
let (hash, hash_ty) = Tap::sighash(psbt, input_index, Some(lh))?;
sign_psbt_schnorr(
&self.inner,
x_only_pubkey,
Some(lh),
&mut psbt.inputs[input_index],
hash,
hash_ty,
secp,
);
}
}
}

Expand Down Expand Up @@ -675,6 +696,33 @@ pub struct SignOptions {
///
/// Defaults to `true` which will try fianlizing psbt after inputs are signed.
pub try_finalize: bool,

/// Specifies which Taproot script-spend leaves we should sign for. This option is ignored if we're signing
/// a non-taproot PSBT.
///
/// Defaults to All, i.e., the wallet will sign all the leaves it has a key for.
pub tap_leaves_options: TapLeavesOptions,
}

/// Customize which taproot script-path leaves the signer should sign.
#[derive(Debug, Clone, PartialEq)]
pub enum TapLeavesOptions {
/// The signer will sign all the leaves it has a key for.
All,
/// The signer won't sign leaves other than the ones specified. Note that it could still ignore
/// some of the specified leaves, if it doesn't have the right key to sign them.
Include(Vec<taproot::TapLeafHash>),
/// The signer won't sign the specified leaves.
Exclude(Vec<taproot::TapLeafHash>),
/// The signer won't sign any leaf. This effectively means that it will try signing only for
/// the key-path.
None,
}

impl Default for TapLeavesOptions {
fn default() -> Self {
TapLeavesOptions::All
}
}

#[allow(clippy::derivable_impls)]
Expand All @@ -686,6 +734,7 @@ impl Default for SignOptions {
allow_all_sighashes: false,
remove_partial_sigs: true,
try_finalize: true,
tap_leaves_options: TapLeavesOptions::default(),
}
}
}
Expand Down Expand Up @@ -1016,6 +1065,7 @@ mod signers_container_tests {
fn sign_transaction(
&self,
_psbt: &mut psbt::PartiallySignedTransaction,
_tap_leaves_options: &TapLeavesOptions,
_secp: &SecpCtx,
) -> Result<(), SignerError> {
Ok(())
Expand Down

0 comments on commit b8459a8

Please sign in to comment.