Introduce WitnessProgram struct and cleanup Address validity invariants

Addresses with Segwitv0 not having len 20/32 are invalid and cannot be
constructed. Also cleans up a API bug in
ScriptBuf::new_witness_prog(ver, prog) allowing prog of invalid lenghts.
This commit is contained in:
sanket1729 2023-01-19 15:30:17 -08:00
parent 41652caf05
commit 6ebc9de252
3 changed files with 119 additions and 67 deletions

View File

@ -392,12 +392,41 @@ pub enum Payload {
/// P2SH address. /// P2SH address.
ScriptHash(ScriptHash), ScriptHash(ScriptHash),
/// Segwit address. /// Segwit address.
WitnessProgram { WitnessProgram(WitnessProgram),
/// The witness program version. }
version: WitnessVersion,
/// The witness program. /// Witness program as defined in BIP141.
program: Vec<u8>, #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
}, pub struct WitnessProgram {
/// The witness program version.
version: WitnessVersion,
/// The witness program. (Between 2 and 40 bytes)
program: Vec<u8>,
}
impl WitnessProgram {
/// Creates a new witness program.
pub fn new(version: WitnessVersion, program: Vec<u8>) -> Result<Self, Error> {
if program.len() < 2 || program.len() > 40 {
return Err(Error::InvalidWitnessProgramLength(program.len()));
}
// Specific segwit v0 check. These addresses can never spend funds sent to them.
if version == WitnessVersion::V0 && (program.len() != 20 && program.len() != 32) {
return Err(Error::InvalidSegwitV0ProgramLength(program.len()));
}
Ok(WitnessProgram { version, program })
}
/// Returns the witness program version.
pub fn version(&self) -> WitnessVersion {
self.version
}
/// Returns the witness program.
pub fn program(&self) -> &[u8] {
&self.program
}
} }
impl Payload { impl Payload {
@ -412,18 +441,13 @@ impl Payload {
hash_inner.copy_from_slice(&script.as_bytes()[2..22]); hash_inner.copy_from_slice(&script.as_bytes()[2..22]);
Payload::ScriptHash(ScriptHash::from_inner(hash_inner)) Payload::ScriptHash(ScriptHash::from_inner(hash_inner))
} else if script.is_witness_program() { } else if script.is_witness_program() {
if script.witness_version() == Some(WitnessVersion::V0)
&& !(script.is_v0_p2wpkh() || script.is_v0_p2wsh())
{
return Err(Error::InvalidSegwitV0ProgramLength(script.len() - 2));
}
let opcode = script.first_opcode().expect("witness_version guarantees len() > 4"); let opcode = script.first_opcode().expect("witness_version guarantees len() > 4");
Payload::WitnessProgram { let witness_program = WitnessProgram::new(
version: WitnessVersion::try_from(opcode)?, WitnessVersion::try_from(opcode)?,
program: script.as_bytes()[2..].to_vec(), script.as_bytes()[2..].to_vec(),
} )?;
Payload::WitnessProgram(witness_program)
} else { } else {
return Err(Error::UnrecognizedScript); return Err(Error::UnrecognizedScript);
}) })
@ -434,8 +458,8 @@ impl Payload {
match *self { match *self {
Payload::PubkeyHash(ref hash) => script::ScriptBuf::new_p2pkh(hash), Payload::PubkeyHash(ref hash) => script::ScriptBuf::new_p2pkh(hash),
Payload::ScriptHash(ref hash) => script::ScriptBuf::new_p2sh(hash), Payload::ScriptHash(ref hash) => script::ScriptBuf::new_p2sh(hash),
Payload::WitnessProgram { version, program: ref prog } => Payload::WitnessProgram(ref prog) =>
script::ScriptBuf::new_witness_program(version, prog) script::ScriptBuf::new_witness_program(prog)
} }
} }
@ -454,10 +478,11 @@ impl Payload {
/// Create a witness pay to public key payload from a public key /// Create a witness pay to public key payload from a public key
pub fn p2wpkh(pk: &PublicKey) -> Result<Payload, Error> { pub fn p2wpkh(pk: &PublicKey) -> Result<Payload, Error> {
Ok(Payload::WitnessProgram { let prog = WitnessProgram::new(
version: WitnessVersion::V0, WitnessVersion::V0,
program: pk.wpubkey_hash().ok_or(Error::UncompressedPubkey)?.as_ref().to_vec(), pk.wpubkey_hash().ok_or(Error::UncompressedPubkey)?.as_ref().to_vec()
}) )?;
Ok(Payload::WitnessProgram(prog))
} }
/// Create a pay to script payload that embeds a witness pay to public key /// Create a pay to script payload that embeds a witness pay to public key
@ -471,10 +496,11 @@ impl Payload {
/// Create a witness pay to script hash payload. /// Create a witness pay to script hash payload.
pub fn p2wsh(script: &script::Script) -> Payload { pub fn p2wsh(script: &script::Script) -> Payload {
Payload::WitnessProgram { let prog = WitnessProgram::new(
version: WitnessVersion::V0, WitnessVersion::V0,
program: script.wscript_hash().as_ref().to_vec(), script.wscript_hash().as_ref().to_vec()
} ).expect("wscript_hash has len 32 compatible with segwitv0");
Payload::WitnessProgram(prog)
} }
/// Create a pay to script payload that embeds a witness pay to script hash address /// Create a pay to script payload that embeds a witness pay to script hash address
@ -492,20 +518,22 @@ impl Payload {
merkle_root: Option<TapNodeHash>, merkle_root: Option<TapNodeHash>,
) -> Payload { ) -> Payload {
let (output_key, _parity) = internal_key.tap_tweak(secp, merkle_root); let (output_key, _parity) = internal_key.tap_tweak(secp, merkle_root);
Payload::WitnessProgram { let prog = WitnessProgram::new(
version: WitnessVersion::V1, WitnessVersion::V1,
program: output_key.to_inner().serialize().to_vec(), output_key.to_inner().serialize().to_vec()
} ).expect("taproot output key has len 32 <= 40");
Payload::WitnessProgram(prog)
} }
/// Create a pay to taproot payload from a pre-tweaked output key. /// Create a pay to taproot payload from a pre-tweaked output key.
/// ///
/// This method is not recommended for use and [Payload::p2tr()] should be used where possible. /// This method is not recommended for use and [Payload::p2tr()] should be used where possible.
pub fn p2tr_tweaked(output_key: TweakedPublicKey) -> Payload { pub fn p2tr_tweaked(output_key: TweakedPublicKey) -> Payload {
Payload::WitnessProgram { let prog = WitnessProgram::new(
version: WitnessVersion::V1, WitnessVersion::V1,
program: output_key.to_inner().serialize().to_vec(), output_key.to_inner().serialize().to_vec()
} ).expect("taproot output key has len 32 <= 40");
Payload::WitnessProgram(prog)
} }
/// Returns a byte slice of the payload /// Returns a byte slice of the payload
@ -513,7 +541,7 @@ impl Payload {
match self { match self {
Payload::ScriptHash(hash) => hash.as_ref(), Payload::ScriptHash(hash) => hash.as_ref(),
Payload::PubkeyHash(hash) => hash.as_ref(), Payload::PubkeyHash(hash) => hash.as_ref(),
Payload::WitnessProgram { program, .. } => program, Payload::WitnessProgram(prog) => prog.program(),
} }
} }
} }
@ -546,7 +574,8 @@ impl<'a> fmt::Display for AddressEncoding<'a> {
prefixed[1..].copy_from_slice(&hash[..]); prefixed[1..].copy_from_slice(&hash[..]);
base58::encode_check_to_fmt(fmt, &prefixed[..]) base58::encode_check_to_fmt(fmt, &prefixed[..])
} }
Payload::WitnessProgram { version, program: prog } => { Payload::WitnessProgram(witness_prog) => {
let (version, prog) = (witness_prog.version(), witness_prog.program());
let mut upper_writer; let mut upper_writer;
let writer = if fmt.alternate() { let writer = if fmt.alternate() {
upper_writer = UpperWriter(fmt); upper_writer = UpperWriter(fmt);
@ -556,7 +585,7 @@ impl<'a> fmt::Display for AddressEncoding<'a> {
}; };
let mut bech32_writer = let mut bech32_writer =
bech32::Bech32Writer::new(self.bech32_hrp, version.bech32_variant(), writer)?; bech32::Bech32Writer::new(self.bech32_hrp, version.bech32_variant(), writer)?;
bech32::WriteBase32::write_u5(&mut bech32_writer, (*version).into())?; bech32::WriteBase32::write_u5(&mut bech32_writer, version.into())?;
bech32::ToBase32::write_base32(&prog, &mut bech32_writer) bech32::ToBase32::write_base32(&prog, &mut bech32_writer)
} }
} }
@ -725,15 +754,15 @@ impl<V: NetworkValidation> Address<V> {
match self.payload { match self.payload {
Payload::PubkeyHash(_) => Some(AddressType::P2pkh), Payload::PubkeyHash(_) => Some(AddressType::P2pkh),
Payload::ScriptHash(_) => Some(AddressType::P2sh), Payload::ScriptHash(_) => Some(AddressType::P2sh),
Payload::WitnessProgram { version, program: ref prog } => { Payload::WitnessProgram(ref prog) => {
// BIP-141 p2wpkh or p2wsh addresses. // BIP-141 p2wpkh or p2wsh addresses.
match version { match prog.version() {
WitnessVersion::V0 => match prog.len() { WitnessVersion::V0 => match prog.program().len() {
20 => Some(AddressType::P2wpkh), 20 => Some(AddressType::P2wpkh),
32 => Some(AddressType::P2wsh), 32 => Some(AddressType::P2wsh),
_ => None, _ => unreachable!("Address creation invariant violation: invalid program length")
}, },
WitnessVersion::V1 if prog.len() == 32 => Some(AddressType::P2tr), WitnessVersion::V1 if prog.program().len() == 32 => Some(AddressType::P2tr),
_ => None, _ => None,
} }
} }
@ -849,12 +878,18 @@ impl Address {
} }
/// Checks whether or not the address is following Bitcoin standardness rules when /// Checks whether or not the address is following Bitcoin standardness rules when
/// *spending* from this address. This method should *NOT* called by senders. For forward /// *spending* from this address. *NOT* to be called by senders.
/// compatibility, the senders must send to any [`Address`]. Receivers can use this method to ///
/// check whether or not they can spend from this address. /// <details>
/// <summary>Spending Standardness</summary>
///
/// For forward compatibility, the senders must send to any [`Address`]. Receivers
/// can use this method to check whether or not they can spend from this address.
/// ///
/// SegWit addresses with unassigned witness versions or non-standard program sizes are /// SegWit addresses with unassigned witness versions or non-standard program sizes are
/// considered non-standard. /// considered non-standard.
/// </details>
///
pub fn is_spend_standard(&self) -> bool { self.address_type().is_some() } pub fn is_spend_standard(&self) -> bool { self.address_type().is_some() }
/// Checks whether or not the address is following Bitcoin standardness rules. /// Checks whether or not the address is following Bitcoin standardness rules.
@ -1041,14 +1076,7 @@ impl FromStr for Address<NetworkUnchecked> {
(WitnessVersion::try_from(v[0])?, bech32::FromBase32::from_base32(p5)?) (WitnessVersion::try_from(v[0])?, bech32::FromBase32::from_base32(p5)?)
}; };
if program.len() < 2 || program.len() > 40 { let witness_program = WitnessProgram::new(version, program)?;
return Err(Error::InvalidWitnessProgramLength(program.len()));
}
// Specific segwit v0 check.
if version == WitnessVersion::V0 && (program.len() != 20 && program.len() != 32) {
return Err(Error::InvalidSegwitV0ProgramLength(program.len()));
}
// Encoding check // Encoding check
let expected = version.bech32_variant(); let expected = version.bech32_variant();
@ -1056,7 +1084,7 @@ impl FromStr for Address<NetworkUnchecked> {
return Err(Error::InvalidBech32Variant { expected, found: variant }); return Err(Error::InvalidBech32Variant { expected, found: variant });
} }
return Ok(Address::new(network, Payload::WitnessProgram { version, program })); return Ok(Address::new(network, Payload::WitnessProgram(witness_program)));
} }
// Base58 // Base58
@ -1238,7 +1266,8 @@ mod tests {
let program = hex!( let program = hex!(
"654f6ea368e0acdfd92976b7c2103a1b26313f430654f6ea368e0acdfd92976b7c2103a1b26313f4" "654f6ea368e0acdfd92976b7c2103a1b26313f430654f6ea368e0acdfd92976b7c2103a1b26313f4"
); );
let addr = Address::new(Bitcoin, Payload::WitnessProgram { version: WitnessVersion::V13, program }); let witness_prog = WitnessProgram::new(WitnessVersion::V13, program).unwrap();
let addr = Address::new(Bitcoin, Payload::WitnessProgram(witness_prog));
roundtrips(&addr); roundtrips(&addr);
} }
@ -1449,10 +1478,10 @@ mod tests {
Payload::ScriptHash(ScriptHash::all_zeros()), Payload::ScriptHash(ScriptHash::all_zeros()),
]; ];
let segwit_payload = (0..=16) let segwit_payload = (0..=16)
.map(|version| Payload::WitnessProgram { .map(|version| Payload::WitnessProgram(WitnessProgram::new(
version: WitnessVersion::try_from(version).unwrap(), WitnessVersion::try_from(version).unwrap(),
program: vec![], vec![0xab; 32], // Choose 32 to make test case valid for all witness versions(including v0)
}) ).unwrap()))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
const LEGACY_EQUIVALENCE_CLASSES: &[&[Network]] = const LEGACY_EQUIVALENCE_CLASSES: &[&[Network]] =

View File

@ -74,7 +74,7 @@ use crate::policy::DUST_RELAY_TX_FEE;
use crate::OutPoint; use crate::OutPoint;
use crate::key::PublicKey; use crate::key::PublicKey;
use crate::address::WitnessVersion; use crate::address::{WitnessVersion, WitnessProgram};
use crate::taproot::{LeafVersion, TapNodeHash, TapLeafHash}; use crate::taproot::{LeafVersion, TapNodeHash, TapLeafHash};
use secp256k1::{Secp256k1, Verification, XOnlyPublicKey}; use secp256k1::{Secp256k1, Verification, XOnlyPublicKey};
use crate::schnorr::{TapTweak, TweakedPublicKey, UntweakedPublicKey}; use crate::schnorr::{TapTweak, TweakedPublicKey, UntweakedPublicKey};
@ -1122,28 +1122,47 @@ impl ScriptBuf {
/// Generates P2WPKH-type of scriptPubkey. /// Generates P2WPKH-type of scriptPubkey.
pub fn new_v0_p2wpkh(pubkey_hash: &WPubkeyHash) -> Self { pub fn new_v0_p2wpkh(pubkey_hash: &WPubkeyHash) -> Self {
ScriptBuf::new_witness_program(WitnessVersion::V0, &pubkey_hash[..]) // pubkey hash is 20 bytes long, so it's safe to use `new_witness_program_unchecked` (Segwitv0)
ScriptBuf::new_witness_program_unchecked(WitnessVersion::V0, &pubkey_hash[..])
} }
/// Generates P2WSH-type of scriptPubkey with a given hash of the redeem script. /// Generates P2WSH-type of scriptPubkey with a given hash of the redeem script.
pub fn new_v0_p2wsh(script_hash: &WScriptHash) -> Self { pub fn new_v0_p2wsh(script_hash: &WScriptHash) -> Self {
ScriptBuf::new_witness_program(WitnessVersion::V0, &script_hash[..]) // script hash is 32 bytes long, so it's safe to use `new_witness_program_unchecked` (Segwitv0)
ScriptBuf::new_witness_program_unchecked(WitnessVersion::V0, &script_hash[..])
} }
/// Generates P2TR for script spending path using an internal public key and some optional /// Generates P2TR for script spending path using an internal public key and some optional
/// script tree merkle root. /// script tree merkle root.
pub fn new_v1_p2tr<C: Verification>(secp: &Secp256k1<C>, internal_key: UntweakedPublicKey, merkle_root: Option<TapNodeHash>) -> Self { pub fn new_v1_p2tr<C: Verification>(secp: &Secp256k1<C>, internal_key: UntweakedPublicKey, merkle_root: Option<TapNodeHash>) -> Self {
let (output_key, _) = internal_key.tap_tweak(secp, merkle_root); let (output_key, _) = internal_key.tap_tweak(secp, merkle_root);
ScriptBuf::new_witness_program(WitnessVersion::V1, &output_key.serialize()) // output key is 32 bytes long, so it's safe to use `new_witness_program_unchecked` (Segwitv1)
ScriptBuf::new_witness_program_unchecked(WitnessVersion::V1, &output_key.serialize())
} }
/// Generates P2TR for key spending path for a known [`TweakedPublicKey`]. /// Generates P2TR for key spending path for a known [`TweakedPublicKey`].
pub fn new_v1_p2tr_tweaked(output_key: TweakedPublicKey) -> Self { pub fn new_v1_p2tr_tweaked(output_key: TweakedPublicKey) -> Self {
ScriptBuf::new_witness_program(WitnessVersion::V1, &output_key.serialize()) // output key is 32 bytes long, so it's safe to use `new_witness_program_unchecked` (Segwitv1)
ScriptBuf::new_witness_program_unchecked(WitnessVersion::V1, &output_key.serialize())
} }
/// Generates P2WSH-type of scriptPubkey with a given hash of the redeem script. /// Generates P2WSH-type of scriptPubkey with a given [`WitnessProgram`].
pub fn new_witness_program(version: WitnessVersion, program: &[u8]) -> Self { pub fn new_witness_program(witness_program: &WitnessProgram) -> Self {
Builder::new()
.push_opcode(witness_program.version().into())
.push_slice(witness_program.program())
.into_script()
}
/// Generates P2WSH-type of scriptPubkey with a given [`WitnessVersion`] and the program bytes.
/// Does not do any checks on version or program length.
///
/// Convenience method used by `new_v0_p2wpkh`, `new_v0_p2wsh`, `new_v1_p2tr`, and
/// `new_v1_p2tr_tweaked`.
fn new_witness_program_unchecked(version: WitnessVersion, program: &[u8]) -> Self {
debug_assert!(program.len() >= 2 && program.len() <= 40);
// In segwit v0, the program must be 20 or 32 bytes long.
debug_assert!(version != WitnessVersion::V0 || program.len() == 20 || program.len() == 32);
Builder::new() Builder::new()
.push_opcode(version.into()) .push_opcode(version.into())
.push_slice(program) .push_slice(program)

View File

@ -1742,6 +1742,7 @@ mod tests {
#[cfg(feature = "rand-std")] #[cfg(feature = "rand-std")]
fn sign_psbt() { fn sign_psbt() {
use crate::WPubkeyHash; use crate::WPubkeyHash;
use crate::address::WitnessProgram;
use crate::bip32::{Fingerprint, DerivationPath}; use crate::bip32::{Fingerprint, DerivationPath};
let unsigned_tx = Transaction { let unsigned_tx = Transaction {
@ -1771,9 +1772,12 @@ mod tests {
psbt.inputs[0].bip32_derivation = map; psbt.inputs[0].bip32_derivation = map;
// Second input is unspendable by us e.g., from another wallet that supports future upgrades. // Second input is unspendable by us e.g., from another wallet that supports future upgrades.
let unknown_prog = WitnessProgram::new(
crate::address::WitnessVersion::V4, vec![0xaa; 34]
).unwrap();
let txout_unknown_future = TxOut{ let txout_unknown_future = TxOut{
value: 10, value: 10,
script_pubkey: ScriptBuf::new_witness_program(crate::address::WitnessVersion::V4, &[0xaa; 34]), script_pubkey: ScriptBuf::new_witness_program(&unknown_prog),
}; };
psbt.inputs[1].witness_utxo = Some(txout_unknown_future); psbt.inputs[1].witness_utxo = Some(txout_unknown_future);