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};
32
33pub 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
40pub const HTTP_REQUESTS_TOTAL: &str = "http_requests_total";
42pub const HTTP_REQUEST_DURATION_SECONDS: &str = "http_request_duration_seconds";
44pub const HTTP_REQUESTS_IN_FLIGHT: &str = "http_requests_in_flight";
46
47pub 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#[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 #[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#[must_use]
117pub fn extract_trace_context(headers: &HeaderMap) -> Context {
118 global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)))
119}
120
121pub 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
131pub fn inject_current_trace_context(headers: &mut HeaderMap) {
133 inject_trace_context(&tracing::Span::current().context(), headers);
134}
135
136pub 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(options.json_rejection_mode.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::configured_payload_too_large(&options.json_rejection_mode)
246 .into_response()
247 } else {
248 response
249 }
250 }
251 Err(_) => {
252 tracing::warn!("request timed out");
253 ApiError::timeout().into_response()
254 }
255 };
256
257 let status = response.status();
258 let status_label = status.as_u16().to_string();
259 tracing::Span::current().record("http.response.status_code", status.as_u16());
260 metrics::counter!(
261 HTTP_REQUESTS_TOTAL,
262 "method" => method,
263 "route" => route.clone(),
264 "status" => status_label.clone()
265 )
266 .increment(1);
267 metrics::histogram!(
268 HTTP_REQUEST_DURATION_SECONDS,
269 "method" => method,
270 "route" => route,
271 "status" => status_label
272 )
273 .record(started.elapsed().as_secs_f64());
274
275 response.headers_mut().insert(
276 X_REQUEST_ID,
277 current_request_id()
278 .header_value()
279 .expect("request IDs originate from a valid header or UUID"),
280 );
281 inject_current_trace_context(response.headers_mut());
282 response
283 }
284 .instrument(span)
285 .await
286}
287
288pub(crate) fn current_request_id() -> RequestId {
289 CURRENT_REQUEST_ID
290 .try_with(Clone::clone)
291 .unwrap_or_else(|_| RequestId::generate())
292}
293
294fn bounded_method(method: &Method) -> &'static str {
295 match *method {
296 Method::CONNECT => "CONNECT",
297 Method::DELETE => "DELETE",
298 Method::GET => "GET",
299 Method::HEAD => "HEAD",
300 Method::OPTIONS => "OPTIONS",
301 Method::PATCH => "PATCH",
302 Method::POST => "POST",
303 Method::PUT => "PUT",
304 Method::TRACE => "TRACE",
305 _ => "OTHER",
306 }
307}
308
309#[derive(Clone, Copy, Debug)]
310struct PanicCaught;
311
312fn panic_response(_panic: Box<dyn Any + Send + 'static>) -> Response<Body> {
313 let mut response = StatusCode::INTERNAL_SERVER_ERROR.into_response();
314 response.extensions_mut().insert(PanicCaught);
315 response
316}
317
318struct InFlightGuard(Gauge);
319
320impl Drop for InFlightGuard {
321 fn drop(&mut self) {
322 self.0.decrement(1.0);
323 }
324}
325
326struct HeaderExtractor<'a>(&'a HeaderMap);
327
328impl Extractor for HeaderExtractor<'_> {
329 fn get(&self, key: &str) -> Option<&str> {
330 self.0.get(key).and_then(|value| value.to_str().ok())
331 }
332
333 fn keys(&self) -> Vec<&str> {
334 self.0.keys().map(HeaderName::as_str).collect()
335 }
336}
337
338struct HeaderInjector<'a>(&'a mut HeaderMap);
339
340impl Injector for HeaderInjector<'_> {
341 fn set(&mut self, key: &str, value: String) {
342 if let (Ok(name), Ok(value)) = (HeaderName::from_bytes(key.as_bytes()), value.parse()) {
343 self.0.insert(name, value);
344 }
345 }
346}