1use std::{
2 future::Future,
3 pin::Pin,
4 sync::Arc,
5 task::{Context, Poll},
6};
7
8use axum::{body::Body, response::IntoResponse};
9use http::{HeaderValue, Request, Response, StatusCode, header::HeaderName};
10use serde::Serialize;
11use tower::{Layer, Service};
12use uuid::{Uuid, Version};
13
14use crate::context::RequestContext;
15
16#[derive(Clone, Copy, Debug, Default)]
27pub struct RequestIdLayer;
28
29impl<S> Layer<S> for RequestIdLayer {
30 type Service = RequestIdService<S>;
31
32 fn layer(&self, inner: S) -> Self::Service {
33 RequestIdService { inner }
34 }
35}
36
37#[derive(Clone, Debug)]
39pub struct RequestIdService<S> {
40 inner: S,
41}
42
43impl<S, RequestBody, ResponseBody> Service<Request<RequestBody>> for RequestIdService<S>
44where
45 S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
46 S::Future: Send + 'static,
47 S::Error: Send + 'static,
48 RequestBody: Send + 'static,
49 ResponseBody: Send + 'static,
50{
51 type Response = Response<ResponseBody>;
52 type Error = S::Error;
53 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
54
55 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
56 self.inner.poll_ready(cx)
57 }
58
59 fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
60 let request_id = request.headers().get(&REQUEST_ID_HEADER).cloned();
61 let future = self.inner.call(request);
62 Box::pin(async move {
63 let mut response = future.await?;
64 response
65 .headers_mut()
66 .entry(&REQUEST_ID_HEADER)
67 .or_insert_with(|| request_id.unwrap_or_else(new_request_id));
68 Ok(response)
69 })
70 }
71}
72
73static REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
74
75fn new_request_id() -> HeaderValue {
76 HeaderValue::from_str(&Uuid::new_v4().to_string())
77 .expect("generated request id contains only valid header characters")
78}
79
80#[derive(Clone, Copy, Debug, Eq, PartialEq)]
86pub enum RequestIdMode {
87 Permissive,
94 Strict,
100}
101
102pub type RequestIdPolicy = RequestIdMode;
104
105#[derive(Clone)]
116pub struct RequestIdConfig {
117 header_name: HeaderName,
118 mode: RequestIdMode,
119 generator: Arc<dyn Fn() -> String + Send + Sync>,
120}
121
122impl std::fmt::Debug for RequestIdConfig {
123 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 formatter
125 .debug_struct("RequestIdConfig")
126 .field("header_name", &self.header_name)
127 .field("mode", &self.mode)
128 .finish_non_exhaustive()
129 }
130}
131
132impl RequestIdConfig {
133 pub fn production() -> Self {
143 Self {
144 header_name: REQUEST_ID_HEADER.clone(),
145 mode: RequestIdMode::Strict,
146 generator: Arc::new(|| Uuid::new_v4().to_string()),
147 }
148 }
149
150 pub fn development() -> Self {
157 Self::production().mode(RequestIdMode::Permissive)
158 }
159
160 pub fn header_name(mut self, header_name: HeaderName) -> Self {
162 self.header_name = header_name;
163 self
164 }
165
166 pub fn mode(mut self, mode: RequestIdMode) -> Self {
168 self.mode = mode;
169 self
170 }
171
172 pub fn generator(mut self, generator: impl Fn() -> String + Send + Sync + 'static) -> Self {
180 self.generator = Arc::new(generator);
181 self
182 }
183
184 pub fn header(&self) -> &HeaderName {
186 &self.header_name
187 }
188
189 pub const fn validation_mode(&self) -> RequestIdMode {
191 self.mode
192 }
193
194 fn generate(&self) -> String {
195 (self.generator)()
196 }
197}
198
199impl Default for RequestIdConfig {
200 fn default() -> Self {
201 Self::production()
202 }
203}
204
205pub fn validated_request_id_layer(config: RequestIdConfig) -> ValidatedRequestIdLayer {
227 ValidatedRequestIdLayer::new(config)
228}
229
230#[derive(Clone, Debug)]
243pub struct ValidatedRequestIdLayer {
244 config: RequestIdConfig,
245}
246
247impl ValidatedRequestIdLayer {
248 pub fn new(config: RequestIdConfig) -> Self {
250 Self { config }
251 }
252}
253
254impl<S> Layer<S> for ValidatedRequestIdLayer {
255 type Service = ValidatedRequestIdService<S>;
256
257 fn layer(&self, inner: S) -> Self::Service {
258 ValidatedRequestIdService {
259 inner,
260 config: self.config.clone(),
261 }
262 }
263}
264
265#[derive(Clone, Debug)]
267pub struct ValidatedRequestIdService<S> {
268 inner: S,
269 config: RequestIdConfig,
270}
271
272impl<S> Service<Request<Body>> for ValidatedRequestIdService<S>
273where
274 S: Service<Request<Body>, Response = Response<Body>> + Send + 'static,
275 S::Future: Send + 'static,
276 S::Error: Send + 'static,
277{
278 type Response = Response<Body>;
279 type Error = S::Error;
280 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
281
282 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
283 self.inner.poll_ready(cx)
284 }
285
286 fn call(&mut self, request: Request<Body>) -> Self::Future {
287 let config = &self.config;
288 let (mut parts, body) = request.into_parts();
289
290 enum Inbound {
291 Valid(String, HeaderValue),
292 Invalid,
293 Missing,
294 }
295
296 let inbound = match parts.headers.get(config.header()) {
297 None => Inbound::Missing,
298 Some(value) => match value.to_str() {
299 Ok(text) if is_valid_request_id(text) => {
302 Inbound::Valid(text.to_owned(), value.clone())
303 }
304 _ => Inbound::Invalid,
305 },
306 };
307
308 let (request_id, header_value) = match inbound {
309 Inbound::Valid(request_id, header_value) => (request_id, header_value),
310 Inbound::Invalid if config.validation_mode() == RequestIdMode::Strict => {
311 let (request_id, header_value) = match generated_request_id_header(config) {
312 Some(generated) => generated,
313 None => {
314 return Box::pin(async move {
315 Ok(invalid_generated_request_id_response(parts.uri.path()))
316 });
317 }
318 };
319 let mut response = invalid_request_id_response(&request_id, parts.uri.path());
320 response
321 .headers_mut()
322 .insert(config.header().clone(), header_value);
323 return Box::pin(async move { Ok(response) });
324 }
325 Inbound::Invalid | Inbound::Missing => match generated_request_id_header(config) {
326 Some(generated) => generated,
327 None => {
328 return Box::pin(async move {
329 Ok(invalid_generated_request_id_response(parts.uri.path()))
330 });
331 }
332 },
333 };
334
335 let response_header = config.header().clone();
336 parts
337 .headers
338 .insert(response_header.clone(), header_value.clone());
339 let context = RequestContext::from_parts(&parts, request_id);
340 parts.extensions.insert(context);
341 let future = self.inner.call(Request::from_parts(parts, body));
342
343 Box::pin(async move {
344 let mut response = future.await?;
345 response
346 .headers_mut()
347 .entry(response_header)
348 .or_insert(header_value);
349 Ok(response)
350 })
351 }
352}
353
354fn generated_request_id_header(config: &RequestIdConfig) -> Option<(String, HeaderValue)> {
355 let request_id = config.generate();
356 let header_value = HeaderValue::from_str(&request_id).ok()?;
357 Some((request_id, header_value))
358}
359
360fn is_valid_request_id(value: &str) -> bool {
361 Uuid::parse_str(value)
362 .ok()
363 .and_then(|uuid| uuid.get_version())
364 == Some(Version::Random)
365}
366
367fn invalid_request_id_response(request_id: &str, path: &str) -> Response<Body> {
368 let timestamp = crate::error::timestamp_now();
369 (
370 StatusCode::BAD_REQUEST,
371 axum::Json(RequestIdErrorBody {
372 error: RequestIdErrorDetails {
373 status_code: StatusCode::BAD_REQUEST.as_u16(),
374 code: "invalid_request_id",
375 message: "invalid request id",
376 details: serde_json::Value::Null,
377 timestamp,
378 path: path.to_owned(),
379 request_id: Some(request_id.to_owned()),
380 },
381 }),
382 )
383 .into_response()
384}
385
386fn invalid_generated_request_id_response(path: &str) -> Response<Body> {
387 let timestamp = crate::error::timestamp_now();
388 (
389 StatusCode::INTERNAL_SERVER_ERROR,
390 axum::Json(RequestIdErrorBody {
391 error: RequestIdErrorDetails {
392 status_code: StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
393 code: "invalid_generated_request_id",
394 message: "generated request id was not a valid HTTP header value",
395 details: serde_json::Value::Null,
396 timestamp,
397 path: path.to_owned(),
398 request_id: None,
399 },
400 }),
401 )
402 .into_response()
403}
404
405#[derive(Debug, Serialize)]
406struct RequestIdErrorBody {
407 error: RequestIdErrorDetails,
408}
409
410#[derive(Debug, Serialize)]
411#[serde(rename_all = "camelCase")]
412struct RequestIdErrorDetails {
413 status_code: u16,
414 code: &'static str,
415 message: &'static str,
416 details: serde_json::Value,
417 timestamp: String,
418 path: String,
419 #[serde(skip_serializing_if = "Option::is_none")]
420 request_id: Option<String>,
421}
422
423#[cfg(test)]
424mod tests {
425 use std::{convert::Infallible, sync::Arc};
426
427 use tower::{Service, service_fn};
428
429 use super::*;
430
431 #[tokio::test]
432 async fn service_does_not_retain_an_extra_generator_owner_across_the_inner_future() {
433 let config = RequestIdConfig::production();
434 let generator = Arc::clone(&config.generator);
435 let mut service = ValidatedRequestIdService {
436 inner: service_fn(|_request: Request<Body>| async {
437 Ok::<_, Infallible>(Response::new(Body::empty()))
438 }),
439 config,
440 };
441 assert_eq!(Arc::strong_count(&generator), 2);
442
443 let future = service.call(Request::new(Body::empty()));
444 assert_eq!(Arc::strong_count(&generator), 2);
445 let response = future.await.unwrap();
446
447 assert_eq!(response.status(), StatusCode::OK);
448 assert_eq!(Arc::strong_count(&generator), 2);
449 }
450}