1use std::{
4 future::Future,
5 net::{IpAddr, SocketAddr},
6};
7
8use axum::extract::FromRequestParts;
9use http::{HeaderMap, Method, request::Parts};
10use serde::Serialize;
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
14#[serde(rename_all = "snake_case")]
15pub enum ClientKind {
16 ApiKey,
18 Authenticated,
20 Anonymous,
22}
23
24impl ClientKind {
25 pub const fn as_str(self) -> &'static str {
27 match self {
28 Self::ApiKey => "api_key",
29 Self::Authenticated => "authenticated",
30 Self::Anonymous => "anonymous",
31 }
32 }
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
63pub struct RequestContext {
64 request_id: String,
65 correlation_id: Option<String>,
66 method: Method,
67 route: Option<String>,
68 path: String,
69 trace_id: Option<String>,
70 span_id: Option<String>,
71 client_kind: ClientKind,
72 user_id: Option<String>,
73 tenant_id: Option<String>,
74 session_id: Option<String>,
75}
76
77impl RequestContext {
78 pub fn new(request_id: impl Into<String>, method: Method, path: impl Into<String>) -> Self {
84 Self {
85 request_id: request_id.into(),
86 correlation_id: None,
87 method,
88 route: None,
89 path: path.into(),
90 trace_id: None,
91 span_id: None,
92 client_kind: ClientKind::Anonymous,
93 user_id: None,
94 tenant_id: None,
95 session_id: None,
96 }
97 }
98
99 pub fn from_parts(parts: &Parts, request_id: impl Into<String>) -> Self {
106 let request_id = request_id.into();
107 let correlation_id = header_to_string(&parts.headers, "x-correlation-id")
108 .or_else(|| (!request_id.is_empty()).then(|| request_id.clone()));
109 let mut context = Self::new(request_id, parts.method.clone(), parts.uri.path());
110 context.correlation_id = correlation_id;
111 context.route = parts
112 .extensions
113 .get::<axum::extract::MatchedPath>()
114 .map(|path| path.as_str().to_owned());
115 context.client_kind = infer_client_kind(&parts.headers);
116 if let Some(trace_context) = parts
117 .headers
118 .get("traceparent")
119 .and_then(|value| value.to_str().ok())
120 .and_then(parse_traceparent)
121 {
122 context.trace_id = Some(trace_context.trace_id.to_owned());
123 context.span_id = Some(trace_context.parent_id.to_owned());
124 }
125 context
126 }
127
128 pub fn request_id(&self) -> &str {
133 &self.request_id
134 }
135
136 pub(crate) fn into_request_id(self) -> String {
137 self.request_id
138 }
139
140 pub fn correlation_id(&self) -> Option<&str> {
145 self.correlation_id.as_deref()
146 }
147
148 pub const fn method(&self) -> &Method {
150 &self.method
151 }
152
153 pub fn route(&self) -> Option<&str> {
159 self.route.as_deref()
160 }
161
162 pub fn path(&self) -> &str {
164 &self.path
165 }
166
167 pub fn trace_id(&self) -> Option<&str> {
171 self.trace_id.as_deref()
172 }
173
174 pub fn span_id(&self) -> Option<&str> {
176 self.span_id.as_deref()
177 }
178
179 pub const fn client_kind(&self) -> ClientKind {
184 self.client_kind
185 }
186
187 pub fn user_id(&self) -> Option<&str> {
189 self.user_id.as_deref()
190 }
191
192 pub fn tenant_id(&self) -> Option<&str> {
194 self.tenant_id.as_deref()
195 }
196
197 pub fn session_id(&self) -> Option<&str> {
199 self.session_id.as_deref()
200 }
201
202 pub fn with_route(mut self, route: impl Into<String>) -> Self {
204 self.route = Some(route.into());
205 self
206 }
207
208 pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
210 self.user_id = Some(user_id.into());
211 self
212 }
213
214 pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
216 self.tenant_id = Some(tenant_id.into());
217 self
218 }
219
220 pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
222 self.session_id = Some(session_id.into());
223 self
224 }
225}
226
227impl<S> FromRequestParts<S> for RequestContext
228where
229 S: Send + Sync,
230{
231 type Rejection = axum::http::StatusCode;
232
233 fn from_request_parts(
234 parts: &mut Parts,
235 _state: &S,
236 ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
237 let context = parts.extensions.get::<Self>().cloned();
238 async move { context.ok_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR) }
239 }
240}
241
242#[derive(Clone, Debug, Eq, Hash, PartialEq)]
244pub struct RequestIdentity(String);
245
246impl RequestIdentity {
247 pub fn new(value: impl Into<String>) -> Self {
249 Self(value.into())
250 }
251
252 pub fn as_str(&self) -> &str {
254 &self.0
255 }
256}
257
258pub trait IdentityExtractor: Clone + Send + Sync + 'static {
260 fn extract(&self, parts: &Parts) -> Option<RequestIdentity>;
262}
263
264impl<F> IdentityExtractor for F
265where
266 F: Fn(&Parts) -> Option<RequestIdentity> + Clone + Send + Sync + 'static,
267{
268 fn extract(&self, parts: &Parts) -> Option<RequestIdentity> {
269 self(parts)
270 }
271}
272
273pub fn context_identity() -> impl IdentityExtractor {
275 |parts: &Parts| {
276 if let Some(context) = parts.extensions.get::<RequestContext>()
277 && let Some(value) = context.user_id().or_else(|| context.tenant_id())
278 {
279 return Some(RequestIdentity::new(value.to_owned()));
280 }
281 header_to_string(&parts.headers, "x-api-key").map(RequestIdentity::new)
282 }
283}
284
285pub fn api_key_identity() -> impl IdentityExtractor {
287 |parts: &Parts| header_to_string(&parts.headers, "x-api-key").map(RequestIdentity::new)
288}
289
290pub fn client_ip_identity() -> impl IdentityExtractor {
297 |parts: &Parts| {
298 peer_ip(parts)
299 .map(|ip| RequestIdentity::new(ip.to_string()))
300 .or_else(|| Some(RequestIdentity::new("anonymous")))
301 }
302}
303
304pub fn trusted_proxy_client_ip_identity(
312 trusted_proxies: impl IntoIterator<Item = IpAddr>,
313) -> impl IdentityExtractor {
314 let trusted_proxies = trusted_proxies.into_iter().collect::<Vec<_>>();
315 move |parts: &Parts| {
316 peer_ip(parts)
317 .map(|peer| {
318 if trusted_proxies.contains(&peer)
319 && let Some(forwarded_ip) = forwarded_for_ip(&parts.headers)
320 {
321 RequestIdentity::new(forwarded_ip.to_string())
322 } else {
323 RequestIdentity::new(peer.to_string())
324 }
325 })
326 .or_else(|| Some(RequestIdentity::new("anonymous")))
327 }
328}
329
330fn peer_ip(parts: &Parts) -> Option<IpAddr> {
331 parts
332 .extensions
333 .get::<axum::extract::ConnectInfo<SocketAddr>>()
334 .map(|connect| connect.0.ip())
335}
336
337fn forwarded_for_ip(headers: &HeaderMap) -> Option<IpAddr> {
338 header_to_string(headers, "x-forwarded-for").and_then(|value| {
339 value
340 .split(',')
341 .next()
342 .map(str::trim)
343 .filter(|value| !value.is_empty())
344 .and_then(|value| value.parse().ok())
345 })
346}
347
348pub(crate) fn header_to_string(headers: &HeaderMap, name: &'static str) -> Option<String> {
349 header_to_str(headers, name).map(str::to_owned)
350}
351
352pub(crate) fn header_to_str<'a>(headers: &'a HeaderMap, name: &'static str) -> Option<&'a str> {
353 headers
354 .get(name)
355 .and_then(|value| value.to_str().ok())
356 .filter(|value| !value.is_empty())
357}
358
359#[derive(Clone, Copy, Debug, Eq, PartialEq)]
360pub(crate) struct ParsedTraceParent<'a> {
361 pub(crate) trace_id: &'a str,
362 pub(crate) parent_id: &'a str,
363 pub(crate) flags: u8,
364}
365
366pub(crate) fn parse_traceparent(value: &str) -> Option<ParsedTraceParent<'_>> {
367 let bytes = value.as_bytes();
368 if bytes.len() < 55 || bytes[2] != b'-' || bytes[35] != b'-' || bytes[52] != b'-' {
369 return None;
370 }
371
372 let version = &bytes[..2];
373 let trace_id = &bytes[3..35];
374 let parent_id = &bytes[36..52];
375 let flags = &bytes[53..55];
376 if !is_lower_hex(version)
377 || version == b"ff"
378 || !is_valid_id(trace_id)
379 || !is_valid_id(parent_id)
380 || !is_lower_hex(flags)
381 || (version == b"00" && bytes.len() != 55)
382 || (bytes.len() > 55 && bytes[55] != b'-')
383 {
384 return None;
385 }
386
387 Some(ParsedTraceParent {
388 trace_id: &value[3..35],
389 parent_id: &value[36..52],
390 flags: u8::from_str_radix(&value[53..55], 16).ok()?,
391 })
392}
393
394fn is_valid_id(value: &[u8]) -> bool {
395 is_lower_hex(value) && value.iter().any(|byte| *byte != b'0')
396}
397
398fn is_lower_hex(value: &[u8]) -> bool {
399 value
400 .iter()
401 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(byte))
402}
403
404fn infer_client_kind(headers: &HeaderMap) -> ClientKind {
405 if headers.contains_key("x-api-key") {
406 ClientKind::ApiKey
407 } else if headers.contains_key(http::header::AUTHORIZATION) {
408 ClientKind::Authenticated
409 } else {
410 ClientKind::Anonymous
411 }
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417 use http::Request;
418
419 #[test]
420 fn request_context_can_consume_request_id() {
421 let context = RequestContext::new("req-123", Method::GET, "/users");
422
423 assert_eq!(context.into_request_id(), "req-123");
424 }
425
426 #[test]
427 fn request_context_extracts_valid_trace_and_parent_span_ids() {
428 let (parts, ()) = Request::builder()
429 .uri("/users/42")
430 .header(
431 "traceparent",
432 "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
433 )
434 .body(())
435 .unwrap()
436 .into_parts();
437
438 let context = RequestContext::from_parts(&parts, "request-1");
439
440 assert_eq!(context.trace_id(), Some("4bf92f3577b34da6a3ce929d0e0e4736"));
441 assert_eq!(context.span_id(), Some("00f067aa0ba902b7"));
442 }
443
444 #[test]
445 fn request_context_ignores_invalid_traceparent_values() {
446 for traceparent in [
447 "not-a-traceparent",
448 "ff-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
449 "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
450 "00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01",
451 "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01-extra",
452 ] {
453 let (parts, ()) = Request::builder()
454 .uri("/")
455 .header("traceparent", traceparent)
456 .body(())
457 .unwrap()
458 .into_parts();
459 let context = RequestContext::from_parts(&parts, "request-1");
460
461 assert_eq!(context.trace_id(), None, "{traceparent}");
462 assert_eq!(context.span_id(), None, "{traceparent}");
463 }
464 }
465
466 #[test]
467 fn client_ip_identity_ignores_forwarded_headers_without_peer_info() {
468 let parts = request_parts(None, Some("203.0.113.10"));
469 let identity = client_ip_identity().extract(&parts).unwrap();
470
471 assert_eq!(identity.as_str(), "anonymous");
472 }
473
474 #[test]
475 fn trusted_proxy_client_ip_identity_uses_forwarded_header_from_trusted_peer() {
476 let parts = request_parts(Some("127.0.0.1:5000"), Some("203.0.113.10, 10.0.0.5"));
477 let trusted_proxy = "127.0.0.1".parse::<IpAddr>().unwrap();
478 let identity = trusted_proxy_client_ip_identity([trusted_proxy])
479 .extract(&parts)
480 .unwrap();
481
482 assert_eq!(identity.as_str(), "203.0.113.10");
483 }
484
485 #[test]
486 fn trusted_proxy_client_ip_identity_ignores_forwarded_header_from_untrusted_peer() {
487 let parts = request_parts(Some("127.0.0.1:5000"), Some("203.0.113.10"));
488 let trusted_proxy = "10.0.0.1".parse::<IpAddr>().unwrap();
489 let identity = trusted_proxy_client_ip_identity([trusted_proxy])
490 .extract(&parts)
491 .unwrap();
492
493 assert_eq!(identity.as_str(), "127.0.0.1");
494 }
495
496 fn request_parts(peer: Option<&str>, forwarded_for: Option<&str>) -> Parts {
497 let mut builder = Request::builder().uri("/");
498 if let Some(forwarded_for) = forwarded_for {
499 builder = builder.header("x-forwarded-for", forwarded_for);
500 }
501 let (mut parts, ()) = builder.body(()).unwrap().into_parts();
502 if let Some(peer) = peer {
503 parts.extensions.insert(axum::extract::ConnectInfo(
504 peer.parse::<SocketAddr>().unwrap(),
505 ));
506 }
507 parts
508 }
509}