From f692c4a9383c318d01cce64f06aab92960a9c19d Mon Sep 17 00:00:00 2001
From: Riccardo Casatta <riccardo@casatta.it>
Date: Tue, 20 Apr 2021 13:59:37 +0200
Subject: [PATCH] Limit bytes read with Take

---
 src/blockdata/transaction.rs |  4 +++-
 src/consensus/encode.rs      | 10 ++++++++++
 src/internal_macros.rs       |  3 ++-
 src/network/message.rs       |  3 +--
 src/util/psbt/map/global.rs  |  5 +++--
 src/util/psbt/mod.rs         |  4 +++-
 6 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/blockdata/transaction.rs b/src/blockdata/transaction.rs
index f45b8908..ba3b5975 100644
--- a/src/blockdata/transaction.rs
+++ b/src/blockdata/transaction.rs
@@ -34,6 +34,7 @@ use blockdata::constants::WITNESS_SCALE_FACTOR;
 #[cfg(feature="bitcoinconsensus")] use blockdata::script;
 use blockdata::script::Script;
 use consensus::{encode, Decodable, Encodable};
+use consensus::encode::MAX_VEC_SIZE;
 use hash_types::{SigHash, Txid, Wtxid};
 use VarInt;
 
@@ -566,7 +567,8 @@ impl Encodable for Transaction {
 }
 
 impl Decodable for Transaction {
-    fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, encode::Error> {
+    fn consensus_decode<D: io::Read>(d: D) -> Result<Self, encode::Error> {
+        let mut d = d.take(MAX_VEC_SIZE as u64);
         let version = i32::consensus_decode(&mut d)?;
         let input = Vec::<TxIn>::consensus_decode(&mut d)?;
         // segwit
diff --git a/src/consensus/encode.rs b/src/consensus/encode.rs
index 54c8366f..6be7d5b6 100644
--- a/src/consensus/encode.rs
+++ b/src/consensus/encode.rs
@@ -575,6 +575,7 @@ macro_rules! impl_vec {
                     return Err(self::Error::OversizedVectorAllocation { requested: byte_size, max: MAX_VEC_SIZE })
                 }
                 let mut ret = Vec::with_capacity(len as usize);
+                let mut d = d.take(MAX_VEC_SIZE as u64);
                 for _ in 0..len {
                     ret.push(Decodable::consensus_decode(&mut d)?);
                 }
@@ -997,6 +998,15 @@ mod tests {
         assert_eq!(cd.ok(), Some(CheckedData(vec![1u8, 2, 3, 4, 5])));
     }
 
+    #[test]
+    fn limit_read_test() {
+        let witness = vec![vec![0u8; 3_999_999]; 2];
+        let ser = serialize(&witness);
+        let mut reader = io::Cursor::new(ser);
+        let err  = Vec::<Vec<u8>>::consensus_decode(&mut reader);
+        assert!(err.is_err());
+    }
+
     #[test]
     fn serialization_round_trips() {
         macro_rules! round_trip {
diff --git a/src/internal_macros.rs b/src/internal_macros.rs
index dc66ec60..6a9a3912 100644
--- a/src/internal_macros.rs
+++ b/src/internal_macros.rs
@@ -33,8 +33,9 @@ macro_rules! impl_consensus_encoding {
         impl $crate::consensus::Decodable for $thing {
             #[inline]
             fn consensus_decode<D: ::std::io::Read>(
-                mut d: D,
+                d: D,
             ) -> Result<$thing, $crate::consensus::encode::Error> {
+                let mut d = d.take($crate::consensus::encode::MAX_VEC_SIZE as u64);
                 Ok($thing {
                     $($field: $crate::consensus::Decodable::consensus_decode(&mut d)?),+
                 })
diff --git a/src/network/message.rs b/src/network/message.rs
index 1ca3a839..a19a9c5a 100644
--- a/src/network/message.rs
+++ b/src/network/message.rs
@@ -29,9 +29,8 @@ use network::address::{Address, AddrV2Message};
 use network::message_network;
 use network::message_blockdata;
 use network::message_filter;
-use consensus::encode::{CheckedData, Decodable, Encodable, VarInt};
+use consensus::encode::{CheckedData, Decodable, Encodable, VarInt, MAX_VEC_SIZE};
 use consensus::{encode, serialize};
-use consensus::encode::MAX_VEC_SIZE;
 
 /// The maximum number of [Inventory] items in an `inv` message.
 ///
diff --git a/src/util/psbt/map/global.rs b/src/util/psbt/map/global.rs
index 7369afa5..6587943e 100644
--- a/src/util/psbt/map/global.rs
+++ b/src/util/psbt/map/global.rs
@@ -19,6 +19,7 @@ use std::cmp;
 
 use blockdata::transaction::Transaction;
 use consensus::{encode, Encodable, Decodable};
+use consensus::encode::MAX_VEC_SIZE;
 use util::psbt::map::Map;
 use util::psbt::raw;
 use util::psbt;
@@ -228,8 +229,8 @@ impl Map for Global {
 impl_psbtmap_consensus_encoding!(Global);
 
 impl Decodable for Global {
-    fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, encode::Error> {
-
+    fn consensus_decode<D: io::Read>(d: D) -> Result<Self, encode::Error> {
+        let mut d = d.take(MAX_VEC_SIZE as u64);
         let mut tx: Option<Transaction> = None;
         let mut version: Option<u32> = None;
         let mut unknowns: BTreeMap<raw::Key, Vec<u8>> = Default::default();
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
index 3cd2ef89..1eb60d48 100644
--- a/src/util/psbt/mod.rs
+++ b/src/util/psbt/mod.rs
@@ -21,6 +21,7 @@
 use blockdata::script::Script;
 use blockdata::transaction::Transaction;
 use consensus::{encode, Encodable, Decodable};
+use consensus::encode::MAX_VEC_SIZE;
 
 use std::io;
 
@@ -162,7 +163,8 @@ impl Encodable for PartiallySignedTransaction {
 }
 
 impl Decodable for PartiallySignedTransaction {
-    fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, encode::Error> {
+    fn consensus_decode<D: io::Read>(d: D) -> Result<Self, encode::Error> {
+        let mut d = d.take(MAX_VEC_SIZE as u64);
         let magic: [u8; 4] = Decodable::consensus_decode(&mut d)?;
 
         if *b"psbt" != magic {