ferron/util/
read_to_end_move.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4use tokio::io::{self, AsyncRead, ReadBuf};
5
6pub struct ReadToEndFuture<R> {
7  reader: R,
8  buffer: Vec<u8>,
9}
10
11impl<R> ReadToEndFuture<R> {
12  pub fn new(reader: R) -> Self {
13    Self {
14      reader,
15      buffer: Vec::new(),
16    }
17  }
18}
19
20impl<R> Future for ReadToEndFuture<R>
21where
22  R: AsyncRead + Unpin,
23{
24  type Output = io::Result<Vec<u8>>;
25
26  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
27    let mut buf = [0; 1024];
28    let mut read_buf = ReadBuf::new(&mut buf);
29
30    loop {
31      match Pin::new(&mut self.reader).poll_read(cx, &mut read_buf) {
32        Poll::Ready(Ok(())) => {
33          let n = read_buf.filled().len();
34          if n == 0 {
35            return Poll::Ready(Ok(self.buffer.clone()));
36          }
37          self.buffer.extend_from_slice(read_buf.filled());
38          read_buf.clear();
39        }
40        Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
41        Poll::Pending => return Poll::Pending,
42      }
43    }
44  }
45}
46
47#[cfg(test)]
48mod tests {
49  use super::*;
50  use std::pin::Pin;
51  use std::task::{Context, Poll};
52  use tokio::io::{self, AsyncRead};
53
54  struct MockReader {
55    data: Vec<u8>,
56    position: usize,
57  }
58
59  impl MockReader {
60    fn new(data: &[u8]) -> Self {
61      Self {
62        data: data.to_vec(),
63        position: 0,
64      }
65    }
66  }
67
68  impl AsyncRead for MockReader {
69    fn poll_read(
70      mut self: Pin<&mut Self>,
71      _cx: &mut Context<'_>,
72      buf: &mut ReadBuf<'_>,
73    ) -> Poll<io::Result<()>> {
74      if self.position >= self.data.len() {
75        return Poll::Ready(Ok(()));
76      }
77
78      let end = (self.position + buf.remaining()).min(self.data.len());
79      buf.put_slice(&self.data[self.position..end]);
80      self.position = end;
81
82      Poll::Ready(Ok(()))
83    }
84  }
85
86  #[tokio::test]
87  async fn test_read_to_end_empty_reader() {
88    let reader = MockReader::new(&[]);
89    let future = ReadToEndFuture::new(reader);
90    let result = future.await;
91    assert_eq!(result.unwrap(), Vec::<u8>::new());
92  }
93
94  #[tokio::test]
95  async fn test_read_to_end_non_empty_reader() {
96    let reader = MockReader::new(b"hello world");
97    let future = ReadToEndFuture::new(reader);
98    let result = future.await;
99    assert_eq!(result.unwrap(), b"hello world");
100  }
101
102  struct ErrorReader;
103
104  impl AsyncRead for ErrorReader {
105    fn poll_read(
106      self: Pin<&mut Self>,
107      _cx: &mut Context<'_>,
108      _buf: &mut ReadBuf<'_>,
109    ) -> Poll<io::Result<()>> {
110      Poll::Ready(Err(io::Error::other("read error")))
111    }
112  }
113
114  #[tokio::test]
115  async fn test_read_to_end_error() {
116    let reader = ErrorReader;
117    let future = ReadToEndFuture::new(reader);
118    let result = future.await;
119    assert!(result.is_err());
120    assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Other);
121  }
122}