baukit_http/
middleware.rs

1use std::{any::Any, fmt, time::Instant};
2
3use axum::{
4    Router,
5    body::Body,
6    extract::{FromRequestParts, MatchedPath, Request, State},
7    http::{
8        HeaderMap, HeaderName, HeaderValue, Method, StatusCode,
9        header::{AUTHORIZATION, CONTENT_TYPE},
10        request::Parts,
11    },
12    middleware::{self, Next},
13    response::{IntoResponse, Response},
14};
15use metrics::Gauge;
16use opentelemetry::{
17    Context, global,
18    propagation::{Extractor, Injector},
19};
20use tokio::time;
21use tower::limit::ConcurrencyLimitLayer;
22use tower_http::{
23    catch_panic::CatchPanicLayer,
24    cors::{AllowOrigin, CorsLayer},
25    limit::RequestBodyLimitLayer,
26};
27use tracing::Instrument as _;
28use tracing_opentelemetry::OpenTelemetrySpanExt as _;
29use uuid::Uuid;
30
31use crate::{ApiError, HttpOptions, extract::JsonRejectionCode};
32
33/// The standard request ID header.
34pub const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
35
36const TRACEPARENT: HeaderName = HeaderName::from_static("traceparent");
37const TRACESTATE: HeaderName = HeaderName::from_static("tracestate");
38const UNMATCHED_ROUTE: &str = "unmatched";
39
40/// Counter metric for completed HTTP requests, labeled `method`, `route`, and `status`.
41pub const HTTP_REQUESTS_TOTAL: &str = "http_requests_total";
42/// Histogram metric for HTTP request seconds, labeled `method`, `route`, and `status`.
43pub const HTTP_REQUEST_DURATION_SECONDS: &str = "http_request_duration_seconds";
44/// Gauge metric for active HTTP requests, labeled `method` and `route`.
45pub const HTTP_REQUESTS_IN_FLIGHT: &str = "http_requests_in_flight";
46
47/// Required Prometheus buckets for [`HTTP_REQUEST_DURATION_SECONDS`], in seconds.
48///
49/// This crate records through the metrics facade and deliberately does not own a
50/// recorder. The recorder setup (normally `baukit-telemetry`) must pass this
51/// slice to `metrics_exporter_prometheus::PrometheusBuilder::set_buckets` or a
52/// metric-specific equivalent before installing the recorder.
53pub const DURATION_BUCKETS: &[f64] = &[
54    0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
55];
56
57tokio::task_local! {
58    static CURRENT_REQUEST_ID: RequestId;
59}
60
61/// A request identifier accepted from the client or generated as a UUID.
62///
63/// The defaults stack inserts this value into request extensions, so handlers
64/// can use it directly as an Axum extractor.
65#[derive(Clone, Debug, Eq, Hash, PartialEq)]
66pub struct RequestId(String);
67
68impl RequestId {
69    fn from_headers(headers: &HeaderMap) -> Self {
70        headers
71            .get(&X_REQUEST_ID)
72            .and_then(|value| value.to_str().ok())
73            .filter(|value| !value.is_empty())
74            .map_or_else(Self::generate, |value| Self(value.to_owned()))
75    }
76
77    fn generate() -> Self {
78        Self(Uuid::new_v4().to_string())
79    }
80
81    /// Returns the request identifier text.
82    #[must_use]
83    pub fn as_str(&self) -> &str {
84        &self.0
85    }
86
87    pub(crate) fn header_value(
88        &self,
89    ) -> Result<HeaderValue, axum::http::header::InvalidHeaderValue> {
90        HeaderValue::from_str(&self.0)
91    }
92}
93
94impl fmt::Display for RequestId {
95    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
96        formatter.write_str(&self.0)
97    }
98}
99
100impl<S> FromRequestParts<S> for RequestId
101where
102    S: Send + Sync,
103{
104    type Rejection = ApiError;
105
106    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
107        parts.extensions.get::<Self>().cloned().ok_or_else(|| {
108            ApiError::internal(std::io::Error::other(
109                "RequestId extractor used without baukit_http::layers",
110            ))
111        })
112    }
113}
114
115/// Extracts W3C trace context from request headers using the global propagator.
116#[must_use]
117pub fn extract_trace_context(headers: &HeaderMap) -> Context {
118    global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)))
119}
120
121/// Injects `context` into outbound headers using the global W3C propagator.
122///
123/// Use this for outbound HTTP requests. The response path uses the same helper
124/// to expose applicable trace headers to callers.
125pub fn inject_trace_context(context: &Context, headers: &mut HeaderMap) {
126    global::get_text_map_propagator(|propagator| {
127        propagator.inject_context(context, &mut HeaderInjector(headers));
128    });
129}
130
131/// Injects the current tracing span's OpenTelemetry context into outbound headers.
132pub fn inject_current_trace_context(headers: &mut HeaderMap) {
133    inject_trace_context(&tracing::Span::current().context(), headers);
134}
135
136/// Applies the complete Baukit HTTP defaults stack to an Axum router.
137///
138/// Construct `options` through [`HttpOptions::default`] or
139/// [`HttpOptions::from_config`]. Invalid zero-valued public limits are replaced
140/// by their sane defaults; validated constructors reject them earlier.
141pub fn layers<S>(router: Router<S>, options: HttpOptions) -> Router<S>
142where
143    S: Clone + Send + Sync + 'static,
144{
145    let options = options.validate().unwrap_or_default();
146    let cors = cors_layer(&options);
147    let body_size_limit = options.body_size_limit;
148    let concurrency_limit = options.concurrency_limit;
149
150    router
151        .layer(RequestBodyLimitLayer::new(body_size_limit))
152        .layer(ConcurrencyLimitLayer::new(concurrency_limit))
153        .layer(cors)
154        .layer(CatchPanicLayer::custom(panic_response))
155        .layer(middleware::from_fn_with_state(options, request_lifecycle))
156}
157
158fn cors_layer(options: &HttpOptions) -> CorsLayer {
159    let mut allowed_headers = vec![
160        AUTHORIZATION,
161        CONTENT_TYPE,
162        X_REQUEST_ID,
163        TRACEPARENT,
164        TRACESTATE,
165    ];
166    for header in &options.additional_allowed_headers {
167        if !allowed_headers.contains(header) {
168            allowed_headers.push(header.clone());
169        }
170    }
171    let mut cors = CorsLayer::new()
172        .allow_methods([
173            Method::GET,
174            Method::HEAD,
175            Method::POST,
176            Method::PUT,
177            Method::PATCH,
178            Method::DELETE,
179            Method::OPTIONS,
180        ])
181        .allow_headers(allowed_headers)
182        .expose_headers([X_REQUEST_ID, TRACEPARENT, TRACESTATE]);
183    if !options.allowed_origins.is_empty() {
184        cors = cors.allow_origin(AllowOrigin::list(options.allowed_origins.clone()));
185    }
186    if options.allow_credentials {
187        cors = cors.allow_credentials(true);
188    }
189    cors
190}
191
192async fn request_lifecycle(
193    State(options): State<HttpOptions>,
194    mut request: Request,
195    next: Next,
196) -> Response {
197    let request_id = RequestId::from_headers(request.headers());
198    request.extensions_mut().insert(request_id.clone());
199    request
200        .extensions_mut()
201        .insert(JsonRejectionCode(options.json_rejection_code.clone()));
202    CURRENT_REQUEST_ID
203        .scope(request_id, lifecycle_inner(options, request, next))
204        .await
205}
206
207async fn lifecycle_inner(options: HttpOptions, request: Request, next: Next) -> Response {
208    let request_id = current_request_id();
209    let method = bounded_method(request.method());
210    let route = request
211        .extensions()
212        .get::<MatchedPath>()
213        .map_or(UNMATCHED_ROUTE, MatchedPath::as_str)
214        .to_owned();
215    let parent_context = extract_trace_context(request.headers());
216    let span_name = format!("{method} {route}");
217    let span = tracing::info_span!(
218        "http.request",
219        otel.name = %span_name,
220        http.request.method = method,
221        http.route = %route,
222        http.response.status_code = tracing::field::Empty,
223        request_id = %request_id,
224    );
225    if let Err(error) = span.set_parent(parent_context) {
226        tracing::debug!(error = %error, "could not attach inbound trace context");
227    }
228
229    async move {
230        let started = Instant::now();
231        let in_flight = metrics::gauge!(
232            HTTP_REQUESTS_IN_FLIGHT,
233            "method" => method,
234            "route" => route.clone()
235        );
236        in_flight.increment(1.0);
237        let _in_flight = InFlightGuard(in_flight);
238
239        let mut response = match time::timeout(options.request_timeout, next.run(request)).await {
240            Ok(mut response) => {
241                if response.extensions_mut().remove::<PanicCaught>().is_some() {
242                    tracing::error!("request handler panicked");
243                    ApiError::internal_without_cause().into_response()
244                } else if response.status() == StatusCode::PAYLOAD_TOO_LARGE {
245                    ApiError::payload_too_large().into_response()
246                } else {
247                    response
248                }
249            }
250            Err(_) => {
251                tracing::warn!("request timed out");
252                ApiError::timeout().into_response()
253            }
254        };
255
256        let status = response.status();
257        let status_label = status.as_u16().to_string();
258        tracing::Span::current().record("http.response.status_code", status.as_u16());
259        metrics::counter!(
260            HTTP_REQUESTS_TOTAL,
261            "method" => method,
262            "route" => route.clone(),
263            "status" => status_label.clone()
264        )
265        .increment(1);
266        metrics::histogram!(
267            HTTP_REQUEST_DURATION_SECONDS,
268            "method" => method,
269            "route" => route,
270            "status" => status_label
271        )
272        .record(started.elapsed().as_secs_f64());
273
274        response.headers_mut().insert(
275            X_REQUEST_ID,
276            current_request_id()
277                .header_value()
278                .expect("request IDs originate from a valid header or UUID"),
279        );
280        inject_current_trace_context(response.headers_mut());
281        response
282    }
283    .instrument(span)
284    .await
285}
286
287pub(crate) fn current_request_id() -> RequestId {
288    CURRENT_REQUEST_ID
289        .try_with(Clone::clone)
290        .unwrap_or_else(|_| RequestId::generate())
291}
292
293fn bounded_method(method: &Method) -> &'static str {
294    match *method {
295        Method::CONNECT => "CONNECT",
296        Method::DELETE => "DELETE",
297        Method::GET => "GET",
298        Method::HEAD => "HEAD",
299        Method::OPTIONS => "OPTIONS",
300        Method::PATCH => "PATCH",
301        Method::POST => "POST",
302        Method::PUT => "PUT",
303        Method::TRACE => "TRACE",
304        _ => "OTHER",
305    }
306}
307
308#[derive(Clone, Copy, Debug)]
309struct PanicCaught;
310
311fn panic_response(_panic: Box<dyn Any + Send + 'static>) -> Response<Body> {
312    let mut response = StatusCode::INTERNAL_SERVER_ERROR.into_response();
313    response.extensions_mut().insert(PanicCaught);
314    response
315}
316
317struct InFlightGuard(Gauge);
318
319impl Drop for InFlightGuard {
320    fn drop(&mut self) {
321        self.0.decrement(1.0);
322    }
323}
324
325struct HeaderExtractor<'a>(&'a HeaderMap);
326
327impl Extractor for HeaderExtractor<'_> {
328    fn get(&self, key: &str) -> Option<&str> {
329        self.0.get(key).and_then(|value| value.to_str().ok())
330    }
331
332    fn keys(&self) -> Vec<&str> {
333        self.0.keys().map(HeaderName::as_str).collect()
334    }
335}
336
337struct HeaderInjector<'a>(&'a mut HeaderMap);
338
339impl Injector for HeaderInjector<'_> {
340    fn set(&mut self, key: &str, value: String) {
341        if let (Ok(name), Ok(value)) = (HeaderName::from_bytes(key.as_bytes()), value.parse()) {
342            self.0.insert(name, value);
343        }
344    }
345}
baukit_http/middleware.rs

baukit_http/
middleware.rs
