diff --git a/core/lib/src/data/data_stream.rs b/core/lib/src/data/data_stream.rs index 719e4b3a38..32a0a6c773 100644 --- a/core/lib/src/data/data_stream.rs +++ b/core/lib/src/data/data_stream.rs @@ -16,6 +16,43 @@ pub type InnerStream = Chain>, BodyReader>; /// be used as an opaque [`Read`] structure. pub struct DataStream(crate InnerStream); +impl DataStream { + pub fn read_to_string_with_limit( + &mut self, + buf: &mut String, + limit: usize, + ) -> Result { + self.do_with_limit(limit, |r| r.read_to_string(buf)) + } + + pub fn read_to_end_with_limit( + &mut self, + buf: &mut Vec, + limit: usize, + ) -> Result { + self.do_with_limit(limit, |r| r.read_to_end(buf)) + } + + fn do_with_limit) -> io::Result, T>( + &mut self, + limit: usize, + f: F, + ) -> Result { + let mut r = self.by_ref().take(limit as u64 + 1); + let s = f(&mut r).map_err(LimitReadError::Io)?; + if r.limit() > 0 { + Ok(s) + } else { + Err(LimitReadError::LimitReached) + } + } +} + +pub enum LimitReadError { + LimitReached, + Io(io::Error), +} + // TODO: Have a `BufRead` impl for `DataStream`. At the moment, this isn't // possible since Hyper's `HttpReader` doesn't implement `BufRead`. impl Read for DataStream {