Implement forward proxy with OpenTelemetry observability

src/proxy.rs

@@ -0,0 +1,680 @@
+use std::{convert::Infallible, net::SocketAddr, time::Instant};
+
+use bytes::Bytes;
+use http::{HeaderValue, Method, Request, Response, StatusCode, Uri, header::HOST, uri::Authority};
+use http_body_util::{BodyExt, Empty, Full, combinators::UnsyncBoxBody};
+use hyper::{body::Incoming, server::conn::http1, service::service_fn};
+use hyper_util::{
+    client::legacy::{Client, connect::HttpConnector},
+    rt::{TokioExecutor, TokioIo},
+};
+use opentelemetry::global;
+use tokio::{
+    net::{TcpListener, TcpStream},
+    task::JoinSet,
+    time::timeout,
+};
+use tracing::{Instrument, Span, debug, error, field, info, info_span};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+use crate::{
+    BoxError,
+    config::Config,
+    metrics::Metrics,
+    telemetry::{HeaderExtractor, HeaderInjector, remove_hop_by_hop_headers},
+};
+
+type ProxyBody = UnsyncBoxBody<Bytes, BoxError>;
+type ProxyResponse = Response<ProxyBody>;
+type HttpClient = Client<HttpConnector, Incoming>;
+
+#[derive(Clone)]
+struct ProxyState {
+    client: HttpClient,
+    metrics: Metrics,
+    connect_timeout: std::time::Duration,
+    upstream_timeout: std::time::Duration,
+}
+
+impl ProxyState {
+    fn new(config: &Config) -> Self {
+        let mut connector = HttpConnector::new();
+        connector.enforce_http(true);
+        connector.set_connect_timeout(Some(config.connect_timeout));
+
+        Self {
+            client: Client::builder(TokioExecutor::new()).build(connector),
+            metrics: Metrics::new(global::meter("hermes")),
+            connect_timeout: config.connect_timeout,
+            upstream_timeout: config.upstream_timeout,
+        }
+    }
+}
+
+pub async fn run(config: Config) -> Result<(), BoxError> {
+    let listener = TcpListener::bind(config.listen_addr).await?;
+    let state = ProxyState::new(&config);
+    let mut connections = JoinSet::new();
+    let shutdown = tokio::signal::ctrl_c();
+    tokio::pin!(shutdown);
+
+    info!(listen.address = %config.listen_addr, "proxy listening");
+
+    loop {
+        tokio::select! {
+            _ = &mut shutdown => {
+                info!("shutdown signal received");
+                break;
+            }
+            accepted = listener.accept() => {
+                match accepted {
+                    Ok((stream, peer)) => {
+                        let state = state.clone();
+                        connections.spawn(async move {
+                            serve_connection(stream, peer, state).await;
+                        });
+                    }
+                    Err(error) => {
+                        error!(error.message = %error, "failed to accept downstream connection");
+                    }
+                }
+            }
+            completed = connections.join_next(), if !connections.is_empty() => {
+                if let Some(Err(error)) = completed {
+                    error!(error.message = %error, "connection task failed");
+                }
+            }
+        }
+    }
+
+    info!("proxy listener stopped");
+    Ok(())
+}
+
+async fn serve_connection(stream: TcpStream, peer: SocketAddr, state: ProxyState) {
+    let service = service_fn(move |request| {
+        let state = state.clone();
+        async move { Ok::<_, Infallible>(handle_request(request, state).await) }
+    });
+
+    let connection = http1::Builder::new()
+        .serve_connection(TokioIo::new(stream), service)
+        .with_upgrades();
+
+    if let Err(error) = connection.await {
+        debug!(
+            client.address = %peer,
+            error.message = %error,
+            "downstream connection ended with an error"
+        );
+    }
+}
+
+async fn handle_request(request: Request<Incoming>, state: ProxyState) -> ProxyResponse {
+    if request.method() == Method::CONNECT {
+        handle_connect(request, state).await
+    } else {
+        handle_http(request, state).await
+    }
+}
+
+async fn handle_http(request: Request<Incoming>, state: ProxyState) -> ProxyResponse {
+    let method = request.method().to_string();
+    let uri = request.uri().clone();
+    let parent = global::get_text_map_propagator(|propagator| {
+        propagator.extract(&HeaderExtractor(request.headers()))
+    });
+    let span = info_span!(
+        "proxy.request",
+        otel.name = %format!("HTTP {method}"),
+        otel.kind = "server",
+        http.request.method = %method,
+        url.scheme = %uri.scheme_str().unwrap_or_default(),
+        url.path = %uri.path(),
+        server.address = %uri.authority().map(Authority::host).unwrap_or_default(),
+        http.response.status_code = field::Empty,
+        error.type = field::Empty,
+        otel.status_code = field::Empty,
+    );
+    let _ = span.set_parent(parent);
+
+    async move {
+        let _active = state.metrics.request_started();
+        let started = Instant::now();
+        let result = forward_http(request, &state).await;
+        finish_request(result, &method, started, &state.metrics)
+    }
+    .instrument(span)
+    .await
+}
+
+async fn forward_http(
+    mut request: Request<Incoming>,
+    state: &ProxyState,
+) -> Result<ProxyResponse, ProxyError> {
+    let authority = validate_http_target(request.uri())?;
+    let host = HeaderValue::from_str(authority.as_str())
+        .map_err(|_| ProxyError::bad_request("invalid_target", "invalid upstream Host header"))?;
+
+    remove_hop_by_hop_headers(request.headers_mut());
+    request.headers_mut().insert(HOST, host);
+
+    let response = send_upstream(request, state).await?;
+    let (parts, body) = response.into_parts();
+    let mut response = Response::from_parts(parts, box_incoming_body(body));
+    remove_hop_by_hop_headers(response.headers_mut());
+    Ok(response)
+}
+
+async fn send_upstream(
+    mut request: Request<Incoming>,
+    state: &ProxyState,
+) -> Result<Response<Incoming>, ProxyError> {
+    let method = request.method().to_string();
+    let host = request
+        .uri()
+        .authority()
+        .map(Authority::host)
+        .unwrap_or_default()
+        .to_owned();
+    let port = request.uri().port_u16().unwrap_or(80);
+    let span = info_span!(
+        "proxy.upstream",
+        otel.name = %format!("HTTP {method}"),
+        otel.kind = "client",
+        http.request.method = %method,
+        server.address = %host,
+        server.port = port,
+        http.response.status_code = field::Empty,
+        error.type = field::Empty,
+        otel.status_code = field::Empty,
+    );
+
+    async move {
+        let context = Span::current().context();
+        global::get_text_map_propagator(|propagator| {
+            propagator.inject_context(&context, &mut HeaderInjector(request.headers_mut()));
+        });
+
+        let started = Instant::now();
+        let result = match timeout(state.upstream_timeout, state.client.request(request)).await {
+            Ok(Ok(response)) => Ok(response),
+            Ok(Err(error)) => Err(ProxyError::new(
+                StatusCode::BAD_GATEWAY,
+                "upstream_request",
+                "upstream",
+                error.to_string(),
+            )),
+            Err(_) => Err(ProxyError::new(
+                StatusCode::GATEWAY_TIMEOUT,
+                "upstream_timeout",
+                "upstream",
+                "upstream request timed out".to_owned(),
+            )),
+        };
+
+        let (status, error_type) = match &result {
+            Ok(response) => {
+                let status = response.status();
+                Span::current().record("http.response.status_code", status.as_u16());
+                if status.is_server_error() {
+                    Span::current().record("otel.status_code", "ERROR");
+                }
+                (Some(status.as_u16()), None)
+            }
+            Err(error) => {
+                Span::current().record("error.type", error.error_type);
+                Span::current().record("otel.status_code", "ERROR");
+                (None, Some(error.error_type))
+            }
+        };
+        state
+            .metrics
+            .upstream_finished(&method, status, error_type, started.elapsed());
+
+        result
+    }
+    .instrument(span)
+    .await
+}
+
+async fn handle_connect(request: Request<Incoming>, state: ProxyState) -> ProxyResponse {
+    let authority = request.uri().authority().cloned();
+    let parent = global::get_text_map_propagator(|propagator| {
+        propagator.extract(&HeaderExtractor(request.headers()))
+    });
+    let span = info_span!(
+        "proxy.connect",
+        otel.name = "CONNECT",
+        otel.kind = "server",
+        http.request.method = "CONNECT",
+        server.address = %authority.as_ref().map(Authority::host).unwrap_or_default(),
+        server.port = authority.as_ref().and_then(Authority::port_u16).unwrap_or_default(),
+        http.response.status_code = field::Empty,
+        error.type = field::Empty,
+        tunnel.bytes.client_to_upstream = field::Empty,
+        tunnel.bytes.upstream_to_client = field::Empty,
+        otel.status_code = field::Empty,
+    );
+    let _ = span.set_parent(parent);
+
+    async move {
+        let _active = state.metrics.request_started();
+        let started = Instant::now();
+        let result = prepare_connect(request, &state).await;
+        finish_request(result, "CONNECT", started, &state.metrics)
+    }
+    .instrument(span)
+    .await
+}
+
+async fn prepare_connect(
+    request: Request<Incoming>,
+    state: &ProxyState,
+) -> Result<ProxyResponse, ProxyError> {
+    let authority = validate_connect_target(request.uri())?;
+    let host = authority.host().to_owned();
+    let port = authority.port_u16().unwrap_or_default();
+    let connect_span = info_span!(
+        "proxy.upstream_connect",
+        otel.name = "TCP CONNECT",
+        otel.kind = "client",
+        server.address = %host,
+        server.port = port,
+        error.type = field::Empty,
+        otel.status_code = field::Empty,
+    );
+
+    let upstream = async {
+        let result = match timeout(
+            state.connect_timeout,
+            TcpStream::connect(authority.as_str()),
+        )
+        .await
+        {
+            Ok(Ok(stream)) => Ok(stream),
+            Ok(Err(error)) => Err(ProxyError::new(
+                StatusCode::BAD_GATEWAY,
+                "upstream_connect",
+                "upstream",
+                error.to_string(),
+            )),
+            Err(_) => Err(ProxyError::new(
+                StatusCode::GATEWAY_TIMEOUT,
+                "upstream_timeout",
+                "upstream",
+                "upstream connection timed out".to_owned(),
+            )),
+        };
+
+        if let Err(error) = &result {
+            Span::current().record("error.type", error.error_type);
+            Span::current().record("otel.status_code", "ERROR");
+        }
+        result
+    }
+    .instrument(connect_span)
+    .await?;
+
+    let tunnel_span = Span::current();
+    let metrics = state.metrics.clone();
+    tokio::spawn(run_tunnel(request, upstream, metrics).instrument(tunnel_span));
+
+    Ok(Response::new(empty_body()))
+}
+
+async fn run_tunnel(request: Request<Incoming>, mut upstream: TcpStream, metrics: Metrics) {
+    let upgraded = match hyper::upgrade::on(request).await {
+        Ok(upgraded) => upgraded,
+        Err(error) => {
+            Span::current().record("error.type", "downstream_upgrade");
+            Span::current().record("otel.status_code", "ERROR");
+            metrics.error("downstream", "downstream_upgrade");
+            error!(error.message = %error, "failed to upgrade CONNECT request");
+            return;
+        }
+    };
+
+    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");
+        }
+    }
+}
+
+fn finish_request(
+    result: Result<ProxyResponse, ProxyError>,
+    method: &str,
+    started: Instant,
+    metrics: &Metrics,
+) -> ProxyResponse {
+    let mut error_type = None;
+    let response = match result {
+        Ok(response) => response,
+        Err(error) => {
+            error_type = Some(error.error_type);
+            Span::current().record("error.type", error.error_type);
+            Span::current().record("otel.status_code", "ERROR");
+            metrics.error(error.phase, error.error_type);
+            error!(
+                error.type = error.error_type,
+                proxy.phase = error.phase,
+                error.message = %error.message,
+                "proxy request failed"
+            );
+            error_response(error.status)
+        }
+    };
+
+    let status = response.status();
+    Span::current().record("http.response.status_code", status.as_u16());
+    if status.is_server_error() {
+        Span::current().record("otel.status_code", "ERROR");
+    }
+    let duration = started.elapsed();
+    metrics.request_finished(method, status.as_u16(), error_type, duration);
+    info!(
+        http.request.method = method,
+        http.response.status_code = status.as_u16(),
+        request.duration_ms = duration.as_millis(),
+        "proxy request completed"
+    );
+    response
+}
+
+fn validate_http_target(uri: &Uri) -> Result<Authority, ProxyError> {
+    if uri.scheme_str() != Some("http") {
+        return Err(ProxyError::bad_request(
+            "unsupported_scheme",
+            "only absolute http:// request targets are supported",
+        ));
+    }
+
+    let authority = uri.authority().ok_or_else(|| {
+        ProxyError::bad_request("invalid_target", "absolute request target has no authority")
+    })?;
+    if authority.as_str().contains('@') {
+        return Err(ProxyError::bad_request(
+            "invalid_target",
+            "userinfo is not allowed in request targets",
+        ));
+    }
+
+    Ok(authority.clone())
+}
+
+fn validate_connect_target(uri: &Uri) -> Result<Authority, ProxyError> {
+    if uri.scheme().is_some() || !uri.path().is_empty() {
+        return Err(ProxyError::bad_request(
+            "invalid_target",
+            "CONNECT target must use authority-form",
+        ));
+    }
+
+    let authority = uri.authority().ok_or_else(|| {
+        ProxyError::bad_request("invalid_target", "CONNECT target has no authority")
+    })?;
+    if authority.port_u16().is_none() || authority.as_str().contains('@') {
+        return Err(ProxyError::bad_request(
+            "invalid_target",
+            "CONNECT target must contain a valid host and explicit port",
+        ));
+    }
+
+    Ok(authority.clone())
+}
+
+fn error_response(status: StatusCode) -> ProxyResponse {
+    let mut response = Response::new(full_body(format!(
+        "{}\n",
+        status.canonical_reason().unwrap_or("Proxy Error")
+    )));
+    *response.status_mut() = status;
+    response
+}
+
+fn box_incoming_body(body: Incoming) -> ProxyBody {
+    body.map_err(|error| -> BoxError { Box::new(error) })
+        .boxed_unsync()
+}
+
+fn full_body(value: impl Into<Bytes>) -> ProxyBody {
+    Full::new(value.into())
+        .map_err(|never| match never {})
+        .boxed_unsync()
+}
+
+fn empty_body() -> ProxyBody {
+    Empty::new().map_err(|never| match never {}).boxed_unsync()
+}
+
+#[derive(Debug)]
+struct ProxyError {
+    status: StatusCode,
+    error_type: &'static str,
+    phase: &'static str,
+    message: String,
+}
+
+impl ProxyError {
+    fn new(
+        status: StatusCode,
+        error_type: &'static str,
+        phase: &'static str,
+        message: String,
+    ) -> Self {
+        Self {
+            status,
+            error_type,
+            phase,
+            message,
+        }
+    }
+
+    fn bad_request(error_type: &'static str, message: &'static str) -> Self {
+        Self::new(
+            StatusCode::BAD_REQUEST,
+            error_type,
+            "request_validation",
+            message.to_owned(),
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{
+        sync::{Arc, Mutex},
+        time::Duration,
+    };
+
+    use hyper::client::conn::http1 as client_http1;
+    use tokio::{
+        io::{AsyncReadExt, AsyncWriteExt},
+        task::JoinHandle,
+    };
+
+    use super::*;
+
+    #[test]
+    fn validates_absolute_http_targets() {
+        let uri = Uri::from_static("http://example.com/path?q=1");
+        assert_eq!(validate_http_target(&uri).unwrap(), "example.com");
+
+        assert!(validate_http_target(&Uri::from_static("/path")).is_err());
+        assert!(validate_http_target(&Uri::from_static("https://example.com/")).is_err());
+        assert!(validate_http_target(&Uri::from_static("http://user@example.com/")).is_err());
+    }
+
+    #[test]
+    fn validates_connect_authority() {
+        let uri = Uri::from_static("example.com:443");
+        assert_eq!(validate_connect_target(&uri).unwrap(), "example.com:443");
+
+        assert!(validate_connect_target(&Uri::from_static("example.com")).is_err());
+        assert!(validate_connect_target(&Uri::from_static("http://example.com:443")).is_err());
+    }
+
+    #[tokio::test]
+    async fn forwards_http_and_rewrites_connection_headers() {
+        #[derive(Default)]
+        struct Observed {
+            target: String,
+            host: String,
+            forwarded_proxy_authorization: bool,
+        }
+
+        let observed = Arc::new(Mutex::new(Observed::default()));
+        let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let upstream_addr = upstream_listener.local_addr().unwrap();
+        let upstream_observed = observed.clone();
+        let upstream_task = tokio::spawn(async move {
+            let (stream, _) = upstream_listener.accept().await.unwrap();
+            let service = service_fn(move |request: Request<Incoming>| {
+                let mut observed = upstream_observed.lock().unwrap();
+                observed.target = request.uri().to_string();
+                observed.host = request
+                    .headers()
+                    .get(HOST)
+                    .and_then(|value| value.to_str().ok())
+                    .unwrap_or_default()
+                    .to_owned();
+                observed.forwarded_proxy_authorization =
+                    request.headers().contains_key("proxy-authorization");
+                drop(observed);
+
+                async move {
+                    Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"proxied"))))
+                }
+            });
+            http1::Builder::new()
+                .serve_connection(TokioIo::new(stream), service)
+                .await
+                .unwrap();
+        });
+
+        let (proxy_addr, proxy_task) = spawn_proxy().await;
+        let (mut sender, connection_task) = connect_client(proxy_addr).await;
+        let request = Request::builder()
+            .method(Method::GET)
+            .uri(format!("http://{upstream_addr}/hello?q=1"))
+            .header(HOST, "wrong.invalid")
+            .header("proxy-authorization", "must-not-be-forwarded")
+            .body(Full::new(Bytes::new()))
+            .unwrap();
+
+        let response = sender.send_request(request).await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+        let body = response.into_body().collect().await.unwrap().to_bytes();
+        assert_eq!(body, Bytes::from_static(b"proxied"));
+
+        upstream_task.abort();
+        let observed = observed.lock().unwrap();
+        assert_eq!(observed.target, "/hello?q=1");
+        assert_eq!(observed.host, upstream_addr.to_string());
+        assert!(!observed.forwarded_proxy_authorization);
+        drop(observed);
+
+        proxy_task.abort();
+        connection_task.abort();
+    }
+
+    #[tokio::test]
+    async fn relays_bytes_through_connect_tunnel() {
+        let echo_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let echo_addr = echo_listener.local_addr().unwrap();
+        let echo_task = tokio::spawn(async move {
+            let (mut stream, _) = echo_listener.accept().await.unwrap();
+            let mut request = [0_u8; 4];
+            stream.read_exact(&mut request).await.unwrap();
+            assert_eq!(&request, b"ping");
+            stream.write_all(b"pong").await.unwrap();
+            stream.shutdown().await.unwrap();
+        });
+
+        let (proxy_addr, proxy_task) = spawn_proxy().await;
+        let (mut sender, connection_task) = connect_client(proxy_addr).await;
+        let request = Request::builder()
+            .method(Method::CONNECT)
+            .uri(echo_addr.to_string())
+            .body(Full::new(Bytes::new()))
+            .unwrap();
+
+        let response = sender.send_request(request).await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let upgraded = hyper::upgrade::on(response).await.unwrap();
+        let mut tunnel = TokioIo::new(upgraded);
+        tunnel.write_all(b"ping").await.unwrap();
+        let mut response = [0_u8; 4];
+        tunnel.read_exact(&mut response).await.unwrap();
+        assert_eq!(&response, b"pong");
+
+        echo_task.await.unwrap();
+        proxy_task.abort();
+        connection_task.abort();
+    }
+
+    fn test_state() -> ProxyState {
+        ProxyState::new(&Config {
+            listen_addr: "127.0.0.1:0".parse().unwrap(),
+            connect_timeout: Duration::from_secs(2),
+            upstream_timeout: Duration::from_secs(2),
+        })
+    }
+
+    async fn spawn_proxy() -> (SocketAddr, JoinHandle<()>) {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let state = test_state();
+        let task = tokio::spawn(async move {
+            while let Ok((stream, peer)) = listener.accept().await {
+                let state = state.clone();
+                tokio::spawn(serve_connection(stream, peer, state));
+            }
+        });
+        (address, task)
+    }
+
+    async fn connect_client(
+        address: SocketAddr,
+    ) -> (
+        client_http1::SendRequest<Full<Bytes>>,
+        JoinHandle<Result<(), hyper::Error>>,
+    ) {
+        let stream = TcpStream::connect(address).await.unwrap();
+        let (sender, connection) = client_http1::handshake(TokioIo::new(stream)).await.unwrap();
+        let connection_task = tokio::spawn(connection.with_upgrades());
+        (sender, connection_task)
+    }
+}