tokio: add read_exact method (#1202)

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};

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<A> Future for ReadExact<'_, A>
+where
+    A: AsyncRead + Unpin + ?Sized,
+{
+    type Output = io::Result<usize>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
+        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));
+            }
+        }
+    }
+}

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<io::Result<usize>> {
+            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"[..]);
+    });
+}