hyper_util/common/
rewind.rs

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

hyper_util/common/
rewind.rs
