temporalio_client/
proxy.rs

1use base64::prelude::*;
2use http_body_util::Empty;
3use hyper::{body::Bytes, header};
4use hyper_util::{
5    client::legacy::{
6        Client,
7        connect::{Connected, Connection},
8    },
9    rt::{TokioExecutor, TokioIo},
10};
11use std::{
12    future::Future,
13    io,
14    pin::Pin,
15    task::{Context, Poll},
16};
17use tokio::{
18    io::{AsyncRead, AsyncWrite, ReadBuf},
19    net::TcpStream,
20};
21use tonic::transport::{Channel, Endpoint};
22use tower::{Service, service_fn};
23
24#[cfg(unix)]
25use tokio::net::UnixStream;
26
27/// Options for HTTP CONNECT proxy.
28#[derive(Clone, Debug, bon::Builder)]
29#[builder(start_fn = new, on(String, into))]
30#[non_exhaustive]
31pub struct HttpConnectProxyOptions {
32    /// The host:port to proxy through for TCP, or unix:/path/to/unix.sock for
33    /// Unix socket (which means it must start with "unix:/").
34    #[builder(start_fn)]
35    pub target_addr: String,
36    /// Optional HTTP basic auth for the proxy as user/pass tuple.
37    pub basic_auth: Option<(String, String)>,
38}
39
40impl HttpConnectProxyOptions {
41    /// Create a channel from the given endpoint that uses the HTTP CONNECT proxy.
42    pub async fn connect_endpoint(
43        &self,
44        endpoint: &Endpoint,
45    ) -> Result<Channel, tonic::transport::Error> {
46        let proxy_options = self.clone();
47        let svc_fn = service_fn(move |uri: tonic::transport::Uri| {
48            let proxy_options = proxy_options.clone();
49            async move { proxy_options.connect(uri).await }
50        });
51        endpoint.connect_with_connector(svc_fn).await
52    }
53
54    async fn connect(
55        &self,
56        uri: tonic::transport::Uri,
57    ) -> anyhow::Result<hyper::upgrade::Upgraded> {
58        let uri = ensure_connect_authority_port(uri);
59        debug!("Connecting to {} via proxy at {}", uri, self.target_addr);
60        // Create CONNECT request
61        let mut req_build = hyper::Request::builder().method("CONNECT").uri(uri);
62        if let Some((user, pass)) = &self.basic_auth {
63            let creds = BASE64_STANDARD.encode(format!("{user}:{pass}"));
64            req_build = req_build.header(header::PROXY_AUTHORIZATION, format!("Basic {creds}"));
65        }
66        let req = req_build.body(Empty::<Bytes>::new())?;
67
68        // We have to create a client with a specific connector because Hyper is
69        // not letting us change the HTTP/2 authority
70        let client = Client::builder(TokioExecutor::new())
71            .build(OverrideAddrConnector(self.target_addr.clone()));
72
73        // Send request
74        let res = client.request(req).await?;
75        if res.status().is_success() {
76            Ok(hyper::upgrade::on(res).await?)
77        } else {
78            Err(anyhow::anyhow!(
79                "CONNECT call failed with status: {}",
80                res.status()
81            ))
82        }
83    }
84}
85
86#[derive(Clone)]
87struct OverrideAddrConnector(String);
88
89impl Service<hyper::Uri> for OverrideAddrConnector {
90    type Response = TokioIo<ProxyStream>;
91
92    type Error = anyhow::Error;
93
94    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
95
96    fn poll_ready(&mut self, _ctx: &mut Context<'_>) -> Poll<anyhow::Result<()>> {
97        Poll::Ready(Ok(()))
98    }
99
100    fn call(&mut self, _uri: hyper::Uri) -> Self::Future {
101        let target_addr = self.0.clone();
102        let fut = async move {
103            Ok(TokioIo::new(
104                ProxyStream::connect(target_addr.as_str()).await?,
105            ))
106        };
107        Box::pin(fut)
108    }
109}
110
111/// Visible only for tests
112#[doc(hidden)]
113pub enum ProxyStream {
114    Tcp(TcpStream),
115    #[cfg(unix)]
116    Unix(UnixStream),
117}
118
119impl ProxyStream {
120    async fn connect(target_addr: &str) -> anyhow::Result<Self> {
121        if target_addr.starts_with("unix:/") {
122            #[cfg(unix)]
123            {
124                Ok(ProxyStream::Unix(
125                    UnixStream::connect(&target_addr[5..]).await?,
126                ))
127            }
128            #[cfg(not(unix))]
129            {
130                Err(anyhow::anyhow!(
131                    "Unix sockets are not supported on this platform"
132                ))
133            }
134        } else {
135            Ok(ProxyStream::Tcp(TcpStream::connect(target_addr).await?))
136        }
137    }
138}
139
140impl AsyncRead for ProxyStream {
141    fn poll_read(
142        self: Pin<&mut Self>,
143        cx: &mut Context<'_>,
144        buf: &mut ReadBuf<'_>,
145    ) -> Poll<io::Result<()>> {
146        match self.get_mut() {
147            ProxyStream::Tcp(s) => Pin::new(s).poll_read(cx, buf),
148            #[cfg(unix)]
149            ProxyStream::Unix(s) => Pin::new(s).poll_read(cx, buf),
150        }
151    }
152}
153
154impl AsyncWrite for ProxyStream {
155    fn poll_write(
156        self: Pin<&mut Self>,
157        cx: &mut Context<'_>,
158        buf: &[u8],
159    ) -> Poll<io::Result<usize>> {
160        match self.get_mut() {
161            ProxyStream::Tcp(s) => Pin::new(s).poll_write(cx, buf),
162            #[cfg(unix)]
163            ProxyStream::Unix(s) => Pin::new(s).poll_write(cx, buf),
164        }
165    }
166
167    fn poll_write_vectored(
168        self: Pin<&mut Self>,
169        cx: &mut Context<'_>,
170        bufs: &[io::IoSlice<'_>],
171    ) -> Poll<io::Result<usize>> {
172        match self.get_mut() {
173            ProxyStream::Tcp(s) => Pin::new(s).poll_write_vectored(cx, bufs),
174            #[cfg(unix)]
175            ProxyStream::Unix(s) => Pin::new(s).poll_write_vectored(cx, bufs),
176        }
177    }
178
179    fn is_write_vectored(&self) -> bool {
180        match self {
181            ProxyStream::Tcp(s) => s.is_write_vectored(),
182            #[cfg(unix)]
183            ProxyStream::Unix(s) => s.is_write_vectored(),
184        }
185    }
186
187    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
188        match self.get_mut() {
189            ProxyStream::Tcp(s) => Pin::new(s).poll_flush(cx),
190            #[cfg(unix)]
191            ProxyStream::Unix(s) => Pin::new(s).poll_flush(cx),
192        }
193    }
194
195    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
196        match self.get_mut() {
197            ProxyStream::Tcp(s) => Pin::new(s).poll_shutdown(cx),
198            #[cfg(unix)]
199            ProxyStream::Unix(s) => Pin::new(s).poll_shutdown(cx),
200        }
201    }
202}
203
204impl Connection for ProxyStream {
205    fn connected(&self) -> Connected {
206        match self {
207            ProxyStream::Tcp(s) => s.connected(),
208            // There is no special connected metadata for Unix sockets
209            #[cfg(unix)]
210            ProxyStream::Unix(_) => Connected::new(),
211        }
212    }
213}
214
215/// Ensure the URI authority includes an explicit port so that hyper emits a
216/// RFC 9110-compliant CONNECT request-target (authority-form requires host:port).
217fn ensure_connect_authority_port(uri: tonic::transport::Uri) -> tonic::transport::Uri {
218    if uri.port().is_some() {
219        return uri;
220    }
221    let port = match uri.scheme_str() {
222        Some("https") => 443,
223        Some("http") => 80,
224        _ => return uri,
225    };
226    let mut parts = uri.into_parts();
227    if let Some(ref authority) = parts.authority
228        && let Ok(new_auth) = format!("{}:{}", authority.host(), port).parse()
229    {
230        parts.authority = Some(new_auth);
231    }
232    tonic::transport::Uri::from_parts(parts).expect("adding port to valid URI should not fail")
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use tokio::{
239        io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
240        net::TcpListener,
241    };
242
243    struct CapturedConnect {
244        request_line: String,
245        headers: Vec<String>,
246    }
247
248    async fn mock_proxy() -> (String, tokio::task::JoinHandle<CapturedConnect>) {
249        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
250        let addr = listener.local_addr().unwrap().to_string();
251        let handle = tokio::spawn(async move {
252            let (stream, _) = listener.accept().await.unwrap();
253            let mut reader = BufReader::new(stream);
254            let mut request_line = String::new();
255            reader.read_line(&mut request_line).await.unwrap();
256            let mut headers = Vec::new();
257            loop {
258                let mut line = String::new();
259                reader.read_line(&mut line).await.unwrap();
260                if line == "\r\n" {
261                    break;
262                }
263                headers.push(line.trim_end().to_string());
264            }
265            reader
266                .into_inner()
267                .write_all(b"HTTP/1.1 200 OK\r\n\r\n")
268                .await
269                .unwrap();
270            CapturedConnect {
271                request_line,
272                headers,
273            }
274        });
275        (addr, handle)
276    }
277
278    #[rstest::rstest]
279    #[case("https://example.com/some/path", "CONNECT example.com:443 HTTP/1.1")]
280    #[case("http://example.com", "CONNECT example.com:80 HTTP/1.1")]
281    #[case("https://example.com:7233", "CONNECT example.com:7233 HTTP/1.1")]
282    #[tokio::test]
283    async fn connect_request_line(#[case] uri: &str, #[case] expected: &str) {
284        let (proxy_addr, handle) = mock_proxy().await;
285        let opts = HttpConnectProxyOptions::new(proxy_addr).build();
286        let uri: tonic::transport::Uri = uri.parse().unwrap();
287        let _ = opts.connect(uri).await;
288
289        let captured = handle.await.unwrap();
290        assert_eq!(captured.request_line.trim(), expected);
291    }
292
293    #[tokio::test]
294    async fn connect_includes_basic_auth() {
295        let (proxy_addr, handle) = mock_proxy().await;
296        let opts = HttpConnectProxyOptions::new(proxy_addr)
297            .basic_auth(("user".to_string(), "pass".to_string()))
298            .build();
299        let uri: tonic::transport::Uri = "https://example.com:7233".parse().unwrap();
300        let _ = opts.connect(uri).await;
301
302        let captured = handle.await.unwrap();
303        let creds = BASE64_STANDARD.encode("user:pass");
304        let auth_header = captured
305            .headers
306            .iter()
307            .find(|h| h.to_lowercase().starts_with("proxy-authorization:"))
308            .expect("missing proxy-authorization header");
309        assert_eq!(
310            auth_header.trim(),
311            format!("proxy-authorization: Basic {creds}")
312        );
313    }
314}
temporalio_client/proxy.rs

temporalio_client/
proxy.rs
