Skip to main content

tokio_core/io/
read_exact.rs

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