hudsucker 0.25.0

MITM HTTP/S proxy
Documentation
use crate::rewind::Rewind;
use std::{
    io::{self, IoSlice},
    pin::Pin,
    task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

pub(crate) struct Tee<T> {
    prefix: Vec<u8>,
    inner: T,
}

impl<T> Tee<T> {
    pub(crate) fn new(io: T) -> Self {
        Tee {
            prefix: Vec::new(),
            inner: io,
        }
    }

    pub(crate) fn into_inner(self) -> T {
        self.inner
    }

    pub(crate) fn rewind(self) -> Rewind<Vec<u8>, T> {
        Rewind::new(self.prefix, self.inner)
    }
}

impl<T> AsyncRead for Tee<T>
where
    T: AsyncRead + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let filled_before = buf.filled().len();
        let res = Pin::new(&mut self.inner).poll_read(cx, buf);

        if res.is_ready() {
            self.prefix
                .extend_from_slice(&buf.filled()[filled_before..]);
        }

        res
    }
}

impl<T> AsyncWrite for Tee<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()
    }
}