Merge rust-bitcoin/rust-bitcoin#2339: io: Add read_to_end trait method to Read

3333dbab24 Use new read_to_limit function (Tobin C. Harding)
f29da57ef6 io: Add functions to read to the end of a reader (Tobin C. Harding)

Pull request description:

  Total re-write, now does:

  - Add a method to the `io::Take` trait `read_to_end`
  - Add a method to the `io::Read` trait `read_to_limit` (with default impl that calls through to `Take::read_to_end`)
  - Use the new methods and remove `psbt::read_to_end`

ACKs for top commit:
  apoelstra:
    ACK 3333dbab24

Tree-SHA512: ce4439a85296db36013e284ebce6f91665b1daa21e6d5a570746a3dea789aec3e6a32d295ab88a8c72b56ee8d50b18e142cbd684980076f2e88c3ae7acf5c322

bitcoin/src/psbt/raw.rs

@@ -8,7 +8,7 @@
 
 use core::fmt;
 
-use io::{BufRead, Read, Write};
+use io::{BufRead, Write};
 
 use super::serialize::{Deserialize, Serialize};
 use crate::consensus::encode::{
@@ -162,7 +162,11 @@ where
     fn consensus_decode<R: BufRead + ?Sized>(r: &mut R) -> Result<Self, encode::Error> {
         let prefix = Vec::<u8>::consensus_decode(r)?;
         let subtype = Subtype::from(r.read_u8()?);
-        let key = read_to_end(r)?;
+
+        // The limit is a DOS protection mechanism the exact value is not
+        // important, 1024 bytes is bigger than any key should be.
+        let mut key = vec![];
+        let _ = r.read_to_limit(&mut key, 1024)?;
 
         Ok(ProprietaryKey { prefix, subtype, key })
     }
@@ -194,17 +198,3 @@ where
         Ok(deserialize(&key.key)?)
     }
 }
-
-pub(crate) fn read_to_end<D: Read + ?Sized>(d: &mut D) -> Result<Vec<u8>, io::Error> {
-    let mut result = vec![];
-    let mut buf = [0u8; 64];
-    loop {
-        match d.read(&mut buf) {
-            Ok(0) => break,
-            Ok(n) => result.extend_from_slice(&buf[0..n]),
-            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
-            Err(e) => return Err(e),
-        };
-    }
-    Ok(result)
-}

io/src/lib.rs

@@ -18,6 +18,8 @@ extern crate alloc;
 mod error;
 mod macros;
 
+#[cfg(feature = "alloc")]
+use alloc::vec::Vec;
 use core::cmp;
 use core::convert::TryInto;
 
@@ -45,6 +47,19 @@ pub trait Read {
 
     #[inline]
     fn take(&mut self, limit: u64) -> Take<Self> { Take { reader: self, remaining: limit } }
+
+    /// Attempts to read up to limit bytes from the reader, allocating space in `buf` as needed.
+    ///
+    /// `limit` is used to prevent a denial of service attack vector since an unbounded reader will
+    /// exhaust all memory.
+    ///
+    /// Similar to `std::io::Read::read_to_end` but with the DOS protection.
+    #[doc(alias = "read_to_end")]
+    #[cfg(any(feature = "alloc", feature = "std"))]
+    #[inline]
+    fn read_to_limit(&mut self, buf: &mut Vec<u8>, limit: u64) -> Result<usize> {
+        self.take(limit).read_to_end(buf)
+    }
 }
 
 /// A trait describing an input stream that uses an internal buffer when reading.
@@ -65,6 +80,27 @@ pub struct Take<'a, R: Read + ?Sized> {
     remaining: u64,
 }
 
+impl<'a, R: Read + ?Sized> Take<'a, R> {
+    #[cfg(any(feature = "alloc", feature = "std"))]
+    #[inline]
+    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
+        let mut read: usize = 0;
+        let mut chunk = [0u8; 64];
+        loop {
+            match self.read(&mut chunk) {
+                Ok(0) => break,
+                Ok(n) => {
+                    buf.extend_from_slice(&chunk[0..n]);
+                    read += n;
+                }
+                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
+                Err(e) => return Err(e),
+            };
+        }
+        Ok(read)
+    }
+}
+
 impl<'a, R: Read + ?Sized> Read for Take<'a, R> {
     #[inline]
     fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
@@ -293,4 +329,29 @@ mod tests {
         assert_eq!(fill.len(), 0);
         assert_eq!(fill, &[]);
     }
+
+    #[test]
+    #[cfg(any(feature = "alloc", feature = "std"))]
+    fn read_to_limit_greater_than_total_length() {
+        let s = "16-byte-string!!".to_string();
+        let mut reader = Cursor::new(&s);
+        let mut buf = vec![];
+
+        // 32 is greater than the reader length.
+        let read = reader.read_to_limit(&mut buf, 32).expect("failed to read to limit");
+        assert_eq!(read, s.len());
+        assert_eq!(&buf, s.as_bytes())
+    }
+
+    #[test]
+    #[cfg(any(feature = "alloc", feature = "std"))]
+    fn read_to_limit_less_than_total_length() {
+        let s = "16-byte-string!!".to_string();
+        let mut reader = Cursor::new(&s);
+        let mut buf = vec![];
+
+        let read = reader.read_to_limit(&mut buf, 2).expect("failed to read to limit");
+        assert_eq!(read, 2);
+        assert_eq!(&buf, "16".as_bytes())
+    }
 }