packages/ak-axum/tracing: init (goauthentik#21316)

Author: rissson

Cargo.lock

@@ -14,6 +14,7 @@ ak-common.workspace = true
 axum.workspace = true
 durstr.workspace = true
 eyre.workspace = true
+tokio.workspace = true
 tower-http.workspace = true
 tower.workspace = true
 tracing.workspace = true

packages/ak-axum/Cargo.toml

@@ -2,3 +2,4 @@
 
 pub mod error;
 pub mod router;
+pub mod tracing;

packages/ak-axum/src/lib.rs

@@ -1,22 +1,32 @@
 //! Utilities for working with [`Router`].
 
 use ak_common::config;
-use axum::{Router, http::StatusCode};
+use axum::{Router, http::StatusCode, middleware::from_fn};
 use tower::ServiceBuilder;
 use tower_http::timeout::TimeoutLayer;
 
+use crate::tracing::{span_middleware, tracing_middleware};
+
 /// Wrap a [`Router`] with common middlewares.
+///
+/// Set `with_tracing` to [`true`] to log requests.
 #[inline]
-pub fn wrap_router(router: Router) -> Router {
+pub fn wrap_router(router: Router, with_tracing: bool) -> Router {
     let config = config::get();
     let timeout = durstr::parse(&config.web.timeout_http_read_header)
         .expect("Invalid duration in http timeout")
         + durstr::parse(&config.web.timeout_http_read).expect("Invalid duration in http timeout")
         + durstr::parse(&config.web.timeout_http_write).expect("Invalid duration in http timeout")
         + durstr::parse(&config.web.timeout_http_idle).expect("Invalid duration in http timeout");
-    let service_builder = ServiceBuilder::new().layer(TimeoutLayer::with_status_code(
-        StatusCode::REQUEST_TIMEOUT,
-        timeout,
-    ));
-    router.layer(service_builder)
+    let service_builder = ServiceBuilder::new()
+        .layer(TimeoutLayer::with_status_code(
+            StatusCode::REQUEST_TIMEOUT,
+            timeout,
+        ))
+        .layer(from_fn(span_middleware));
+    if with_tracing {
+        router.layer(service_builder.layer(from_fn(tracing_middleware)))
+    } else {
+        router.layer(service_builder)
+    }
 }

packages/ak-axum/src/router.rs

@@ -0,0 +1,48 @@
+//! Middlewares for tracing requests.
+
+use std::collections::HashMap;
+
+use ak_common::config;
+use axum::{extract::Request, middleware::Next, response::Response};
+use tokio::time::Instant;
+use tracing::{Instrument as _, info, info_span, trace};
+
+/// Create a [`tracing::Span`] for requests.
+pub(crate) async fn span_middleware(request: Request, next: Next) -> Response {
+    let config = config::get();
+    let http_headers = request
+        .headers()
+        .iter()
+        .filter(|(name, _)| {
+            for header in &config.log.http_headers {
+                if header.eq_ignore_ascii_case(name.as_str()) {
+                    return true;
+                }
+            }
+            false
+        })
+        .map(|(name, value)| (name.to_string().to_lowercase().replace('-', "_"), value))
+        .collect::<HashMap<_, _>>();
+    let span = info_span!(
+        "request",
+        path = %request.uri(),
+        method = %request.method(),
+        http_headers = ?http_headers,
+    );
+    next.run(request).instrument(span).await
+}
+
+/// Log requests.
+pub(crate) async fn tracing_middleware(request: Request, next: Next) -> Response {
+    let event = request.uri().clone();
+    trace!("request start");
+
+    let start = Instant::now();
+    let response = next.run(request).await;
+    let runtime = start.elapsed();
+    let status = response.status().as_u16();
+
+    info!(status = status, runtime = runtime.as_millis(), "{event}");
+
+    response
+}