rama/http/client/
proxy_connector.rs

1use crate::{
2    Layer, Service,
3    error::{BoxError, BoxErrorExt, ErrorContext as _, ErrorExt as _},
4    extensions::{Extensions, ExtensionsRef},
5    http::client::proxy::layer::{
6        HttpProxyConnector, HttpProxyConnectorLayer, MaybeHttpProxiedConnection,
7    },
8    io::Io,
9    net::{
10        AuthorityInputExt, Protocol, ProtocolInputExt,
11        address::ProxyAddress,
12        client::{ConnectorService, EstablishedClientConnection},
13    },
14    proxy::socks5::{Socks5ProxyConnector, Socks5ProxyConnectorLayer},
15    telemetry::tracing,
16};
17use pin_project_lite::pin_project;
18use std::{
19    fmt::Debug,
20    pin::Pin,
21    task::{self, Poll},
22};
23use tokio::io::{AsyncRead, AsyncWrite};
24
25/// Proxy connector which supports http(s) and socks5(h) proxy address
26///
27/// Connector will look at [`ProxyAddress`] to determine which proxy
28/// connector to use if one is configured
29#[derive(Debug, Clone)]
30pub struct ProxyConnector<S> {
31    inner: S,
32    socks: Socks5ProxyConnector<S>,
33    http: HttpProxyConnector<S>,
34    required: bool,
35}
36
37impl<S: Clone> ProxyConnector<S> {
38    /// Creates a new [`ProxyConnector`].
39    fn new(
40        inner: S,
41        socks_proxy_layer: Socks5ProxyConnectorLayer,
42        http_proxy_layer: HttpProxyConnectorLayer,
43        required: bool,
44    ) -> Self {
45        Self {
46            socks: socks_proxy_layer.into_layer(inner.clone()),
47            http: http_proxy_layer.into_layer(inner.clone()),
48            inner,
49            required,
50        }
51    }
52
53    #[inline]
54    /// Creates a new required [`ProxyConnector`].
55    ///
56    /// This connector will fail if no [`ProxyAddress`] is configured
57    pub fn required(
58        inner: S,
59        socks_proxy_layer: Socks5ProxyConnectorLayer,
60        http_proxy_layer: HttpProxyConnectorLayer,
61    ) -> Self {
62        Self::new(inner, socks_proxy_layer, http_proxy_layer, true)
63    }
64
65    #[inline]
66    /// Creates a new optional [`ProxyConnector`].
67    ///
68    /// This connector will forward to the inner connector if no [`ProxyAddress`] is configured
69    pub fn optional(
70        inner: S,
71        socks_proxy_layer: Socks5ProxyConnectorLayer,
72        http_proxy_layer: HttpProxyConnectorLayer,
73    ) -> Self {
74        Self::new(inner, socks_proxy_layer, http_proxy_layer, false)
75    }
76}
77
78impl<Input, S> Service<Input> for ProxyConnector<S>
79where
80    S: ConnectorService<Input, Connection: Io + Unpin>,
81    Input: AuthorityInputExt + ProtocolInputExt + Send + ExtensionsRef + 'static,
82{
83    type Output = EstablishedClientConnection<MaybeProxiedConnection<S::Connection>, Input>;
84    type Error = BoxError;
85
86    async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
87        let proxy = input.extensions().get_ref::<ProxyAddress>();
88
89        match proxy {
90            None => {
91                if self.required {
92                    return Err(BoxError::from_static_str(
93                        "proxy required but none is defined",
94                    ));
95                }
96                tracing::trace!("no proxy detected in ctx, using inner connector");
97                let EstablishedClientConnection { input, conn } =
98                    self.inner.connect(input).await.into_box_error()?;
99
100                let conn = MaybeProxiedConnection::direct(conn);
101                Ok(EstablishedClientConnection { input, conn })
102            }
103            Some(proxy) => {
104                let protocol = proxy.protocol.as_ref();
105                tracing::trace!(?protocol, "proxy detected in ctx");
106
107                let protocol = protocol.unwrap_or_else(|| {
108                    tracing::trace!("no protocol detected, using http as protocol");
109                    &Protocol::HTTP
110                });
111
112                if protocol.is_socks5() {
113                    tracing::trace!(
114                        target = %&proxy.address,
115                        "using socks proxy connector",
116                    );
117
118                    let EstablishedClientConnection { input, conn } =
119                        self.socks.connect(input).await?;
120
121                    let conn = MaybeProxiedConnection::socks(conn);
122                    Ok(EstablishedClientConnection { input, conn })
123                } else if protocol.is_http() {
124                    tracing::trace!(
125                        target = %&proxy.address,
126                        "using http proxy connector"
127                    );
128
129                    let EstablishedClientConnection { input, conn } =
130                        self.http.connect(input).await?;
131
132                    let conn = MaybeProxiedConnection::http(conn);
133                    Ok(EstablishedClientConnection { input, conn })
134                } else {
135                    Err(
136                        BoxError::from_static_str("received unsupport proxy protocol")
137                            .with_context_debug_field("protocol", || protocol.clone()),
138                    )
139                }
140            }
141        }
142    }
143}
144
145pin_project! {
146    /// A connection which will be proxied if a [`ProxyAddress`] was configured
147    pub struct MaybeProxiedConnection<S> {
148        #[pin]
149        inner: Connection<S>,
150    }
151}
152
153impl<S: ExtensionsRef> MaybeProxiedConnection<S> {
154    pub fn direct(conn: S) -> Self {
155        Self {
156            inner: Connection::Direct { conn },
157        }
158    }
159
160    pub fn socks(conn: S) -> Self {
161        Self {
162            inner: Connection::Socks { conn },
163        }
164    }
165
166    pub fn http(conn: MaybeHttpProxiedConnection<S>) -> Self {
167        Self {
168            inner: Connection::Http { conn },
169        }
170    }
171}
172
173impl<S: Debug> Debug for MaybeProxiedConnection<S> {
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        f.debug_struct("MaybeProxiedConnection")
176            .field("inner", &self.inner)
177            .finish()
178    }
179}
180
181impl<S: ExtensionsRef> ExtensionsRef for MaybeProxiedConnection<S> {
182    fn extensions(&self) -> &Extensions {
183        match &self.inner {
184            Connection::Direct { conn } | Connection::Socks { conn } => conn.extensions(),
185            Connection::Http { conn } => conn.extensions(),
186        }
187    }
188}
189
190pin_project! {
191    #[project = ConnectionProj]
192    enum Connection<S> {
193        Direct{ #[pin] conn: S },
194        Socks{ #[pin] conn: S },
195        Http{ #[pin] conn: MaybeHttpProxiedConnection<S> },
196
197    }
198}
199
200impl<S: Debug> Debug for Connection<S> {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        match self {
203            Self::Direct { conn } => f.debug_struct("Direct").field("conn", conn).finish(),
204            Self::Socks { conn } => f.debug_struct("Socks").field("conn", conn).finish(),
205            Self::Http { conn } => f.debug_struct("Http").field("conn", conn).finish(),
206        }
207    }
208}
209
210#[warn(clippy::missing_trait_methods)]
211impl<Conn: AsyncWrite> AsyncWrite for MaybeProxiedConnection<Conn> {
212    fn poll_write(
213        self: Pin<&mut Self>,
214        cx: &mut task::Context<'_>,
215        buf: &[u8],
216    ) -> Poll<Result<usize, std::io::Error>> {
217        match self.project().inner.project() {
218            ConnectionProj::Direct { conn } | ConnectionProj::Socks { conn } => {
219                conn.poll_write(cx, buf)
220            }
221            ConnectionProj::Http { conn } => conn.poll_write(cx, buf),
222        }
223    }
224
225    fn poll_flush(
226        self: Pin<&mut Self>,
227        cx: &mut task::Context<'_>,
228    ) -> Poll<Result<(), std::io::Error>> {
229        match self.project().inner.project() {
230            ConnectionProj::Direct { conn } | ConnectionProj::Socks { conn } => conn.poll_flush(cx),
231            ConnectionProj::Http { conn } => conn.poll_flush(cx),
232        }
233    }
234
235    fn poll_shutdown(
236        self: Pin<&mut Self>,
237        cx: &mut task::Context<'_>,
238    ) -> Poll<Result<(), std::io::Error>> {
239        match self.project().inner.project() {
240            ConnectionProj::Direct { conn } | ConnectionProj::Socks { conn } => {
241                conn.poll_shutdown(cx)
242            }
243            ConnectionProj::Http { conn } => conn.poll_shutdown(cx),
244        }
245    }
246
247    fn is_write_vectored(&self) -> bool {
248        match &self.inner {
249            Connection::Direct { conn } | Connection::Socks { conn } => conn.is_write_vectored(),
250            Connection::Http { conn } => conn.is_write_vectored(),
251        }
252    }
253
254    fn poll_write_vectored(
255        self: Pin<&mut Self>,
256        cx: &mut task::Context<'_>,
257        bufs: &[std::io::IoSlice<'_>],
258    ) -> Poll<Result<usize, std::io::Error>> {
259        match self.project().inner.project() {
260            ConnectionProj::Direct { conn } | ConnectionProj::Socks { conn } => {
261                conn.poll_write_vectored(cx, bufs)
262            }
263            ConnectionProj::Http { conn } => conn.poll_write_vectored(cx, bufs),
264        }
265    }
266}
267
268#[warn(clippy::missing_trait_methods)]
269impl<Conn: AsyncRead> AsyncRead for MaybeProxiedConnection<Conn> {
270    fn poll_read(
271        self: Pin<&mut Self>,
272        cx: &mut task::Context<'_>,
273        buf: &mut tokio::io::ReadBuf<'_>,
274    ) -> Poll<std::io::Result<()>> {
275        match self.project().inner.project() {
276            ConnectionProj::Direct { conn } | ConnectionProj::Socks { conn } => {
277                conn.poll_read(cx, buf)
278            }
279            ConnectionProj::Http { conn } => conn.poll_read(cx, buf),
280        }
281    }
282}
283
284/// Proxy connector layer which supports http(s) and socks5(h) proxy address
285///
286/// Connector will look at [`ProxyAddress`] to determine which proxy
287/// connector to use if one is configured
288pub struct ProxyConnectorLayer {
289    socks_layer: Socks5ProxyConnectorLayer,
290    http_layer: HttpProxyConnectorLayer,
291    required: bool,
292}
293
294impl ProxyConnectorLayer {
295    #[must_use]
296    /// Creates a new required [`ProxyConnectorLayer`].
297    ///
298    /// This connector will fail if no [`ProxyAddress`] is configured
299    pub fn required(
300        socks_proxy_layer: Socks5ProxyConnectorLayer,
301        http_proxy_layer: HttpProxyConnectorLayer,
302    ) -> Self {
303        Self {
304            socks_layer: socks_proxy_layer,
305            http_layer: http_proxy_layer,
306            required: true,
307        }
308    }
309
310    #[must_use]
311    /// Creates a new optional [`ProxyConnectorLayer`].
312    ///
313    /// This connector will forward to the inner connector if no [`ProxyAddress`] is configured
314    pub fn optional(
315        socks_proxy_layer: Socks5ProxyConnectorLayer,
316        http_proxy_layer: HttpProxyConnectorLayer,
317    ) -> Self {
318        Self {
319            socks_layer: socks_proxy_layer,
320            http_layer: http_proxy_layer,
321            required: false,
322        }
323    }
324}
325
326impl<S: Clone> Layer<S> for ProxyConnectorLayer {
327    type Service = ProxyConnector<S>;
328
329    fn layer(&self, inner: S) -> Self::Service {
330        ProxyConnector::new(
331            inner,
332            self.socks_layer.clone(),
333            self.http_layer.clone(),
334            self.required,
335        )
336    }
337
338    fn into_layer(self, inner: S) -> Self::Service {
339        ProxyConnector::new(inner, self.socks_layer, self.http_layer, self.required)
340    }
341}
rama/http/client/proxy_connector.rs

rama/http/client/
proxy_connector.rs
