HTTP CONNECT proxy support (#714)

Fixes #309

Author: cretz

client/Cargo.toml

@@ -17,11 +17,13 @@ telemetry = ["dep:opentelemetry"]
 anyhow = "1.0"
 async-trait = "0.1"
 backoff = "0.4"
+base64 = "0.21.7"
 derive_builder = { workspace = true }
 derive_more = "0.99"
 futures = "0.3"
 futures-retry = "0.6.0"
 http = "0.2"
+hyper = { version = "0.14.28" }
 once_cell = { workspace = true }
 opentelemetry = { workspace = true, features = ["metrics"], optional = true  }
 parking_lot = "0.12"

client/src/lib.rs

@@ -8,11 +8,13 @@
 extern crate tracing;
 
 mod metrics;
+mod proxy;
 mod raw;
 mod retry;
 mod worker_registry;
 mod workflow_handle;
 
+pub use crate::proxy::HttpConnectProxyOptions;
 pub use crate::retry::{CallType, RetryClient, RETRYABLE_ERROR_CODES};
 pub use raw::{HealthService, OperatorService, TestService, WorkflowService};
 pub use temporal_sdk_core_protos::temporal::api::{
@@ -142,6 +144,10 @@ pub struct ClientOptions {
     /// be applied if the headers don't already have an "Authorization" header.
     #[builder(default)]
     pub api_key: Option<String>,
+
+    /// HTTP CONNECT proxy to use for this client.
+    #[builder(default)]
+    pub http_connect_proxy: Option<HttpConnectProxyOptions>,
 }
 
 /// Configuration options for TLS
@@ -403,7 +409,12 @@ impl ClientOptions {
         } else {
             channel
         };
-        let channel = channel.connect().await?;
+        // If there is a proxy, we have to connect that way
+        let channel = if let Some(proxy) = self.http_connect_proxy.as_ref() {
+            proxy.connect_endpoint(&channel).await?
+        } else {
+            channel.connect().await?
+        };
         let service = ServiceBuilder::new()
             .layer_fn(move |channel| GrpcMetricSvc {
                 inner: channel,

client/src/proxy.rs

@@ -0,0 +1,85 @@
+use base64::prelude::*;
+use hyper::header;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+use tokio::net::TcpStream;
+use tonic::transport::Channel;
+use tonic::transport::Endpoint;
+use tower::{service_fn, Service};
+
+/// Options for HTTP CONNECT proxy.
+#[derive(Clone, Debug)]
+pub struct HttpConnectProxyOptions {
+    /// The host:port to proxy through.
+    pub target_addr: String,
+    /// Optional HTTP basic auth for the proxy as user/pass tuple.
+    pub basic_auth: Option<(String, String)>,
+}
+
+impl HttpConnectProxyOptions {
+    /// Create a channel from the given endpoint that uses the HTTP CONNECT proxy.
+    pub async fn connect_endpoint(
+        &self,
+        endpoint: &Endpoint,
+    ) -> Result<Channel, tonic::transport::Error> {
+        let proxy_options = self.clone();
+        let svc_fn = service_fn(move |uri: tonic::transport::Uri| {
+            let proxy_options = proxy_options.clone();
+            async move { proxy_options.connect(uri).await }
+        });
+        endpoint.connect_with_connector(svc_fn).await
+    }
+
+    async fn connect(
+        &self,
+        uri: tonic::transport::Uri,
+    ) -> anyhow::Result<hyper::upgrade::Upgraded> {
+        debug!("Connecting to {} via proxy at {}", uri, self.target_addr);
+        // Create CONNECT request
+        let mut req_build = hyper::Request::builder().method("CONNECT").uri(uri);
+        if let Some((user, pass)) = &self.basic_auth {
+            let creds = BASE64_STANDARD.encode(format!("{}:{}", user, pass));
+            req_build = req_build.header(header::PROXY_AUTHORIZATION, format!("Basic {}", creds));
+        }
+        let req = req_build.body(hyper::Body::empty())?;
+
+        // We have to create a client with a specific connector because Hyper is
+        // not letting us change the HTTP/2 authority
+        let client =
+            hyper::Client::builder().build(OverrideAddrConnector(self.target_addr.clone()));
+
+        // Send request
+        let res = client.request(req).await?;
+        if res.status().is_success() {
+            Ok(hyper::upgrade::on(res).await?)
+        } else {
+            Err(anyhow::anyhow!(
+                "CONNECT call failed with status: {}",
+                res.status()
+            ))
+        }
+    }
+}
+
+#[derive(Clone)]
+struct OverrideAddrConnector(String);
+
+impl Service<hyper::Uri> for OverrideAddrConnector {
+    type Response = TcpStream;
+
+    type Error = anyhow::Error;
+
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+    fn poll_ready(&mut self, _ctx: &mut Context<'_>) -> Poll<anyhow::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, _uri: hyper::Uri) -> Self::Future {
+        let target_addr = self.0.clone();
+        let fut = async move { Ok(TcpStream::connect(target_addr).await?) };
+        Box::pin(fut)
+    }
+}