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
24pub 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}