hyper_util/common/
rewind.rs

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

hyper_util/common/
rewind.rs
