tokio_io/io/
read_exact.rs

1use std::io;
2use std::mem;
3
4use futures::{Future, Poll};
5
6use AsyncRead;
7
8/// A future which can be used to easily read exactly enough bytes to fill
9/// a buffer.
10///
11/// Created by the [`read_exact`] function.
12///
13/// [`read_exact`]: fn.read_exact.html
14#[derive(Debug)]
15pub struct ReadExact<A, T> {
16    state: State<A, T>,
17}
18
19#[derive(Debug)]
20enum State<A, T> {
21    Reading { a: A, buf: T, pos: usize },
22    Empty,
23}
24
25/// Creates a future which will read exactly enough bytes to fill `buf`,
26/// returning an error if EOF is hit sooner.
27///
28/// The returned future will resolve to both the I/O stream as well as the
29/// buffer once the read operation is completed.
30///
31/// In the case of an error the buffer and the object will be discarded, with
32/// the error yielded. In the case of success the object will be destroyed and
33/// the buffer will be returned, with all data read from the stream appended to
34/// the buffer.
35pub fn read_exact<A, T>(a: A, buf: T) -> ReadExact<A, T>
36where
37    A: AsyncRead,
38    T: AsMut<[u8]>,
39{
40    ReadExact {
41        state: State::Reading {
42            a: a,
43            buf: buf,
44            pos: 0,
45        },
46    }
47}
48
49fn eof() -> io::Error {
50    io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")
51}
52
53impl<A, T> Future for ReadExact<A, T>
54where
55    A: AsyncRead,
56    T: AsMut<[u8]>,
57{
58    type Item = (A, T);
59    type Error = io::Error;
60
61    fn poll(&mut self) -> Poll<(A, T), io::Error> {
62        match self.state {
63            State::Reading {
64                ref mut a,
65                ref mut buf,
66                ref mut pos,
67            } => {
68                let buf = buf.as_mut();
69                while *pos < buf.len() {
70                    let n = try_ready!(a.poll_read(&mut buf[*pos..]));
71                    *pos += n;
72                    if n == 0 {
73                        return Err(eof());
74                    }
75                }
76            }
77            State::Empty => panic!("poll a ReadExact after it's done"),
78        }
79
80        match mem::replace(&mut self.state, State::Empty) {
81            State::Reading { a, buf, .. } => Ok((a, buf).into()),
82            State::Empty => panic!(),
83        }
84    }
85}