diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index 0772a098a..9b3d854ff 100644 --- a/tokio/src/io/mod.rs +++ b/tokio/src/io/mod.rs @@ -49,7 +49,9 @@ pub use std::io::{Error, ErrorKind, Result}; mod copy; mod read; mod write; +mod read_exact; pub use self::copy::{copy, Copy}; pub use self::read::{read, Read}; pub use self::write::{write, Write}; +pub use self::read_exact::{read_exact, ReadExact}; diff --git a/tokio/src/io/read_exact.rs b/tokio/src/io/read_exact.rs new file mode 100644 index 000000000..e8a66a47d --- /dev/null +++ b/tokio/src/io/read_exact.rs @@ -0,0 +1,73 @@ +use std::future::Future; +use std::io; +use std::marker::Unpin; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio_io::AsyncRead; + +macro_rules! ready { + ($e:expr) => { + match $e { + ::std::task::Poll::Ready(t) => t, + ::std::task::Poll::Pending => return ::std::task::Poll::Pending, + } + }; +} + + +/// A future which can be used to easily read exactly enough bytes to fill +/// a buffer. +/// +/// Created by the [`read_exact`] function. +/// +/// [`read_exact`]: fn.read_exact.html +pub fn read_exact<'a, A>(reader: &'a mut A, buf: &'a mut[u8]) -> ReadExact<'a, A> +where + A: AsyncRead + Unpin + ?Sized +{ + ReadExact { reader, buf, pos: 0 } +} + +/// Creates a future which will read exactly enough bytes to fill `buf`, +/// returning an error if EOF is hit sooner. +/// +/// On success the number of bytes is returned +#[derive(Debug)] +pub struct ReadExact<'a, A: ?Sized> { + reader: &'a mut A, + buf: &'a mut [u8], + pos: usize, +} + + +fn eof() -> io::Error { + io::Error::new(io::ErrorKind::UnexpectedEof, "early eof") +} + +// forward Unpin +impl<'a, A: Unpin + ?Sized> Unpin for ReadExact<'_, A> {} + +impl Future for ReadExact<'_, A> +where + A: AsyncRead + Unpin + ?Sized, +{ + type Output = io::Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + // if our buffer is empty, then we need to read some data to continue. + if self.pos < self.buf.len() { + let me = &mut *self; + let n = ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf[me.pos..]))?; + me.pos += n; + if n == 0 { + return Err(eof()).into(); + } + } + + if self.pos >= self.buf.len() { + return Poll::Ready(Ok(self.pos)); + } + } + } +} diff --git a/tokio/tests/io.rs b/tokio/tests/io.rs index f708ce605..5bf52d60f 100644 --- a/tokio/tests/io.rs +++ b/tokio/tests/io.rs @@ -132,3 +132,38 @@ fn copy() { assert_eq!(wr.0[..], b"hello world"[..]); }); } + +#[test] +fn read_exact() { + struct Rd { + val: &'static [u8; 11], + } + + impl AsyncRead for Rd { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut [u8] + ) -> Poll> { + let me = &mut *self; + let len = buf.len(); + + buf[..].copy_from_slice(&me.val[..len]); + Poll::Ready(Ok(buf.len())) + } + } + + let mut buf = Box::new([0; 8]); + let mut task = MockTask::new(); + + task.enter(|cx| { + let mut rd = Rd { val: b"hello world" }; + + let read = tokio::io::read_exact(&mut rd, &mut buf[..]); + pin_mut!(read); + + let n = assert_ready_ok!(read.poll(cx)); + assert_eq!(n, 8); + assert_eq!(buf[..], b"hello wo"[..]); + }); +}