1 | use crate::io::{AsyncRead, ReadBuf}; |
2 | |
3 | use pin_project_lite::pin_project; |
4 | use std::future::Future; |
5 | use std::io; |
6 | use std::marker::PhantomPinned; |
7 | use std::marker::Unpin; |
8 | use std::pin::Pin; |
9 | use std::task::{Context, Poll}; |
10 | |
11 | /// A future which can be used to easily read exactly enough bytes to fill |
12 | /// a buffer. |
13 | /// |
14 | /// Created by the [`AsyncReadExt::read_exact`][read_exact]. |
15 | /// [read_exact]: [crate::io::AsyncReadExt::read_exact] |
16 | pub(crate) fn read_exact<'a, A>(reader: &'a mut A, buf: &'a mut [u8]) -> ReadExact<'a, A> |
17 | where |
18 | A: AsyncRead + Unpin + ?Sized, |
19 | { |
20 | ReadExact { |
21 | reader, |
22 | buf: ReadBuf::new(buf), |
23 | _pin: PhantomPinned, |
24 | } |
25 | } |
26 | |
27 | pin_project! { |
28 | /// Creates a future which will read exactly enough bytes to fill `buf`, |
29 | /// returning an error if EOF is hit sooner. |
30 | /// |
31 | /// On success the number of bytes is returned |
32 | #[derive(Debug)] |
33 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
34 | pub struct ReadExact<'a, A: ?Sized> { |
35 | reader: &'a mut A, |
36 | buf: ReadBuf<'a>, |
37 | // Make this future `!Unpin` for compatibility with async trait methods. |
38 | #[pin] |
39 | _pin: PhantomPinned, |
40 | } |
41 | } |
42 | |
43 | fn eof() -> io::Error { |
44 | io::Error::new(io::ErrorKind::UnexpectedEof, "early eof" ) |
45 | } |
46 | |
47 | impl<A> Future for ReadExact<'_, A> |
48 | where |
49 | A: AsyncRead + Unpin + ?Sized, |
50 | { |
51 | type Output = io::Result<usize>; |
52 | |
53 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { |
54 | let me = self.project(); |
55 | |
56 | loop { |
57 | // if our buffer is empty, then we need to read some data to continue. |
58 | let rem = me.buf.remaining(); |
59 | if rem != 0 { |
60 | ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf))?; |
61 | if me.buf.remaining() == rem { |
62 | return Err(eof()).into(); |
63 | } |
64 | } else { |
65 | return Poll::Ready(Ok(me.buf.capacity())); |
66 | } |
67 | } |
68 | } |
69 | } |
70 | |