futures_util/io/
read_exact.rs

1use std::io;
2use std::mem;
3
4use {Poll, Future, task};
5
6use io::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 {
22        a: A,
23        buf: T,
24        pos: usize,
25    },
26    Empty,
27}
28
29pub fn read_exact<A, T>(a: A, buf: T) -> ReadExact<A, T>
30    where A: AsyncRead,
31          T: AsMut<[u8]>,
32{
33    ReadExact {
34        state: State::Reading {
35            a,
36            buf,
37            pos: 0,
38        },
39    }
40}
41
42fn eof() -> io::Error {
43    io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")
44}
45
46impl<A, T> Future for ReadExact<A, T>
47    where A: AsyncRead,
48          T: AsMut<[u8]>,
49{
50    type Item = (A, T);
51    type Error = io::Error;
52
53    fn poll(&mut self, cx: &mut task::Context) -> Poll<(A, T), io::Error> {
54        match self.state {
55            State::Reading { ref mut a, ref mut buf, ref mut pos } => {
56                let buf = buf.as_mut();
57                while *pos < buf.len() {
58                    let n = try_ready!(a.poll_read(cx, &mut buf[*pos..]));
59                    *pos += n;
60                    if n == 0 {
61                        return Err(eof())
62                    }
63                }
64            }
65            State::Empty => panic!("poll a ReadExact after it's done"),
66        }
67
68        match mem::replace(&mut self.state, State::Empty) {
69            State::Reading { a, buf, .. } => Ok((a, buf).into()),
70            State::Empty => panic!(),
71        }
72    }
73}