diff --git a/src/rt/io.rs b/src/rt/io.rs index fbe9ecc901..2adaeda8f1 100644 --- a/src/rt/io.rs +++ b/src/rt/io.rs @@ -244,6 +244,52 @@ impl<'data> ReadBufCursor<'data> { self.buf.init = self.buf.filled.max(self.buf.init); } + /// Read into a byte slice safely. + /// + /// This function provides a safe alternative to the [`as_mut`] and + /// [`advance`] APIS. It initializes the uninitialized part of the + /// buffer (if one exists) and passed it as a slice to the provided + /// closure. Then, it advances the filled part of the buffer by the + /// returned value. + #[inline] + pub fn read_with<E>(&mut self, f: impl FnOnce(&mut [u8]) -> Result<usize, E>) -> Result<(), E> { + // Initialize the unfilled portion of this cursor. + if self.buf.init < self.buf.raw.len() { + self.buf.raw[self.buf.init..].fill(MaybeUninit::new(0)); + self.buf.init = self.buf.raw.len(); + } + + // SAFETY: The full buffer is initialized now. + let slice = unsafe { + core::slice::from_raw_parts_mut( + self.buf.raw[self.buf.filled..].as_mut_ptr() as *mut u8, + self.buf.raw.len() - self.buf.filled, + ) + }; + + // Call the function. + let n = match f(slice) { + Ok(n) => n, + Err(e) => return Err(e), + }; + + // Panic if `n` is greater than the capacity we actually have. + if self + .buf + .filled + .checked_add(n) + .map(|total| total > self.buf.raw.len()) + != Some(true) + { + panic!("closure returned more bytes than are available in the buffer"); + } + + // Advance the buffer. + self.buf.filled = self.buf.filled.checked_add(n).expect("overflow"); + + Ok(()) + } + #[inline] pub(crate) fn remaining(&self) -> usize { self.buf.remaining() diff --git a/tests/rt.rs b/tests/rt.rs new file mode 100644 index 0000000000..c97066ecbf --- /dev/null +++ b/tests/rt.rs @@ -0,0 +1,21 @@ +use hyper::rt::ReadBuf; +use std::io::Read; + +const TEXT: &[u8] = r#"Philosophers often behave like little children who scribble some marks on a piece of paper at random and then ask the grown-up "What's that?" — It happened like this: the grown-up had drawn pictures for the child several times and said "this is a man," "this is a house," etc. And then the child makes some marks too and asks: what's this then?"#.as_bytes(); + +#[test] +fn with_read() { + // Create the buffer. + let mut buffer = [0u8; TEXT.len() + 27]; + let mut read_buf = ReadBuf::new(&mut buffer); + let mut cursor = read_buf.unfilled(); + + // Read into it. + cursor.read_with(|out| { + let mut text = TEXT; + text.read(out) + }).unwrap(); + + assert_eq!(read_buf.filled(), TEXT); + assert_eq!(&buffer[..TEXT.len()], TEXT); +}