fix(stackable-telemetry/webhook): Correctly extract connect and host info in Axum trace layer (#806)

* Add changelog entry

* Update PR link in changelog

* Add support for connect info

* Add support for host info

* Fix numeric fields, include host address and port

* Revert Cargo.toml changes from b7c6118

* Add port and scheme to error message

* Fix typo

* Add quotes around field keys

* chore: remove lifetimes, leverage Snafu's generated Into parameter

---------

Co-authored-by: Nick Larsen <[email protected]>

Author: Techassi

crates/stackable-telemetry/src/instrumentation/axum/mod.rs

@@ -10,16 +10,20 @@
 //!
 //! [1]: https://opentelemetry.io/
 //! [2]: https://opentelemetry.io/docs/specs/semconv/http/http-spans/
-use std::{future::Future, net::SocketAddr, str::FromStr, task::Poll};
+use std::{future::Future, net::SocketAddr, num::ParseIntError, task::Poll};
 
 use axum::{
-    extract::{ConnectInfo, Host, MatchedPath, Request},
-    http::{header::USER_AGENT, HeaderMap},
+    extract::{ConnectInfo, MatchedPath, Request},
+    http::{
+        header::{HOST, USER_AGENT},
+        HeaderMap,
+    },
     response::Response,
 };
 use futures_util::ready;
 use opentelemetry::trace::SpanKind;
 use pin_project::pin_project;
+use snafu::{ResultExt, Snafu};
 use tower::{Layer, Service};
 use tracing::{debug, field::Empty, instrument, trace_span, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
@@ -30,6 +34,10 @@ mod injector;
 pub use extractor::*;
 pub use injector::*;
 
+const X_FORWARDED_HOST_HEADER_KEY: &str = "X-Forwarded-Host";
+const DEFAULT_HTTPS_PORT: u16 = 443;
+const DEFAULT_HTTP_PORT: u16 = 80;
+
 /// A Tower [`Layer`][1] which decorates [`TraceService`].
 ///
 /// ### Example with Axum
@@ -163,13 +171,25 @@ where
     }
 }
 
+#[derive(Debug, Snafu)]
+pub enum ServerHostError {
+    #[snafu(display("failed to parse port {port:?} as u16 from string"))]
+    ParsePort { source: ParseIntError, port: String },
+
+    #[snafu(display("encountered invalid request scheme {scheme:?}"))]
+    InvalidScheme { scheme: String },
+
+    #[snafu(display("failed to extract any host information from request"))]
+    ExtractHost,
+}
+
 /// This trait provides various helper functions to extract data from a
 /// HTTP [`Request`].
 pub trait RequestExt {
     /// Returns the client socket address, if available.
     fn client_socket_address(&self) -> Option<SocketAddr>;
 
-    /// Returns the server socket address, if available.
+    /// Returns the server host, if available.
     ///
     /// ### Value Selection Strategy
     ///
@@ -186,7 +206,7 @@ pub trait RequestExt {
     /// > - The Host header.
     ///
     /// [1]: https://opentelemetry.io/docs/specs/semconv/http/http-spans/#setting-serveraddress-and-serverport-attributes
-    fn server_socket_address(&self) -> Option<SocketAddr>;
+    fn server_host(&self) -> Result<(String, u16), ServerHostError>;
 
     /// Returns the matched path, like `/object/:object_id/tags`.
     ///
@@ -211,9 +231,33 @@ pub trait RequestExt {
 }
 
 impl RequestExt for Request {
-    fn server_socket_address(&self) -> Option<SocketAddr> {
-        let host = self.extensions().get::<Host>()?;
-        SocketAddr::from_str(&host.0).ok()
+    fn server_host(&self) -> Result<(String, u16), ServerHostError> {
+        // There is currently no obvious way to use the Host extractor from Axum
+        // directly. Using that extractor either requires impossible code (async
+        // in the Service's call function, unnecessary cloning or consuming self
+        // and returning a newly created request). That's why the following
+        // section mirrors the Axum extractor implementation. The implementation
+        // currently only looks for the X-Forwarded-Host / Host header and falls
+        // back to the request URI host. The Axum implementation also extracts
+        // data from the Forwarded header.
+
+        if let Some(host) = self
+            .headers()
+            .get(X_FORWARDED_HOST_HEADER_KEY)
+            .and_then(|host| host.to_str().ok())
+        {
+            return server_host_to_tuple(host, self.uri().scheme_str());
+        }
+
+        if let Some(host) = self.headers().get(HOST).and_then(|host| host.to_str().ok()) {
+            return server_host_to_tuple(host, self.uri().scheme_str());
+        }
+
+        if let (Some(host), Some(port)) = (self.uri().host(), self.uri().port_u16()) {
+            return Ok((host.to_owned(), port));
+        }
+
+        ExtractHostSnafu.fail()
     }
 
     fn client_socket_address(&self) -> Option<SocketAddr> {
@@ -242,6 +286,29 @@ impl RequestExt for Request {
     }
 }
 
+fn server_host_to_tuple(
+    host: &str,
+    scheme: Option<&str>,
+) -> Result<(String, u16), ServerHostError> {
+    if let Some((host, port)) = host.split_once(':') {
+        // First, see if the host header value contains a colon indicating that
+        // it includes a non-default port.
+        let port: u16 = port.parse().context(ParsePortSnafu { port })?;
+        Ok((host.to_owned(), port))
+    } else {
+        // If there is no port included in the header value, the port is implied.
+        // Port 443 for HTTPS and port 80 for HTTP.
+        let port = match scheme {
+            Some("https") => DEFAULT_HTTPS_PORT,
+            Some("http") => DEFAULT_HTTP_PORT,
+            Some(scheme) => return InvalidSchemeSnafu { scheme }.fail(),
+            _ => return InvalidSchemeSnafu { scheme: "" }.fail(),
+        };
+
+        Ok((host.to_owned(), port))
+    }
+}
+
 /// This trait provides various helper functions to create a [`Span`] out of
 /// an HTTP [`Request`].
 pub trait SpanExt {
@@ -319,7 +386,7 @@ impl SpanExt for Span {
         // - https://github.com/tokio-rs/tracing/pull/732
         //
         // Additionally we cannot use consts for field names. There was an
-        // upstream PR to add support for it, but it was unexpectingly closed.
+        // upstream PR to add support for it, but it was unexpectedly closed.
         // See https://github.com/tokio-rs/tracing/pull/2254.
         //
         // If this is eventually supported (maybe with our efforts), we can use
@@ -332,22 +399,21 @@ impl SpanExt for Span {
         debug!("create http span");
         let span = trace_span!(
             "HTTP request",
-            otel.name = span_name,
-            otel.kind = ?SpanKind::Server,
-            otel.status_code = Empty,
-            otel.status_message = Empty,
-            http.request.method = http_method,
-            http.response.status_code = Empty,
-            url.path = url.path(),
-            url.query = url.query(),
-            url.scheme = url.scheme_str().unwrap_or_default(),
-            user_agent.original = Empty,
-            server.address = Empty,
-            server.port = Empty,
-            client.address = Empty,
-            client.port = Empty,
-            http.route = Empty,
-            http.response.status_code = Empty,
+            "otel.name" = span_name,
+            "otel.kind" = ?SpanKind::Server,
+            "otel.status_code" = Empty,
+            "otel.status_message" = Empty,
+            "http.request.method" = http_method,
+            "http.response.status_code" = Empty,
+            "http.route" = Empty,
+            "url.path" = url.path(),
+            "url.query" = url.query(),
+            "url.scheme" = url.scheme_str().unwrap_or_default(),
+            "user_agent.original" = Empty,
+            "server.address" = Empty,
+            "server.port" = Empty,
+            "client.address" = Empty,
+            "client.port" = Empty,
             // TODO (@Techassi): Add network.protocol.version
         );
 
@@ -363,9 +429,12 @@ impl SpanExt for Span {
         // Setting server.address and server.port
         // See https://opentelemetry.io/docs/specs/semconv/http/http-spans/#setting-serveraddress-and-serverport-attributes
 
-        if let Some(server_socket_address) = req.server_socket_address() {
-            span.record("server.address", server_socket_address.ip().to_string())
-                .record("server.port", server_socket_address.port());
+        if let Ok((host, port)) = req.server_host() {
+            // NOTE (@Techassi): We cast to i64, because otherwise the field
+            // will NOT be recorded as a number but as a string. This is likely
+            // an issue in the tracing-opentelemetry crate.
+            span.record("server.address", host)
+                .record("server.port", port as i64);
         }
 
         // Setting fields according to the HTTP server semantic conventions
@@ -375,25 +444,22 @@ impl SpanExt for Span {
             span.record("client.address", client_socket_address.ip().to_string());
 
             if opt_in {
-                span.record("client.port", client_socket_address.port());
+                // NOTE (@Techassi): We cast to i64, because otherwise the field
+                // will NOT be recorded as a number but as a string. This is
+                // likely an issue in the tracing-opentelemetry crate.
+                span.record("client.port", client_socket_address.port() as i64);
             }
         }
 
         // Only include the headers if the user opted in, because this might
         // potentially be an expensive operation when many different headers
         // are present. The OpenTelemetry spec also marks this as opt-in.
 
-        // NOTE (@Techassi): Currently, tracing doesn't support recording
-        // fields which are not registered at span creation which thus makes
-        // it impossible to record request headers at runtime.
+        // FIXME (@Techassi): Currently, tracing doesn't support recording
+        // fields which are not registered at span creation which thus makes it
+        // impossible to record request headers at runtime.
         // See: https://github.com/tokio-rs/tracing/issues/1343
 
-        // FIXME (@Techassi): Add support for this when tracing allows dynamic
-        // fields.
-        // if opt_in {
-        //     span.add_header_fields(req.headers())
-        // }
-
         if let Some(http_route) = req.matched_path() {
             span.record("http.route", http_route.as_str());
         }
@@ -420,7 +486,11 @@ impl SpanExt for Span {
 
     fn finalize_with_response(&self, response: &mut Response) {
         let status_code = response.status();
-        self.record("http.response.status_code", status_code.as_u16());
+
+        // NOTE (@Techassi): We cast to i64, because otherwise the field will
+        // NOT be recorded as a number but as a string. This is likely an issue
+        // in the tracing-opentelemetry crate.
+        self.record("http.response.status_code", status_code.as_u16() as i64);
 
         // Only set the span status to "Error" when we encountered an server
         // error. See:

crates/stackable-webhook/CHANGELOG.md

@@ -4,6 +4,15 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix the extraction of `ConnectInfo` (data about the connection client) and
+  the `Host` info (data about the server) in the `AxumTraceLayer`. This was
+  previously not extracted correctly and thus not included in the OpenTelemetry
+  compatible traces ([#806]).
+
+[#806]: https://github.com/stackabletech/operator-rs/pull/806
+
 ## [0.3.0] - 2024-05-08
 
 ### Added

crates/stackable-webhook/src/lib.rs

@@ -173,6 +173,7 @@ impl WebhookServer {
     async fn run_server(self) -> Result<()> {
         debug!("run webhook server");
 
+        // TODO (@Techassi): Make opt-in configurable from the outside
         // Create an OpenTelemetry tracing layer
         debug!("create tracing service (layer)");
         let trace_layer = AxumTraceLayer::new().with_opt_in();

crates/stackable-webhook/src/tls.rs

@@ -18,7 +18,7 @@ use tokio_rustls::{
     },
     TlsAcceptor,
 };
-use tower::Service;
+use tower::{Service, ServiceExt};
 use tracing::{instrument, trace, warn};
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -139,10 +139,24 @@ impl TlsServer {
                     socket_addr: self.socket_addr,
                 })?;
 
+        // To be able to extract the connect info from incoming requests, it is
+        // required to turn the router into a Tower service which is capable of
+        // doing that. Calling `into_make_service_with_connect_info` returns a
+        // new struct `IntoMakeServiceWithConnectInfo` which implements the
+        // Tower Service trait. This service is called after the TCP connection
+        // has been accepted.
+        //
+        // Inspired by:
+        // - https://github.com/tokio-rs/axum/discussions/2397
+        // - https://github.com/tokio-rs/axum/blob/b02ce307371a973039018a13fa012af14775948c/examples/serve-with-hyper/src/main.rs#L98
+
+        let mut router = self
+            .router
+            .into_make_service_with_connect_info::<SocketAddr>();
+
         pin_mut!(tcp_listener);
         loop {
             let tls_acceptor = tls_acceptor.clone();
-            let router = self.router.clone();
 
             // Wait for new tcp connection
             let (tcp_stream, remote_addr) = match tcp_listener.accept().await {
@@ -153,6 +167,10 @@ impl TlsServer {
                 }
             };
 
+            // Here, the connect info is extracted by calling Tower's Service
+            // trait function on `IntoMakeServiceWithConnectInfo`
+            let tower_service = router.call(remote_addr).await.unwrap();
+
             tokio::spawn(async move {
                 // Wait for tls handshake to happen
                 let Ok(tls_stream) = tls_acceptor.accept(tcp_stream).await else {
@@ -167,16 +185,13 @@ impl TlsServer {
                 // Hyper also has its own `Service` trait and doesn't use tower. We can use
                 // `hyper::service::service_fn` to create a hyper `Service` that calls our app through
                 // `tower::Service::call`.
-                let service = service_fn(move |request: Request<Incoming>| {
-                    // We have to clone `tower_service` because hyper's `Service` uses `&self` whereas
-                    // tower's `Service` requires `&mut self`.
-                    //
-                    // We don't need to call `poll_ready` since `Router` is always ready.
-                    router.clone().call(request)
+                let hyper_service = service_fn(move |request: Request<Incoming>| {
+                    // We need to clone here, because oneshot consumes self
+                    tower_service.clone().oneshot(request)
                 });
 
                 if let Err(err) = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new())
-                    .serve_connection_with_upgrades(tls_stream, service)
+                    .serve_connection_with_upgrades(tls_stream, hyper_service)
                     .await
                 {
                     warn!(%err, %remote_addr, "failed to serve connection");