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.

bitcoin/src/address.rs

@@ -392,12 +392,41 @@ pub enum Payload {
     /// P2SH address.
     ScriptHash(ScriptHash),
     /// 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 {
@@ -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 {
+            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 {
             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 } =>
+            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
     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
@@ -471,10 +496,11 @@ impl Payload {
 
     /// Create a witness pay to script hash 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
@@ -492,20 +518,22 @@ impl Payload {
         merkle_root: Option<TapNodeHash>,
     ) -> Payload {
         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.
     ///
     /// 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 {
+        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
@@ -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 {
+                match prog.version() {
-                    WitnessVersion::V0 => match prog.len() {
+                    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
+    /// *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
     /// 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 {
+            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
             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 {
+            .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<_>>();
 
         const LEGACY_EQUIVALENCE_CLASSES: &[&[Network]] =

bitcoin/src/blockdata/script.rs

@@ -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.
+    /// 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()
             .push_opcode(version.into())
             .push_slice(program)

bitcoin/src/psbt/mod.rs

@@ -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);