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

baukit_http/
middleware.rs
