diff options
| -rw-r--r-- | tokio/src/io/read_buf.rs | 8 | ||||
| -rw-r--r-- | tokio/src/io/util/take.rs | 5 |
2 files changed, 9 insertions, 4 deletions
diff --git a/tokio/src/io/read_buf.rs b/tokio/src/io/read_buf.rs index 03b5d05c..8b430ee1 100644 --- a/tokio/src/io/read_buf.rs +++ b/tokio/src/io/read_buf.rs @@ -77,6 +77,14 @@ impl<'a> ReadBuf<'a> { unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) } } + /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`. + #[inline] + pub fn take(&mut self, n: usize) -> ReadBuf<'_> { + let max = std::cmp::min(self.remaining(), n); + // Saftey: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. + unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) } + } + /// Returns a shared reference to the initialized portion of the buffer. /// /// This includes the filled portion. diff --git a/tokio/src/io/util/take.rs b/tokio/src/io/util/take.rs index 2abc7693..4e424f6c 100644 --- a/tokio/src/io/util/take.rs +++ b/tokio/src/io/util/take.rs @@ -85,10 +85,7 @@ impl<R: AsyncRead> AsyncRead for Take<R> { } let me = self.project(); - let max = std::cmp::min(buf.remaining() as u64, *me.limit_) as usize; - // Make a ReadBuf of the unfulled section up to max - // Saftey: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. - let mut b = unsafe { ReadBuf::uninit(&mut buf.unfilled_mut()[..max]) }; + let mut b = buf.take(*me.limit_ as usize); ready!(me.inner.poll_read(cx, &mut b))?; let n = b.filled().len(); |