Add SOCKS5 CONNECT support

src/proxy.rs

@@ -20,7 +20,7 @@ use hyper_util::{
 use opentelemetry::global;
 use tokio::{
     net::{TcpListener, TcpStream},
-    task::JoinSet,
+    task::{JoinHandle, JoinSet},
     time::timeout,
 };
 use tracing::{Instrument, Span, debug, error, field, info, info_span, warn};
@@ -31,7 +31,9 @@ use crate::{
     acl::{Acl, Decision, DestinationHost},
     config::Config,
     metrics::Metrics,
+    socks5::{self, SocksState},
     telemetry::{HeaderExtractor, HeaderInjector, remove_hop_by_hop_headers},
+    tunnel::{self, TunnelProtocol},
 };
 
 type ProxyBody = UnsyncBoxBody<Bytes, BoxError>;
@@ -66,6 +68,14 @@ impl ProxyState {
 pub async fn run(config: Config) -> Result<(), BoxError> {
     let listener = TcpListener::bind(config.listen_addr).await?;
     let state = ProxyState::new(&config);
+    let mut socks_task = match config.socks5_listen_addr {
+        Some(address) => {
+            let listener = TcpListener::bind(address).await?;
+            let state = SocksState::new(&config);
+            Some(tokio::spawn(socks5::run(listener, state)))
+        }
+        None => None,
+    };
     let mut connections = JoinSet::new();
     let shutdown = tokio::signal::ctrl_c();
     tokio::pin!(shutdown);
@@ -96,13 +106,33 @@ pub async fn run(config: Config) -> Result<(), BoxError> {
                     error!(error.message = %error, "connection task failed");
                 }
             }
+            completed = wait_for_socks(&mut socks_task) => {
+                match completed {
+                    Ok(Ok(())) => error!("SOCKS5 listener stopped unexpectedly"),
+                    Ok(Err(error)) => error!(error.message = %error, "SOCKS5 listener failed"),
+                    Err(error) => error!(error.message = %error, "SOCKS5 listener task failed"),
+                }
+                break;
+            }
         }
     }
 
+    if let Some(task) = socks_task {
+        task.abort();
+    }
     info!("proxy listener stopped");
     Ok(())
 }
 
+async fn wait_for_socks(
+    task: &mut Option<JoinHandle<Result<(), BoxError>>>,
+) -> Result<Result<(), BoxError>, tokio::task::JoinError> {
+    match task {
+        Some(task) => task.await,
+        None => std::future::pending().await,
+    }
+}
+
 async fn serve_connection(stream: TcpStream, peer: SocketAddr, state: ProxyState) {
     let peer_ip = peer.ip();
     let service = service_fn(move |request| {
@@ -188,9 +218,14 @@ fn authorize_request(
         .acl
         .authorize(peer, request.headers(), destination.as_ref());
 
+    let protocol = if request.method() == Method::CONNECT {
+        "http_connect"
+    } else {
+        "http"
+    };
     state
         .metrics
-        .acl_decision(outcome.decision.as_str(), outcome.reason);
+        .acl_decision(protocol, outcome.decision.as_str(), outcome.reason);
     Span::current().record("acl.decision", outcome.decision.as_str());
     Span::current().record("acl.reason", outcome.reason);
     if let Some(rule) = &outcome.rule {
@@ -424,41 +459,14 @@ async fn run_tunnel(request: Request<Incoming>, mut upstream: TcpStream, metrics
         }
     };
 
-    let started = Instant::now();
-    let _active = metrics.tunnel_started();
     let mut downstream = TokioIo::new(upgraded);
-
-    match tokio::io::copy_bidirectional(&mut downstream, &mut upstream).await {
-        Ok((client_to_upstream, upstream_to_client)) => {
-            Span::current().record(
-                "tunnel.bytes.client_to_upstream",
-                i64::try_from(client_to_upstream).unwrap_or(i64::MAX),
-            );
-            Span::current().record(
-                "tunnel.bytes.upstream_to_client",
-                i64::try_from(upstream_to_client).unwrap_or(i64::MAX),
-            );
-            metrics.tunnel_finished(
-                started.elapsed(),
-                client_to_upstream,
-                upstream_to_client,
-                None,
-            );
-            info!(
-                tunnel.bytes.client_to_upstream = client_to_upstream,
-                tunnel.bytes.upstream_to_client = upstream_to_client,
-                tunnel.duration_ms = started.elapsed().as_millis(),
-                "CONNECT tunnel closed"
-            );
-        }
-        Err(error) => {
-            Span::current().record("error.type", "tunnel_io");
-            Span::current().record("otel.status_code", "ERROR");
-            metrics.error("tunnel", "tunnel_io");
-            metrics.tunnel_finished(started.elapsed(), 0, 0, Some("tunnel_io"));
-            error!(error.message = %error, "CONNECT tunnel failed");
-        }
-    }
+    tunnel::relay(
+        &mut downstream,
+        &mut upstream,
+        &metrics,
+        TunnelProtocol::HttpConnect,
+    )
+    .await;
 }
 
 fn finish_request(
@@ -869,7 +877,9 @@ mod tests {
         let address = listener.local_addr().unwrap();
         let state = ProxyState::new(&Config {
             listen_addr: address,
+            socks5_listen_addr: None,
             connect_timeout: Duration::from_secs(2),
+            socks5_handshake_timeout: Duration::from_secs(2),
             upstream_timeout: Duration::from_secs(2),
             telemetry: TelemetryConfig::default(),
             acl: Arc::new(acl),