Merge rust-bitcoin/rust-bitcoin#6920: primitives: Bound aggregate size and weight while decoding

### benches/bitcoin/block.rs
@@ -124,7 +124,9 @@ fn bench_large_block(c: &mut Criterion) {
     let mut g = c.benchmark_group("large_block");
     g.measurement_time(Duration::from_secs(15)).warm_up_time(Duration::from_secs(3));
 
-    for num_tx in [1000, 10000, 64000] {
+    // The largest count is kept below the number of minimal transactions that fit within the
+    // maximum block weight, otherwise the decoder rejects the block as too heavy.
+    for num_tx in [1000, 10000, 16000] {
         let raw_block = build_test_block(num_tx);
 
         g.bench_function(BenchmarkId::new("decode", format!("{}tx", num_tx)), |b| {

### primitives/src/block.rs
@@ -595,12 +595,33 @@ crate::decoder_newtype! {
     /// Constructs a new [`Block`] decoder.
     pub const fn new() -> Self { Self(Decoder2::new(HeaderDecoder::new(), VecDecoder::new())) }
 
+    fn map_push_bytes_err(err: <BlockInnerDecoder as encoding::Decoder>::Error) -> BlockDecoderError {
+        BlockDecoderError(error::BlockDecoderErrorInner::Decode(err))
+    }
+
     fn end(result: Result<(Header, Vec<Transaction>), <BlockInnerDecoder as encoding::Decoder>::Error>) -> Result<Block, BlockDecoderError> {
-        let (header, transactions) = result.map_err(BlockDecoderError)?;
+        let (header, transactions) =
+            result.map_err(|e| BlockDecoderError(error::BlockDecoderErrorInner::Decode(e)))?;
+        let weight = block_weight_wu(&transactions);
+        if weight > crate::transaction::MAX_BLOCK_WEIGHT.to_wu() {
+            return Err(BlockDecoderError(error::BlockDecoderErrorInner::BlockTooHeavy(weight)));
+        }
         Ok(Self::Output::new_unchecked(header, transactions))
     }
 }
 
+/// Computes the consensus weight, in weight units, of a decoded block.
+#[cfg(feature = "alloc")]
+fn block_weight_wu(transactions: &[Transaction]) -> u64 {
+    // The 80-byte header and the transaction count prefix are non-witness data.
+    let header_and_count = 80 + crate::compact_size_encode(transactions.len()).as_slice().len();
+    let mut weight = (header_and_count as u64) * 4;
+    for tx in transactions {
+        weight = weight.saturating_add(crate::transaction::transaction_weight(tx).to_wu());
+    }
+    weight
+}
+
 /// Computes the Merkle root for a list of transactions.
 ///
 /// Returns [`None`] if the iterator was empty, or if the transaction list contains
@@ -1023,7 +1044,16 @@ pub mod error {
     /// [`Block`]: super::Block
     #[cfg(feature = "alloc")]
     #[derive(Debug, Clone, PartialEq, Eq)]
-    pub struct BlockDecoderError(pub(super) <super::BlockInnerDecoder as encoding::Decoder>::Error);
+    pub struct BlockDecoderError(pub(super) BlockDecoderErrorInner);
+
+    #[cfg(feature = "alloc")]
+    #[derive(Debug, Clone, PartialEq, Eq)]
+    pub(super) enum BlockDecoderErrorInner {
+        /// Error decoding the block header or transactions.
+        Decode(<super::BlockInnerDecoder as encoding::Decoder>::Error),
+        /// Block weight exceeds the maximum block weight.
+        BlockTooHeavy(u64),
+    }
 
     #[cfg(feature = "alloc")]
     impl From<Infallible> for BlockDecoderError {
@@ -1035,15 +1065,28 @@ pub mod error {
     impl fmt::Display for BlockDecoderError {
         #[inline]
         fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-            write_err!(f, "block decoder error"; self.0)
+            match self.0 {
+                BlockDecoderErrorInner::Decode(ref e) => write_err!(f, "block decoder error"; e),
+                BlockDecoderErrorInner::BlockTooHeavy(weight) => write!(
+                    f,
+                    "block weight {} exceeds the maximum of {}",
+                    weight,
+                    crate::transaction::MAX_BLOCK_WEIGHT.to_wu()
+                ),
+            }
         }
     }
 
     #[cfg(feature = "alloc")]
     #[cfg(feature = "std")]
     impl std::error::Error for BlockDecoderError {
         #[inline]
-        fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { Some(&self.0) }
+        fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+            match self.0 {
+                BlockDecoderErrorInner::Decode(ref e) => Some(e),
+                BlockDecoderErrorInner::BlockTooHeavy(_) => None,
+            }
+        }
     }
 
     /// Invalid block error.
@@ -2221,17 +2264,17 @@ mod tests {
     #[cfg(feature = "alloc")]
     fn block_decoder_error() {
         fn is_first(err: &BlockDecoderError) -> bool {
-            match err.0 {
-                encoding::Decoder2Error::First(_) => true,
-                encoding::Decoder2Error::Second(_) => false,
-            }
+            matches!(
+                err.0,
+                error::BlockDecoderErrorInner::Decode(encoding::Decoder2Error::First(_))
+            )
         }
 
         fn is_second(err: &BlockDecoderError) -> bool {
-            match err.0 {
-                encoding::Decoder2Error::First(_) => false,
-                encoding::Decoder2Error::Second(_) => true,
-            }
+            matches!(
+                err.0,
+                error::BlockDecoderErrorInner::Decode(encoding::Decoder2Error::Second(_))
+            )
         }
 
         let err_first = Block::decoder().end().unwrap_err();
@@ -2259,6 +2302,69 @@ mod tests {
         assert!(std::error::Error::source(&err_second).is_some());
     }
 
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn block_decoder_rejects_aggregate_weight_over_limit() {
+        fn large_transaction(tag: u8) -> Transaction {
+            Transaction {
+                version: crate::transaction::Version::ONE,
+                lock_time: crate::absolute::LockTime::ZERO,
+                inputs: vec![crate::TxIn {
+                    previous_output: crate::OutPoint {
+                        txid: crate::Txid::from_byte_array([tag; 32]),
+                        vout: 0,
+                    },
+                    script_sig: crate::ScriptSigBuf::from_bytes(vec![tag; 550_000]),
+                    sequence: crate::Sequence::MAX,
+                    witness: crate::Witness::new(),
+                }],
+                outputs: vec![crate::TxOut {
+                    amount: units::Amount::ZERO,
+                    script_pubkey: crate::ScriptPubKeyBuf::new(),
+                }],
+            }
+        }
+
+        // Each transaction is individually valid but together they exceed a block's weight.
+        let transactions = vec![large_transaction(1), large_transaction(2)];
+        let block = Block::new_unchecked(dummy_header(), transactions);
+        let encoded = encoding::encode_to_vec(&block);
+
+        let err = encoding::decode_from_slice::<Block>(&encoded).unwrap_err();
+        assert!(matches!(
+            err,
+            encoding::DecodeError::Parse(BlockDecoderError(
+                error::BlockDecoderErrorInner::BlockTooHeavy(_)
+            ))
+        ));
+    }
+
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn block_weight_wu_matches_expected() {
+        let tx = Transaction {
+            version: crate::transaction::Version::ONE,
+            lock_time: crate::absolute::LockTime::ZERO,
+            inputs: vec![crate::TxIn {
+                previous_output: crate::OutPoint {
+                    txid: crate::Txid::from_byte_array([1u8; 32]),
+                    vout: 0,
+                },
+                script_sig: crate::ScriptSigBuf::new(),
+                sequence: crate::Sequence::MAX,
+                witness: crate::Witness::new(),
+            }],
+            outputs: vec![crate::TxOut {
+                amount: units::Amount::ZERO,
+                script_pubkey: crate::ScriptPubKeyBuf::new(),
+            }],
+        };
+
+        // header_and_count = 80 + 1 (compact size of a single transaction) = 81
+        // weight = 81 * 4 + transaction_weight(tx) = 324 + 240 = 564
+        assert_eq!(block_weight_wu(&[tx]), 564);
+    }
+
     #[test]
     #[cfg(feature = "alloc")]
     fn header_decoder_error() {

### primitives/src/transaction.rs
@@ -367,6 +367,44 @@ const SEGWIT_MARKER: u8 = 0x00;
 #[cfg(feature = "alloc")]
 const SEGWIT_FLAG: u8 = 0x01;
 
+/// Maximum block weight, and therefore the maximum weight of any single transaction. (BIP-0141)
+#[cfg(feature = "alloc")]
+pub(crate) const MAX_BLOCK_WEIGHT: Weight = Weight::from_wu(4_000_000);
+
+/// Computes the consensus weight of a transaction.
+///
+/// Weight is `4 * base_size + witness_overhead`, where the base size excludes the segwit marker,
+/// flag and witness data, all of which contribute a weight of one.
+#[cfg(feature = "alloc")]
+pub(crate) fn transaction_weight(tx: &Transaction) -> Weight {
+    let compact = |n: usize| crate::compact_size_encode(n).as_slice().len();
+
+    let mut base = 4; // version
+    base += compact(tx.inputs.len());
+    for input in &tx.inputs {
+        let script = input.script_sig.len();
+        base += 36 + compact(script) + script + 4; // outpoint + script + sequence
+    }
+    base += compact(tx.outputs.len());
+    for output in &tx.outputs {
+        let script = output.script_pubkey.len();
+        base += 8 + compact(script) + script; // amount + script
+    }
+    base += 4; // lock_time
+
+    let witness = if tx.uses_segwit_serialization() {
+        let mut w = 2; // segwit marker and flag
+        for input in &tx.inputs {
+            w += input.witness.size();
+        }
+        w
+    } else {
+        0
+    };
+
+    Weight::from_wu((base as u64) * 4 + witness as u64)
+}
+
 // This is equivalent to consensus encoding but hashes the fields manually.
 #[cfg(feature = "alloc")]
 fn hash_transaction(tx: &Transaction, uses_segwit_serialization: bool) -> sha256d::Hash {
@@ -646,6 +684,12 @@ impl encoding::Decoder for TransactionDecoder {
                 if tx.outputs.is_empty() {
                     return Err(E(Inner::NoOutputs));
                 }
+                // A transaction heavier than a whole block can never be valid; reject it before
+                // the more expensive checks below.
+                let weight = transaction_weight(&tx);
+                if weight > MAX_BLOCK_WEIGHT {
+                    return Err(E(Inner::TransactionTooHeavy(weight.to_wu())));
+                }
                 // check for null prevout in non-coinbase txs
                 if tx.inputs.len() > 1 {
                     for (index, input) in tx.inputs.iter().enumerate() {
@@ -1357,6 +1401,8 @@ pub mod error {
         OutputValueSumTooLarge(u64),
         /// Transaction has no outputs.
         NoOutputs,
+        /// Transaction weight exceeds the maximum block weight.
+        TransactionTooHeavy(u64),
     }
 
     #[cfg(feature = "alloc")]
@@ -1392,6 +1438,12 @@ pub mod error {
                 E::OutputValueSumTooLarge(val) =>
                     write!(f, "sum of output values {} satoshis exceeds MAX_MONEY", val),
                 E::NoOutputs => write!(f, "transaction has no outputs"),
+                E::TransactionTooHeavy(wu) => write!(
+                    f,
+                    "transaction weight {} exceeds the maximum block weight of {}",
+                    wu,
+                    super::MAX_BLOCK_WEIGHT.to_wu()
+                ),
             }
         }
     }
@@ -1418,6 +1470,7 @@ pub mod error {
                 E::DuplicateInput(_) => None,
                 E::OutputValueSumTooLarge(_) => None,
                 E::NoOutputs => None,
+                E::TransactionTooHeavy(_) => None,
             }
         }
     }
@@ -2304,6 +2357,89 @@ mod tests {
         assert_eq!(err, TransactionDecoderError(TransactionDecoderErrorInner::NoOutputs));
     }
 
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn reject_transaction_exceeding_maximum_block_weight() {
+        let script_len = 600_000usize;
+        let mut tx_bytes = Vec::new();
+
+        tx_bytes.extend_from_slice(&Version::ONE.to_u32().to_le_bytes());
+        tx_bytes.push(1); // input count
+        tx_bytes.extend_from_slice(&[0x01; 32]); // non-null prevout txid
+        tx_bytes.extend_from_slice(&0u32.to_le_bytes()); // prevout vout
+        tx_bytes.push(0); // empty script_sig
+        tx_bytes.extend_from_slice(&Sequence::MAX.to_consensus_u32().to_le_bytes());
+        tx_bytes.push(2); // output count
+
+        for _ in 0..2 {
+            tx_bytes.extend_from_slice(&0u64.to_le_bytes()); // amount
+            tx_bytes.extend_from_slice(crate::compact_size_encode(script_len).as_slice());
+            tx_bytes.resize(tx_bytes.len() + script_len, 0);
+        }
+        tx_bytes.extend_from_slice(&0u32.to_le_bytes()); // lock_time
+
+        // With no witness discount this transaction weighs more than 4,000,000 WU,
+        // so it cannot fit in a consensus-valid block.
+        let err = encoding::decode_from_slice::<Transaction>(&tx_bytes).unwrap_err();
+        assert!(matches!(
+            err,
+            encoding::DecodeError::Parse(TransactionDecoderError(
+                TransactionDecoderErrorInner::TransactionTooHeavy(_)
+            ))
+        ));
+    }
+
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn transaction_weight_matches_expected() {
+        // A segwit transaction with a non-empty script_sig, a non-empty script_pubkey and
+        // witness data, so that every term in `transaction_weight` contributes to the result.
+        let tx = Transaction {
+            version: Version::ONE,
+            lock_time: absolute::LockTime::ZERO,
+            inputs: vec![TxIn {
+                previous_output: OutPoint { txid: Txid::from_byte_array([1u8; 32]), vout: 0 },
+                script_sig: ScriptSigBuf::from_bytes(vec![0u8; 5]),
+                sequence: Sequence::MAX,
+                witness: Witness::from_slice(&[[0xab; 72].as_slice(), [0xcd; 33].as_slice()]),
+            }],
+            outputs: vec![TxOut {
+                amount: Amount::ONE_SAT,
+                script_pubkey: ScriptPubKeyBuf::from_bytes(vec![0u8; 7]),
+            }],
+        };
+
+        // base    = 4 (version) + 1 (input count) + (36 + 1 + 5 + 4) (input)
+        //           + 1 (output count) + (8 + 1 + 7) (output) + 4 (locktime) = 72
+        // witness = 2 (segwit marker and flag) + 108 (witness.size()) = 110
+        // weight  = 72 * 4 + 110 = 398
+        assert_eq!(transaction_weight(&tx).to_wu(), 398);
+    }
+
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn accept_transaction_at_maximum_block_weight() {
+        // A base size of 1,000,000 bytes gives a weight of exactly MAX_BLOCK_WEIGHT, which the
+        // `weight > MAX_BLOCK_WEIGHT` boundary in `end` must accept rather than reject.
+        let script_len = 999_936usize;
+        let mut tx_bytes = Vec::new();
+
+        tx_bytes.extend_from_slice(&Version::ONE.to_u32().to_le_bytes());
+        tx_bytes.push(1); // input count
+        tx_bytes.extend_from_slice(&[0x01; 32]); // non-null prevout txid
+        tx_bytes.extend_from_slice(&0u32.to_le_bytes()); // prevout vout
+        tx_bytes.push(0); // empty script_sig
+        tx_bytes.extend_from_slice(&Sequence::MAX.to_consensus_u32().to_le_bytes());
+        tx_bytes.push(1); // output count
+        tx_bytes.extend_from_slice(&0u64.to_le_bytes()); // amount
+        tx_bytes.extend_from_slice(crate::compact_size_encode(script_len).as_slice());
+        tx_bytes.resize(tx_bytes.len() + script_len, 0);
+        tx_bytes.extend_from_slice(&0u32.to_le_bytes()); // lock_time
+
+        let tx = encoding::decode_from_slice::<Transaction>(&tx_bytes).unwrap();
+        assert_eq!(transaction_weight(&tx).to_wu(), MAX_BLOCK_WEIGHT.to_wu());
+    }
+
     #[test]
     #[cfg(feature = "alloc")]
     fn compute_ntxid_ignores_script_sig_and_witness() {

### primitives/src/witness.rs
@@ -83,6 +83,12 @@ const MAX_WITNESS_STACK_ITEMS: usize = 4_000_000;
 /// single witness item cannot exceed what fits in a block.
 const MAX_WITNESS_ITEM_SIZE: usize = 4_000_000;
 
+/// Maximum total serialized byte size of a witness stack, including its count and length prefixes.
+///
+/// This is an anti-DoS limit based on Bitcoin's 4MB block weight limit.
+/// The serialized witness stack is part of a block, so its size cannot exceed what fits in a block.
+const MAX_WITNESS_SIZE: usize = 4_000_000;
+
 /// The Witness is the data used to unlock bitcoin since the [SegWit upgrade].
 ///
 /// Can be logically seen as an array of bytestrings, i.e. [`Vec<Vec<u8>>`], and it is serialized on the wire
@@ -401,6 +407,8 @@ pub struct WitnessDecoder {
     /// - `None` means we're currently reading the length.
     /// - `Some(n)` means we're reading element data with `n` bytes remaining.
     element_bytes_remaining: Option<usize>,
+    /// Running total of serialized witness bytes decoded so far.
+    serialized_size: usize,
 }
 
 impl WitnessDecoder {
@@ -414,6 +422,7 @@ impl WitnessDecoder {
             element_idx: 0,
             element_length_decoder: CompactSizeDecoder::new_with_limit(MAX_WITNESS_ITEM_SIZE),
             element_bytes_remaining: None,
+            serialized_size: 0,
         }
     }
 }
@@ -446,6 +455,7 @@ impl encoding::Decoder for WitnessDecoder {
             let decoder = core::mem::take(&mut self.witness_count_decoder);
             let witness_elements = decoder.end().map_err(Inner::LengthPrefixDecode).map_err(E)?;
             self.witness_elements = Some(witness_elements);
+            self.serialized_size = crate::compact_size_encode(witness_elements).as_slice().len();
 
             // Short circuit for zero witness elements.
             if witness_elements == 0 {
@@ -515,6 +525,13 @@ impl encoding::Decoder for WitnessDecoder {
 
                 // keep the element length prefix in the content area.
                 let encoded_compact_size = crate::compact_size_encode(element_length);
+                self.serialized_size = self
+                    .serialized_size
+                    .saturating_add(encoded_compact_size.as_slice().len())
+                    .saturating_add(element_length);
+                if self.serialized_size > MAX_WITNESS_SIZE {
+                    return Err(E(Inner::WitnessTooLarge(self.serialized_size)));
+                }
                 self.content.extend_from_slice(encoded_compact_size.as_slice());
 
                 if element_length == 0 {
@@ -1029,6 +1046,8 @@ pub mod error {
         LengthPrefixDecode(CompactSizeDecoderError),
         /// Not enough bytes given to decoder.
         UnexpectedEof(UnexpectedEofError),
+        /// The combined serialized witness elements exceed the maximum size.
+        WitnessTooLarge(usize),
     }
 
     impl From<Infallible> for WitnessDecoderError {
@@ -1044,6 +1063,12 @@ pub mod error {
             match self.0 {
                 E::LengthPrefixDecode(ref e) => write_err!(f, "vec decoder error"; e),
                 E::UnexpectedEof(ref e) => write_err!(f, "decoder error"; e),
+                E::WitnessTooLarge(size) => write!(
+                    f,
+                    "serialized witness size of {} bytes exceeds the maximum of {}",
+                    size,
+                    super::MAX_WITNESS_SIZE
+                ),
             }
         }
     }
@@ -1057,6 +1082,7 @@ pub mod error {
             match self.0 {
                 E::LengthPrefixDecode(ref e) => Some(e),
                 E::UnexpectedEof(ref e) => Some(e),
+                E::WitnessTooLarge(_) => None,
             }
         }
     }
@@ -1621,17 +1647,20 @@ mod test {
 
     #[test]
     #[cfg(feature = "alloc")]
-    fn decode_max_length() {
+    fn decode_max_serialized_size() {
+        let element_len = MAX_WITNESS_SIZE
+            - crate::compact_size_encode(1usize).as_slice().len()
+            - crate::compact_size_encode(MAX_WITNESS_ITEM_SIZE).as_slice().len();
         let mut encoded = Vec::new();
         encoded.extend_from_slice(crate::compact_size_encode(1usize).as_slice());
-        encoded.extend_from_slice(crate::compact_size_encode(4_000_000usize).as_slice());
-        encoded.resize(encoded.len() + 4_000_000, 0u8);
+        encoded.extend_from_slice(crate::compact_size_encode(element_len).as_slice());
+        encoded.resize(encoded.len() + element_len, 0u8);
 
         let mut slice = encoded.as_slice();
         let mut decoder = WitnessDecoder::new();
         decoder.push_bytes(&mut slice).unwrap();
         let witness = decoder.end().unwrap();
-        assert_eq!(witness[0].len(), 4_000_000);
+        assert_eq!(witness[0].len(), element_len);
     }
 
     #[test]
@@ -1994,9 +2023,12 @@ mod test {
     #[cfg(feature = "alloc")]
     #[test]
     fn dos_protection() {
+        let element_len = MAX_WITNESS_SIZE
+            - crate::compact_size_encode(MAX_WITNESS_STACK_ITEMS).as_slice().len()
+            - crate::compact_size_encode(MAX_WITNESS_ITEM_SIZE).as_slice().len();
         let mut encoded = Vec::new();
         encoded.extend_from_slice(&[0xFE, 0x00, 0x09, 0x3D, 0x00]); // 4_000_000 (witness count)
-        encoded.extend_from_slice(&[0xFE, 0x00, 0x09, 0x3D, 0x00]); // 4_000_000 (1st element length)
+        encoded.extend_from_slice(crate::compact_size_encode(element_len).as_slice());
 
         let mut slice = encoded.as_slice();
         let mut dec = WitnessDecoder::new();
@@ -2030,4 +2062,46 @@ mod test {
         let mut dec = WitnessDecoder::new();
         assert!(dec.push_bytes(&mut encoded.as_slice()).is_err());
     }
+
+    #[cfg(feature = "alloc")]
+    #[test]
+    fn decode_rejects_witness_over_limit_due_to_aggregate_data() {
+        let element_len = MAX_WITNESS_ITEM_SIZE / 2 + 1;
+
+        let mut first = Vec::new();
+        first.extend_from_slice(crate::compact_size_encode(2usize).as_slice());
+        first.extend_from_slice(crate::compact_size_encode(element_len).as_slice());
+        first.resize(first.len() + element_len, 0);
+
+        let mut decoder = WitnessDecoder::new();
+        let mut first_slice = first.as_slice();
+        assert!(decoder.push_bytes(&mut first_slice).unwrap().needs_more());
+        assert!(first_slice.is_empty());
+
+        let mut second = Vec::new();
+        second.extend_from_slice(crate::compact_size_encode(element_len).as_slice());
+        second.resize(second.len() + element_len, 0);
+
+        let mut second_slice = second.as_slice();
+        assert!(
+            decoder.push_bytes(&mut second_slice).is_err(),
+            "a witness whose aggregate data exceeds MAX_WITNESS_SIZE must be rejected"
+        );
+    }
+
+    #[cfg(feature = "alloc")]
+    #[test]
+    fn decode_rejects_witness_over_limit_due_to_prefix_overhead() {
+        let count = MAX_WITNESS_SIZE;
+        let mut encoded = Vec::new();
+        encoded.extend_from_slice(crate::compact_size_encode(count).as_slice());
+        encoded.resize(encoded.len() + count, 0);
+
+        let mut slice = encoded.as_slice();
+        let mut decoder = WitnessDecoder::new();
+        assert!(
+            decoder.push_bytes(&mut slice).is_err(),
+            "empty elements whose length prefixes and count prefix exceed MAX_WITNESS_SIZE must be rejected"
+        );
+    }
 }