hyper/common/io/
rewind.rs

Help
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) = self.pre.take() {
57            // If there are no remaining bytes, let the bytes get dropped.
58            if !prefix.is_empty() {
59                let copy_len = cmp::min(prefix.len(), 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(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}
hyper/common/io/rewind.rs

hyper/common/io/
rewind.rs
