hyper1_proxy/
stream.rs

1use hyper::rt::{Read, Write};
2use hyper_util::client::legacy::connect::{Connected, Connection};
3use hyper_util::rt::TokioIo;
4use std::io;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
8
9#[cfg(feature = "rustls-base")]
10use tokio_rustls::client::TlsStream as RustlsStream;
11
12#[cfg(feature = "tls")]
13use tokio_native_tls::TlsStream;
14
15#[cfg(feature = "openssl-tls")]
16use tokio_openssl::SslStream as OpenSslStream;
17
18#[cfg(feature = "rustls-base")]
19pub type TlsStream<R> = RustlsStream<R>;
20
21#[cfg(feature = "openssl-tls")]
22pub type TlsStream<R> = OpenSslStream<R>;
23
24/// A Proxy Stream wrapper
25pub enum ProxyStream<R> {
26    NoProxy(R),
27    Regular(R),
28    #[cfg(any(feature = "tls", feature = "rustls-base", feature = "openssl-tls"))]
29    Secured(TokioIo<TlsStream<TokioIo<R>>>),
30}
31
32macro_rules! match_fn_pinned {
33    ($self:expr, $fn:ident, $ctx:expr, $buf:expr) => {
34        match $self.get_mut() {
35            ProxyStream::NoProxy(s) => Pin::new(s).$fn($ctx, $buf),
36            ProxyStream::Regular(s) => Pin::new(s).$fn($ctx, $buf),
37            #[cfg(any(feature = "tls", feature = "rustls-base", feature = "openssl-tls"))]
38            ProxyStream::Secured(s) => Pin::new(s).$fn($ctx, $buf),
39        }
40    };
41
42    ($self:expr, $fn:ident, $ctx:expr) => {
43        match $self.get_mut() {
44            ProxyStream::NoProxy(s) => Pin::new(s).$fn($ctx),
45            ProxyStream::Regular(s) => Pin::new(s).$fn($ctx),
46            #[cfg(any(feature = "tls", feature = "rustls-base", feature = "openssl-tls"))]
47            ProxyStream::Secured(s) => Pin::new(s).$fn($ctx),
48        }
49    };
50}
51
52impl<R: Read + Write + Unpin> Read for ProxyStream<R> {
53    fn poll_read(
54        self: Pin<&mut Self>,
55        cx: &mut Context<'_>,
56        buf: hyper::rt::ReadBufCursor<'_>,
57    ) -> Poll<Result<(), std::io::Error>> {
58        match_fn_pinned!(self, poll_read, cx, buf)
59    }
60}
61
62impl<R: Read + Write + Unpin> Write for ProxyStream<R> {
63    fn poll_write(
64        self: Pin<&mut Self>,
65        cx: &mut Context<'_>,
66        buf: &[u8],
67    ) -> Poll<io::Result<usize>> {
68        match_fn_pinned!(self, poll_write, cx, buf)
69    }
70
71    fn poll_write_vectored(
72        self: Pin<&mut Self>,
73        cx: &mut Context<'_>,
74        bufs: &[io::IoSlice<'_>],
75    ) -> Poll<Result<usize, io::Error>> {
76        match_fn_pinned!(self, poll_write_vectored, cx, bufs)
77    }
78
79    fn is_write_vectored(&self) -> bool {
80        match self {
81            ProxyStream::NoProxy(s) => s.is_write_vectored(),
82            ProxyStream::Regular(s) => s.is_write_vectored(),
83            #[cfg(any(feature = "tls", feature = "rustls-base", feature = "openssl-tls"))]
84            ProxyStream::Secured(s) => s.is_write_vectored(),
85        }
86    }
87
88    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
89        match_fn_pinned!(self, poll_flush, cx)
90    }
91
92    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
93        match_fn_pinned!(self, poll_shutdown, cx)
94    }
95}
96
97impl<R: Read + Write + Connection + Unpin> Connection for ProxyStream<R> {
98    fn connected(&self) -> Connected {
99        match self {
100            ProxyStream::NoProxy(s) => s.connected(),
101
102            ProxyStream::Regular(s) => s.connected().proxy(true),
103            #[cfg(feature = "tls")]
104            ProxyStream::Secured(s) => {
105                let inner = s.inner().get_ref().get_ref().get_ref().inner();
106                inner.connected().proxy(true)
107            }
108
109            #[cfg(feature = "rustls-base")]
110            ProxyStream::Secured(s) => s.get_ref().0.connected().proxy(true),
111
112            #[cfg(feature = "openssl-tls")]
113            ProxyStream::Secured(s) => s.get_ref().connected().proxy(true),
114        }
115    }
116}