rama/http/client/
proxy_connector.rs

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

rama/http/client/
proxy_connector.rs
