diff options
Diffstat (limited to 'tokio-util/src/lib.rs')
-rw-r--r-- | tokio-util/src/lib.rs | 34 |
1 files changed, 34 insertions, 0 deletions
diff --git a/tokio-util/src/lib.rs b/tokio-util/src/lib.rs index 1e4b9d40..10b828ef 100644 --- a/tokio-util/src/lib.rs +++ b/tokio-util/src/lib.rs @@ -57,3 +57,37 @@ pub mod either; #[cfg(feature = "time")] pub mod time; + +#[cfg(any(feature = "io", feature = "codec"))] +mod util { + use tokio::io::{AsyncRead, ReadBuf}; + + use bytes::BufMut; + use futures_core::ready; + use std::io; + use std::pin::Pin; + use std::task::{Context, Poll}; + + pub(crate) fn poll_read_buf<T: AsyncRead>( + cx: &mut Context<'_>, + io: Pin<&mut T>, + buf: &mut impl BufMut, + ) -> Poll<io::Result<usize>> { + if !buf.has_remaining_mut() { + return Poll::Ready(Ok(0)); + } + + let orig = buf.bytes_mut().as_ptr() as *const u8; + let mut b = ReadBuf::uninit(buf.bytes_mut()); + + ready!(io.poll_read(cx, &mut b))?; + let n = b.filled().len(); + + // Safety: we can assume `n` bytes were read, since they are in`filled`. + assert_eq!(orig, b.filled().as_ptr()); + unsafe { + buf.advance_mut(n); + } + Poll::Ready(Ok(n)) + } +} |