1use std::pin::Pin;
2use std::task::{Context, Poll};
3use std::{cmp, io};
4
5use bytes::{Buf, Bytes};
6
7use crate::rt::{Read, ReadBufCursor, Write};
8
9/// Combine a buffer with an IO, rewinding reads to use the buffer.
10#[derive(Debug)]
11pub(crate) struct Rewind<T> {
12 pre: Option<Bytes>,
13 inner: T,
14}
15
16impl<T> Rewind<T> {
17 #[cfg(test)]
18 pub(crate) fn new(io: T) -> Self {
19 Rewind {
20 pre: None,
21 inner: io,
22 }
23 }
24
25 pub(crate) fn new_buffered(io: T, buf: Bytes) -> Self {
26 Rewind {
27 pre: Some(buf),
28 inner: io,
29 }
30 }
31
32 #[cfg(test)]
33 pub(crate) fn rewind(&mut self, bs: Bytes) {
34 debug_assert!(self.pre.is_none());
35 self.pre = Some(bs);
36 }
37
38 pub(crate) fn into_inner(self) -> (T, Bytes) {
39 (self.inner, self.pre.unwrap_or_default())
40 }
41
42 // pub(crate) fn get_mut(&mut self) -> &mut T {
43 // &mut self.inner
44 // }
45}
46
47impl<T> Read for Rewind<T>
48where
49 T: Read + Unpin,
50{
51 fn poll_read(
52 mut self: Pin<&mut Self>,
53 cx: &mut Context<'_>,
54 mut buf: ReadBufCursor<'_>,
55 ) -> Poll<io::Result<()>> {
56 if let Some(mut prefix: Bytes) = self.pre.take() {
57 // If there are no remaining bytes, let the bytes get dropped.
58 if !prefix.is_empty() {
59 let copy_len: usize = cmp::min(v1:prefix.len(), v2:buf.remaining());
60 // TODO: There should be a way to do following two lines cleaner...
61 buf.put_slice(&prefix[..copy_len]);
62 prefix.advance(cnt:copy_len);
63 // Put back what's left
64 if !prefix.is_empty() {
65 self.pre = Some(prefix);
66 }
67
68 return Poll::Ready(Ok(()));
69 }
70 }
71 Pin::new(&mut self.inner).poll_read(cx, buf)
72 }
73}
74
75impl<T> Write for Rewind<T>
76where
77 T: Write + Unpin,
78{
79 fn poll_write(
80 mut self: Pin<&mut Self>,
81 cx: &mut Context<'_>,
82 buf: &[u8],
83 ) -> Poll<io::Result<usize>> {
84 Pin::new(&mut self.inner).poll_write(cx, buf)
85 }
86
87 fn poll_write_vectored(
88 mut self: Pin<&mut Self>,
89 cx: &mut Context<'_>,
90 bufs: &[io::IoSlice<'_>],
91 ) -> Poll<io::Result<usize>> {
92 Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
93 }
94
95 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
96 Pin::new(&mut self.inner).poll_flush(cx)
97 }
98
99 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
100 Pin::new(&mut self.inner).poll_shutdown(cx)
101 }
102
103 fn is_write_vectored(&self) -> bool {
104 self.inner.is_write_vectored()
105 }
106}
107
108#[cfg(all(
109 any(feature = "client", feature = "server"),
110 any(feature = "http1", feature = "http2"),
111))]
112#[cfg(test)]
113mod tests {
114 use super::super::Compat;
115 use super::Rewind;
116 use bytes::Bytes;
117 use tokio::io::AsyncReadExt;
118
119 #[cfg(not(miri))]
120 #[tokio::test]
121 async fn partial_rewind() {
122 let underlying = [104, 101, 108, 108, 111];
123
124 let mock = tokio_test::io::Builder::new().read(&underlying).build();
125
126 let mut stream = Compat::new(Rewind::new(Compat::new(mock)));
127
128 // Read off some bytes, ensure we filled o1
129 let mut buf = [0; 2];
130 stream.read_exact(&mut buf).await.expect("read1");
131
132 // Rewind the stream so that it is as if we never read in the first place.
133 stream.0.rewind(Bytes::copy_from_slice(&buf[..]));
134
135 let mut buf = [0; 5];
136 stream.read_exact(&mut buf).await.expect("read1");
137
138 // At this point we should have read everything that was in the MockStream
139 assert_eq!(&buf, &underlying);
140 }
141
142 #[cfg(not(miri))]
143 #[tokio::test]
144 async fn full_rewind() {
145 let underlying = [104, 101, 108, 108, 111];
146
147 let mock = tokio_test::io::Builder::new().read(&underlying).build();
148
149 let mut stream = Compat::new(Rewind::new(Compat::new(mock)));
150
151 let mut buf = [0; 5];
152 stream.read_exact(&mut buf).await.expect("read1");
153
154 // Rewind the stream so that it is as if we never read in the first place.
155 stream.0.rewind(Bytes::copy_from_slice(&buf[..]));
156
157 let mut buf = [0; 5];
158 stream.read_exact(&mut buf).await.expect("read1");
159 }
160}
161