Support for Unix sockets in HTTP connect proxy (#984)

cretz · web-flow · commit eb74c70c2f4b · 2025-08-22T14:19:59.000-05:00
Fixes #981
diff --git a/client/src/lib.rs b/client/src/lib.rs
@@ -9,7 +9,9 @@ extern crate tracing;
 
 pub mod callback_based;
 mod metrics;
-mod proxy;
+/// Visible only for tests
+#[doc(hidden)]
+pub mod proxy;
 mod raw;
 mod retry;
 mod worker_registry;
diff --git a/client/src/proxy.rs b/client/src/proxy.rs
@@ -2,22 +2,33 @@ use base64::prelude::*;
 use http_body_util::Empty;
 use hyper::{body::Bytes, header};
 use hyper_util::{
-    client::legacy::Client,
+    client::legacy::{
+        Client,
+        connect::{Connected, Connection},
+    },
     rt::{TokioExecutor, TokioIo},
 };
 use std::{
     future::Future,
+    io,
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio::net::TcpStream;
+use tokio::{
+    io::{AsyncRead, AsyncWrite, ReadBuf},
+    net::TcpStream,
+};
 use tonic::transport::{Channel, Endpoint};
 use tower::{Service, service_fn};
 
+#[cfg(unix)]
+use tokio::net::UnixStream;
+
 /// Options for HTTP CONNECT proxy.
 #[derive(Clone, Debug)]
 pub struct HttpConnectProxyOptions {
-    /// The host:port to proxy through.
+    /// The host:port to proxy through for TCP, or unix:/path/to/unix.sock for
+    /// Unix socket (which means it must start with "unix:/").
     pub target_addr: String,
     /// Optional HTTP basic auth for the proxy as user/pass tuple.
     pub basic_auth: Option<(String, String)>,
@@ -72,7 +83,7 @@ impl HttpConnectProxyOptions {
 struct OverrideAddrConnector(String);
 
 impl Service<hyper::Uri> for OverrideAddrConnector {
-    type Response = TokioIo<TcpStream>;
+    type Response = TokioIo<ProxyStream>;
 
     type Error = anyhow::Error;
 
@@ -84,7 +95,115 @@ impl Service<hyper::Uri> for OverrideAddrConnector {
 
     fn call(&mut self, _uri: hyper::Uri) -> Self::Future {
         let target_addr = self.0.clone();
-        let fut = async move { Ok(TokioIo::new(TcpStream::connect(target_addr).await?)) };
+        let fut = async move {
+            Ok(TokioIo::new(
+                ProxyStream::connect(target_addr.as_str()).await?,
+            ))
+        };
         Box::pin(fut)
     }
 }
+
+/// Visible only for tests
+#[doc(hidden)]
+pub enum ProxyStream {
+    Tcp(TcpStream),
+    #[cfg(unix)]
+    Unix(UnixStream),
+}
+
+impl ProxyStream {
+    async fn connect(target_addr: &str) -> anyhow::Result<Self> {
+        if target_addr.starts_with("unix:/") {
+            #[cfg(unix)]
+            {
+                Ok(ProxyStream::Unix(
+                    UnixStream::connect(&target_addr[5..]).await?,
+                ))
+            }
+            #[cfg(not(unix))]
+            {
+                Err(anyhow::anyhow!(
+                    "Unix sockets are not supported on this platform"
+                ))
+            }
+        } else {
+            Ok(ProxyStream::Tcp(TcpStream::connect(target_addr).await?))
+        }
+    }
+}
+
+impl AsyncRead for ProxyStream {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            ProxyStream::Tcp(s) => Pin::new(s).poll_read(cx, buf),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => Pin::new(s).poll_read(cx, buf),
+        }
+    }
+}
+
+impl AsyncWrite for ProxyStream {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        match self.get_mut() {
+            ProxyStream::Tcp(s) => Pin::new(s).poll_write(cx, buf),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => Pin::new(s).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        match self.get_mut() {
+            ProxyStream::Tcp(s) => Pin::new(s).poll_write_vectored(cx, bufs),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => Pin::new(s).poll_write_vectored(cx, bufs),
+        }
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        match self {
+            ProxyStream::Tcp(s) => s.is_write_vectored(),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => s.is_write_vectored(),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            ProxyStream::Tcp(s) => Pin::new(s).poll_flush(cx),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => Pin::new(s).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            ProxyStream::Tcp(s) => Pin::new(s).poll_shutdown(cx),
+            #[cfg(unix)]
+            ProxyStream::Unix(s) => Pin::new(s).poll_shutdown(cx),
+        }
+    }
+}
+
+impl Connection for ProxyStream {
+    fn connected(&self) -> Connected {
+        match self {
+            ProxyStream::Tcp(s) => s.connected(),
+            // There is no special connected metadata for Unix sockets
+            #[cfg(unix)]
+            ProxyStream::Unix(_) => Connected::new(),
+        }
+    }
+}
diff --git a/test-utils/Cargo.toml b/test-utils/Cargo.toml
@@ -17,7 +17,11 @@ ephemeral-server = ["temporal-sdk-core/ephemeral-server"]
 anyhow = "1.0"
 assert_matches = "1"
 async-trait = "0.1"
+bytes = "1.10"
 futures-util = { version = "0.3", default-features = false }
+hyper = { version = "1.4.1" }
+http-body-util = "0.1"
+hyper-util = "0.1.6"
 parking_lot = "0.12"
 prost = { workspace = true }
 rand = "0.9"
diff --git a/test-utils/src/http_proxy.rs b/test-utils/src/http_proxy.rs
@@ -0,0 +1,135 @@
+use std::{
+    io,
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+};
+
+use bytes::Bytes;
+use http_body_util::Empty;
+use hyper::{
+    Request, Response, StatusCode, body::Incoming, server::conn::http1, service::service_fn,
+};
+use hyper_util::rt::TokioIo;
+use temporal_client::proxy::ProxyStream;
+#[cfg(unix)]
+use tokio::net::UnixListener;
+use tokio::{
+    net::{TcpListener, TcpStream},
+    sync::oneshot,
+};
+
+pub struct HttpProxy {
+    proxy_hits: Arc<AtomicUsize>,
+    shutdown_tx: oneshot::Sender<()>,
+}
+impl HttpProxy {
+    pub fn spawn_tcp(listener: TcpListener) -> Self {
+        Self::spawn(ProxyListener::Tcp(listener))
+    }
+
+    #[cfg(unix)]
+    pub fn spawn_unix(listener: UnixListener) -> Self {
+        Self::spawn(ProxyListener::Unix(listener))
+    }
+
+    fn spawn(listener: ProxyListener) -> Self {
+        let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>();
+        let proxy_hits = Arc::new(AtomicUsize::new(0));
+        let proxy_hits_cloned = proxy_hits.clone();
+        tokio::spawn(async move {
+            loop {
+                let proxy_hits_cloned = proxy_hits_cloned.clone();
+                tokio::select! {
+                    _ = &mut shutdown_rx => break,
+                    stream = listener.accept() => {
+                        let stream = match stream {
+                            Ok(stream) => stream,
+                            Err(e) => { println!("Proxy accept error: {e}"); continue; }
+                        };
+                        tokio::spawn(async move {
+                            if let Err(e) = http1::Builder::new()
+                                .serve_connection(
+                                    TokioIo::new(stream),
+                                    service_fn(move |req| handle_connect(req, proxy_hits_cloned.clone())),
+                                )
+                                .with_upgrades()
+                                .await
+                            {
+                                println!("Proxy conn error: {e}");
+                            }
+                        });
+                    }
+                }
+            }
+        });
+        Self {
+            proxy_hits,
+            shutdown_tx,
+        }
+    }
+
+    pub fn hit_count(&self) -> usize {
+        self.proxy_hits.load(Ordering::SeqCst)
+    }
+
+    /// Returns before shutdown occurs
+    pub fn shutdown(self) {
+        let _ = self.shutdown_tx.send(());
+    }
+}
+
+async fn handle_connect(
+    req: Request<Incoming>,
+    counter: Arc<AtomicUsize>,
+) -> Result<Response<Empty<Bytes>>, hyper::Error> {
+    if req.method() == hyper::Method::CONNECT {
+        // Increment atomic counter
+        counter.fetch_add(1, Ordering::SeqCst);
+
+        // Tell the client the tunnel is established
+        tokio::spawn(async move {
+            if let Some(addr) = req.uri().authority().map(|a| a.as_str()) {
+                match TcpStream::connect(addr).await {
+                    Ok(mut server_stream) => match hyper::upgrade::on(req).await {
+                        Ok(upgraded) => {
+                            let mut upgraded = TokioIo::new(upgraded);
+                            let _ =
+                                tokio::io::copy_bidirectional(&mut upgraded, &mut server_stream)
+                                    .await;
+                        }
+                        Err(err) => println!("Upgrade failed: {err}"),
+                    },
+                    Err(e) => println!("Failed to connect to {addr}: {e}"),
+                }
+            }
+        });
+
+        Ok(Response::builder()
+            .status(StatusCode::OK)
+            .body(Empty::new())
+            .unwrap())
+    } else {
+        Ok(Response::builder()
+            .status(StatusCode::METHOD_NOT_ALLOWED)
+            .body(Empty::new())
+            .unwrap())
+    }
+}
+
+enum ProxyListener {
+    Tcp(TcpListener),
+    #[cfg(unix)]
+    Unix(UnixListener),
+}
+
+impl ProxyListener {
+    async fn accept(&self) -> io::Result<ProxyStream> {
+        match self {
+            ProxyListener::Tcp(tcp) => tcp.accept().await.map(|(s, _)| ProxyStream::Tcp(s)),
+            #[cfg(unix)]
+            ProxyListener::Unix(unix) => unix.accept().await.map(|(s, _)| ProxyStream::Unix(s)),
+        }
+    }
+}
diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs
@@ -5,9 +5,11 @@
 extern crate tracing;
 
 pub mod canned_histories;
+mod http_proxy;
 pub mod interceptors;
 pub mod workflows;
 
+pub use http_proxy::HttpProxy;
 pub use temporal_sdk_core::replay::HistoryForReplay;
 
 use crate::stream::{Stream, TryStreamExt};
diff --git a/tests/integ_tests/client_tests.rs b/tests/integ_tests/client_tests.rs
