hyper/common/io/
rewind.rs

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

hyper/common/io/
rewind.rs
