From 284c9377c2904760ab5ba33561ef8d8c0c3f0bc4 Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Sat, 25 Apr 2026 22:56:17 +0200 Subject: [PATCH] Add X-Request-Id middleware MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-request ULID minted at the edge, exposed in request extensions and on the response header. Caller-supplied X-Request-Id is echoed when well-formed (1..=128 ASCII printable characters); otherwise rejected and replaced with a fresh ULID so the value is always safe to log. Companion to the TypeScript SDK redesign — clients now correlate logs across the wire by reading X-Request-Id from response headers (and the SDK already surfaces it on every OmnigraphError as `requestId`). No spec change required; the header is a transport-layer concern. Tests: - mint a ULID when no header is provided - echo a valid caller-supplied id - reject overlong header (200 chars), mint a fresh ULID Co-Authored-By: Claude Opus 4.7 --- Cargo.lock | 1 + crates/omnigraph-server/Cargo.toml | 1 + crates/omnigraph-server/src/lib.rs | 2 + crates/omnigraph-server/src/request_id.rs | 59 +++++++++++++++++ crates/omnigraph-server/tests/server.rs | 79 +++++++++++++++++++++++ 5 files changed, 142 insertions(+) create mode 100644 crates/omnigraph-server/src/request_id.rs diff --git a/Cargo.lock b/Cargo.lock index 9eafcf9..0172c49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4700,6 +4700,7 @@ dependencies = [ "tower-http", "tracing", "tracing-subscriber", + "ulid", "utoipa", ] diff --git a/crates/omnigraph-server/Cargo.toml b/crates/omnigraph-server/Cargo.toml index 0f938a6..548330f 100644 --- a/crates/omnigraph-server/Cargo.toml +++ b/crates/omnigraph-server/Cargo.toml @@ -31,6 +31,7 @@ serde_yaml = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } tower-http = { workspace = true } +ulid = { workspace = true } utoipa = { workspace = true } cedar-policy = { workspace = true } futures = { workspace = true } diff --git a/crates/omnigraph-server/src/lib.rs b/crates/omnigraph-server/src/lib.rs index eea4699..3701cad 100644 --- a/crates/omnigraph-server/src/lib.rs +++ b/crates/omnigraph-server/src/lib.rs @@ -2,6 +2,7 @@ pub mod api; pub mod auth; pub mod config; pub mod policy; +pub mod request_id; use std::collections::{HashMap, HashSet}; use std::fs; @@ -460,6 +461,7 @@ pub fn build_app(state: AppState) -> Router { .merge(protected) .layer(DefaultBodyLimit::max(DEFAULT_REQUEST_BODY_LIMIT_BYTES)) .layer(TraceLayer::new_for_http()) + .layer(middleware::from_fn(request_id::request_id_middleware)) .with_state(state) } diff --git a/crates/omnigraph-server/src/request_id.rs b/crates/omnigraph-server/src/request_id.rs new file mode 100644 index 0000000..cae5b8b --- /dev/null +++ b/crates/omnigraph-server/src/request_id.rs @@ -0,0 +1,59 @@ +//! `X-Request-Id` middleware. +//! +//! Mints a ULID per inbound request, or echoes a caller-supplied +//! `X-Request-Id` header if it's well-formed. Stores the value in request +//! extensions so handlers can include it in error bodies, log lines, or +//! audit records, and surfaces it on the response header so SDK clients +//! can correlate logs across the wire. + +use axum::{ + body::Body, + extract::Request, + http::{HeaderName, HeaderValue, header}, + middleware::Next, + response::Response, +}; +use ulid::Ulid; + +pub const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); + +/// Wraps a request id pulled out of (or minted into) request extensions. +#[derive(Clone, Debug)] +pub struct RequestId(pub String); + +impl RequestId { + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Acceptable inbound `X-Request-Id` shape: 1..=128 ASCII printable chars. +/// Rejecting wider input keeps the value safe to log and emit verbatim. +fn is_valid_inbound(raw: &str) -> bool { + !raw.is_empty() + && raw.len() <= 128 + && raw + .bytes() + .all(|b| b.is_ascii_graphic() || b == b' ' || b == b'-' || b == b'_') +} + +pub async fn request_id_middleware(mut req: Request, next: Next) -> Response { + let inbound = req + .headers() + .get(header::HeaderName::from_static("x-request-id")) + .and_then(|v| v.to_str().ok()) + .filter(|raw| is_valid_inbound(raw)); + + let id = match inbound { + Some(raw) => raw.to_owned(), + None => Ulid::new().to_string(), + }; + + req.extensions_mut().insert(RequestId(id.clone())); + + let mut response = next.run(req).await; + if let Ok(value) = HeaderValue::from_str(&id) { + response.headers_mut().insert(X_REQUEST_ID, value); + } + response +} diff --git a/crates/omnigraph-server/tests/server.rs b/crates/omnigraph-server/tests/server.rs index 77b9118..a81141d 100644 --- a/crates/omnigraph-server/tests/server.rs +++ b/crates/omnigraph-server/tests/server.rs @@ -675,6 +675,85 @@ async fn healthz_succeeds_after_startup() { } } +#[tokio::test(flavor = "multi_thread")] +async fn request_id_minted_when_absent() { + let (_temp, app) = app_for_loaded_repo().await; + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/healthz") + .method(Method::GET) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let id = response + .headers() + .get("x-request-id") + .expect("X-Request-Id missing") + .to_str() + .unwrap() + .to_owned(); + // ULIDs are 26 chars Crockford base32. + assert_eq!(id.len(), 26); + assert!(id.chars().all(|c| c.is_ascii_alphanumeric())); +} + +#[tokio::test(flavor = "multi_thread")] +async fn request_id_echoed_when_caller_supplies_valid_value() { + let (_temp, app) = app_for_loaded_repo().await; + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/healthz") + .method(Method::GET) + .header("X-Request-Id", "trace-abc123") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response + .headers() + .get("x-request-id") + .unwrap() + .to_str() + .unwrap(), + "trace-abc123" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn request_id_minted_when_caller_supplies_invalid_value() { + let (_temp, app) = app_for_loaded_repo().await; + // 200-char string is a valid HeaderValue but exceeds the inbound length cap. + let too_long = "a".repeat(200); + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/healthz") + .method(Method::GET) + .header("X-Request-Id", &too_long) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let id = response + .headers() + .get("x-request-id") + .unwrap() + .to_str() + .unwrap(); + assert_ne!(id, too_long); + assert_eq!(id.len(), 26); +} + #[tokio::test(flavor = "multi_thread")] async fn schema_drift_returns_conflict_for_snapshot_read_and_change() { let (temp, app) = app_for_loaded_repo().await;