ferron/util/
read_to_end_move.rs1use 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}