milksad
/
rust-bitcoin-unsafe-fast
14
0
Fork 0
Code Issues Pull Requests Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

141
bitcoin/src/address.rs
View File

@ -392,12 +392,41 @@ pub enum Payload {
/// P2SH address.
ScriptHash(ScriptHash),
/// Segwit address.
WitnessProgram {
WitnessProgram(WitnessProgram),
}
/// Witness program as defined in BIP141.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WitnessProgram {
/// The witness program version.
version: WitnessVersion,
/// The witness program.
/// 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 {
@ -412,18 +441,13 @@ impl Payload {
hash_inner.copy_from_slice(&script.as_bytes()[2..22]);
Payload::ScriptHash(ScriptHash::from_inner(hash_inner))
} 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");
Payload::WitnessProgram {
version: WitnessVersion::try_from(opcode)?,
program: script.as_bytes()[2..].to_vec(),
}
let witness_program = WitnessProgram::new(
WitnessVersion::try_from(opcode)?,
script.as_bytes()[2..].to_vec(),
)?;
Payload::WitnessProgram(witness_program)
} else {
return Err(Error::UnrecognizedScript);
})
@ -434,8 +458,8 @@ impl Payload {
match *self {
Payload::PubkeyHash(ref hash) => script::ScriptBuf::new_p2pkh(hash),
Payload::ScriptHash(ref hash) => script::ScriptBuf::new_p2sh(hash),
Payload::WitnessProgram { version, program: ref prog } =>
script::ScriptBuf::new_witness_program(version, prog)
Payload::WitnessProgram(ref 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
pub fn p2wpkh(pk: &PublicKey) -> Result<Payload, Error> {
Ok(Payload::WitnessProgram {
version: WitnessVersion::V0,
program: pk.wpubkey_hash().ok_or(Error::UncompressedPubkey)?.as_ref().to_vec(),
})
let prog = WitnessProgram::new(
WitnessVersion::V0,
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
@ -471,10 +496,11 @@ impl Payload {
/// Create a witness pay to script hash payload.
pub fn p2wsh(script: &script::Script) -> Payload {
Payload::WitnessProgram {
version: WitnessVersion::V0,
program: script.wscript_hash().as_ref().to_vec(),
}
let prog = WitnessProgram::new(
WitnessVersion::V0,
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
@ -492,20 +518,22 @@ impl Payload {
merkle_root: Option<TapNodeHash>,
) -> Payload {
let (output_key, _parity) = internal_key.tap_tweak(secp, merkle_root);
Payload::WitnessProgram {
version: WitnessVersion::V1,
program: output_key.to_inner().serialize().to_vec(),
}
let prog = WitnessProgram::new(
WitnessVersion::V1,
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.
///
/// This method is not recommended for use and [Payload::p2tr()] should be used where possible.
pub fn p2tr_tweaked(output_key: TweakedPublicKey) -> Payload {
Payload::WitnessProgram {
version: WitnessVersion::V1,
program: output_key.to_inner().serialize().to_vec(),
}
let prog = WitnessProgram::new(
WitnessVersion::V1,
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
@ -513,7 +541,7 @@ impl Payload {
match self {
Payload::ScriptHash(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[..]);
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 writer = if fmt.alternate() {
upper_writer = UpperWriter(fmt);
@ -556,7 +585,7 @@ impl<'a> fmt::Display for AddressEncoding<'a> {
};
let mut bech32_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)
}
}
@ -725,15 +754,15 @@ impl<V: NetworkValidation> Address<V> {
match self.payload {
Payload::PubkeyHash(_) => Some(AddressType::P2pkh),
Payload::ScriptHash(_) => Some(AddressType::P2sh),
Payload::WitnessProgram { version, program: ref prog } => {
Payload::WitnessProgram(ref prog) => {
// BIP-141 p2wpkh or p2wsh addresses.
match version {
WitnessVersion::V0 => match prog.len() {
match prog.version() {
WitnessVersion::V0 => match prog.program().len() {
20 => Some(AddressType::P2wpkh),
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,
}
}
@ -849,12 +878,18 @@ impl Address {
}
/// 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
/// compatibility, the senders must send to any [`Address`]. Receivers can use this method to
/// check whether or not they can spend from this address.
/// *spending* from this address. *NOT* to be called by senders.
///
/// <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
/// considered non-standard.
/// </details>
///
pub fn is_spend_standard(&self) -> bool { self.address_type().is_some() }
/// 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)?)
};
if program.len() < 2 || program.len() > 40 {
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()));
}
let witness_program = WitnessProgram::new(version, program)?;
// Encoding check
let expected = version.bech32_variant();
@ -1056,7 +1084,7 @@ impl FromStr for Address<NetworkUnchecked> {
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
@ -1238,7 +1266,8 @@ mod tests {
let program = hex!(
"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);
}
@ -1449,10 +1478,10 @@ mod tests {
Payload::ScriptHash(ScriptHash::all_zeros()),
];
let segwit_payload = (0..=16)
.map(|version| Payload::WitnessProgram {
version: WitnessVersion::try_from(version).unwrap(),
program: vec![],
})
.map(|version| Payload::WitnessProgram(WitnessProgram::new(
WitnessVersion::try_from(version).unwrap(),
vec![0xab; 32], // Choose 32 to make test case valid for all witness versions(including v0)
).unwrap()))
.collect::<Vec<_>>();
const LEGACY_EQUIVALENCE_CLASSES: &[&[Network]] =

33
bitcoin/src/blockdata/script.rs
View File

@ -74,7 +74,7 @@ use crate::policy::DUST_RELAY_TX_FEE;
use crate::OutPoint;
use crate::key::PublicKey;
use crate::address::WitnessVersion;
use crate::address::{WitnessVersion, WitnessProgram};
use crate::taproot::{LeafVersion, TapNodeHash, TapLeafHash};
use secp256k1::{Secp256k1, Verification, XOnlyPublicKey};
use crate::schnorr::{TapTweak, TweakedPublicKey, UntweakedPublicKey};
@ -1122,28 +1122,47 @@ impl ScriptBuf {
/// Generates P2WPKH-type of scriptPubkey.
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.
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
/// script tree merkle root.
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);
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`].
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.
pub fn new_witness_program(version: WitnessVersion, program: &[u8]) -> Self {
/// Generates P2WSH-type of scriptPubkey with a given [`WitnessProgram`].
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()
.push_opcode(version.into())
.push_slice(program)

6
bitcoin/src/psbt/mod.rs
View File

@ -1742,6 +1742,7 @@ mod tests {
#[cfg(feature = "rand-std")]
fn sign_psbt() {
use crate::WPubkeyHash;
use crate::address::WitnessProgram;
use crate::bip32::{Fingerprint, DerivationPath};
let unsigned_tx = Transaction {
@ -1771,9 +1772,12 @@ mod tests {
psbt.inputs[0].bip32_derivation = map;
// 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{
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);