product-os-proxy 0.0.19

// Adapted from https://github.com/hyperium/hyper/blob/master/src/common/io/rewind.rs
//
//! Rewind IO adapter for buffered reads.
//!
//! This module provides a [`Rewind`] wrapper that allows prepending buffered data
//! to an async IO stream. This is used during HTTP CONNECT handling where we need
//! to peek at the initial bytes of an upgraded connection to determine the protocol
//! (TLS, WebSocket, plain HTTP) and then "rewind" those bytes back into the stream
//! for the actual protocol handler to consume.

use hyper::body::Bytes;
use std::{
    cmp,
    io::{self, IoSlice},
    pin::Pin,
    task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

/// Combine a buffer with an IO stream, rewinding reads to use the buffer first.
///
/// When reading from a `Rewind<T>`, any buffered bytes (the "prefix") are returned
/// first before delegating to the underlying IO stream. This allows "peeking" at
/// data from a stream and then transparently replaying it.
///
/// Writes are always delegated directly to the underlying stream.
///
/// # Type Parameters
///
/// * `T` - The underlying async IO type (must implement [`AsyncRead`] and/or [`AsyncWrite`])
///
/// # Examples
///
/// ```ignore
/// use bytes::Bytes;
///
/// let prefix = Bytes::from_static(b"hello ");
/// let rewind = Rewind::new(inner_stream, prefix);
/// // First read will return "hello " before reading from inner_stream
/// ```
#[derive(Debug)]
pub(crate) struct Rewind<T> {
    pre: Option<Bytes>,
    inner: T,
}

impl<T> Rewind<T> {
    /// Creates a new `Rewind` wrapper with the given prefix buffer.
    ///
    /// # Arguments
    ///
    /// * `io` - The underlying IO stream
    /// * `buf` - The prefix bytes to return before reading from the stream
    pub(crate) fn new(io: T, buf: Bytes) -> Self {
        Rewind {
            pre: Some(buf),
            inner: io,
        }
    }
}

impl<T> AsyncRead for Rewind<T>
where
    T: AsyncRead + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if let Some(mut prefix) = self.pre.take() {
            // If there are no remaining bytes, let the bytes get dropped.
            if !prefix.is_empty() {
                let copy_len = cmp::min(prefix.len(), buf.remaining());
                buf.put_slice(&prefix.split_to(copy_len));
                // Put back what's left
                if !prefix.is_empty() {
                    self.pre = Some(prefix);
                }

                return Poll::Ready(Ok(()));
            }
        }

        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl<T> AsyncWrite for Rewind<T>
where
    T: AsyncWrite + Unpin,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, buf)
    }

    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }

    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::AsyncReadExt;

    #[tokio::test]
    async fn test_rewind_returns_prefix_first() {
        let inner = tokio::io::empty();
        let prefix = Bytes::from_static(b"hello");
        let mut rewind = Rewind::new(inner, prefix);

        let mut buf = [0u8; 10];
        let n = rewind.read(&mut buf).await.unwrap();
        assert_eq!(n, 5);
        assert_eq!(&buf[..n], b"hello");
    }

    #[tokio::test]
    async fn test_rewind_empty_prefix_reads_inner() {
        let inner = &b"world"[..];
        let prefix = Bytes::new();
        let mut rewind = Rewind::new(inner, prefix);

        let mut buf = [0u8; 10];
        let n = rewind.read(&mut buf).await.unwrap();
        assert_eq!(n, 5);
        assert_eq!(&buf[..n], b"world");
    }

    #[tokio::test]
    async fn test_rewind_partial_prefix_read() {
        let inner = &b"world"[..];
        let prefix = Bytes::from_static(b"hello ");
        let mut rewind = Rewind::new(inner, prefix);

        // Read only 3 bytes - should get partial prefix
        let mut buf = [0u8; 3];
        let n = rewind.read(&mut buf).await.unwrap();
        assert_eq!(n, 3);
        assert_eq!(&buf[..n], b"hel");

        // Read more - should get rest of prefix
        let n = rewind.read(&mut buf).await.unwrap();
        assert_eq!(n, 3);
        assert_eq!(&buf[..n], b"lo ");

        // Read more - should get inner stream
        let mut buf = [0u8; 10];
        let n = rewind.read(&mut buf).await.unwrap();
        assert_eq!(n, 5);
        assert_eq!(&buf[..n], b"world");
    }
}