use axum::body::Body;
use axum::extract::{FromRequest, Request};
use axum::http::{HeaderValue, Request as HttpRequest, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response as AxumResponse};
use serde::de::DeserializeOwned;
use std::sync::Arc;
use std::time::Duration;
use crate::module::HARNESS_API;
use crate::ports::{Decision, Defer, IdGen, Quota};
use crate::problem::Problem;
use crate::scope::Scope;
use crate::sidecar::{X_HARNESS_API, X_HARNESS_MODULE};
use crate::usage::Exhausted;
use tracing::info_span;
pub const X_REQUEST_ID: &str = "x-request-id";
pub const MAX_BODY_BYTES: usize = 64 * 1024;
#[allow(clippy::duration_suboptimal_units)]
const CORS_PREFLIGHT_MAX_AGE: Duration = Duration::from_secs(86_400);
pub fn request_id_is_valid(value: &str) -> bool {
(8..=128).contains(&value.len())
&& value
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RequestSummary {
pub request_id: String,
pub method: String,
pub route: String,
pub module: String,
pub status: u16,
pub ip_hash: String,
pub ua_family: String,
}
#[derive(Clone)]
pub(crate) struct ScopeState {
pub defer: Arc<dyn Defer>,
pub id_gen: Arc<dyn IdGen>,
pub module: Option<&'static str>,
}
pub(crate) async fn scope_layer(
axum::extract::State(state): axum::extract::State<ScopeState>,
mut request: Request,
next: Next,
) -> AxumResponse {
use tracing::Instrument as _;
use tracing::field::Empty;
let incoming = request
.headers()
.get(X_REQUEST_ID)
.and_then(|value| value.to_str().ok())
.filter(|value| request_id_is_valid(value));
let request_id = match incoming {
Some(valid) => valid.to_owned(),
None => state.id_gen.ulid(),
};
let method = request.method().as_str().to_owned();
let ip_hash = crate::logging::subject_hash(
&crate::rate_limit::client_ip(request.headers()).unwrap_or_default(),
);
let ua_family = request
.headers()
.get(header::USER_AGENT)
.and_then(|value| value.to_str().ok())
.map_or_else(|| "unknown".to_owned(), ua_family_of);
let span = info_span!(
"request",
request_id = %request_id,
method = %method,
route = Empty,
module = Empty,
status = Empty,
duration_ms = Empty,
ip_hash = %ip_hash,
ua_family = %ua_family,
);
let scope = Scope {
defer: Arc::clone(&state.defer),
span: span.clone(),
request_id: request_id.clone(),
};
request.extensions_mut().insert(scope);
let route = request
.extensions()
.get::<axum::extract::MatchedPath>()
.map(|matched| matched.as_str().to_owned())
.unwrap_or_default();
#[cfg(not(target_arch = "wasm32"))]
let started = std::time::Instant::now();
let future = next.run(request);
let mut response = future.instrument(span.clone()).await;
if let Ok(value) = HeaderValue::from_str(&request_id) {
response.headers_mut().insert(X_REQUEST_ID, value);
}
if let Ok(value) = HeaderValue::from_str(&HARNESS_API.to_string()) {
response.headers_mut().insert(X_HARNESS_API, value);
}
if let Some(name) = state.module
&& let Ok(value) = HeaderValue::from_str(name)
{
response.headers_mut().insert(X_HARNESS_MODULE, value);
}
let module = route
.strip_prefix("/v1/")
.and_then(|rest| rest.split('/').next())
.unwrap_or_default()
.to_owned();
let status = response.status().as_u16();
span.record("route", route.as_str());
span.record("module", module.as_str());
span.record("status", status);
#[cfg(not(target_arch = "wasm32"))]
{
let duration_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
span.record("duration_ms", duration_ms);
}
response.extensions_mut().insert(RequestSummary {
request_id,
method,
route,
module,
status,
ip_hash,
ua_family,
});
response
}
fn ua_family_of(user_agent: &str) -> String {
let token = user_agent
.split(['/', ' ', ';', '('])
.next()
.unwrap_or_default()
.to_lowercase();
let truncated: String = token.chars().take(24).collect();
if truncated.is_empty() {
"unknown".to_owned()
} else {
truncated
}
}
pub(crate) async fn security_headers_layer(request: Request, next: Next) -> AxumResponse {
let is_api = request.uri().path().starts_with("/v1/");
let mut response = next.run(request).await;
if is_api {
insert_no_store_headers(response.headers_mut());
insert_framing_headers(response.headers_mut());
}
response
}
pub(crate) async fn token_response_layer(request: Request, next: Next) -> AxumResponse {
let carries_token = request.uri().query().is_some_and(|query| {
query
.split('&')
.any(|pair| pair == "token" || pair.starts_with("token="))
});
let mut response = next.run(request).await;
if carries_token {
insert_no_store_headers(response.headers_mut());
}
response
}
fn insert_no_store_headers(headers: &mut axum::http::HeaderMap) {
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
headers.insert(
header::X_CONTENT_TYPE_OPTIONS,
HeaderValue::from_static("nosniff"),
);
if !keeps_own_referrer_policy(headers) {
headers.insert(
header::REFERRER_POLICY,
HeaderValue::from_static("no-referrer"),
);
}
}
fn keeps_own_referrer_policy(headers: &axum::http::HeaderMap) -> bool {
let mut values = headers.get_all(header::REFERRER_POLICY).iter();
match (values.next(), values.next()) {
(Some(value), None) => value
.to_str()
.is_ok_and(|policy| policy.trim().eq_ignore_ascii_case("same-origin")),
_ => false,
}
}
fn insert_framing_headers(headers: &mut axum::http::HeaderMap) {
headers.insert(header::X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
headers.append(
header::CONTENT_SECURITY_POLICY,
HeaderValue::from_static("frame-ancestors 'none'"),
);
}
pub struct Json<T>(pub T);
impl<T, S> FromRequest<S> for Json<T>
where
T: DeserializeOwned,
S: Send + Sync,
{
type Rejection = Problem;
async fn from_request(request: HttpRequest<Body>, state: &S) -> Result<Self, Self::Rejection> {
let instance = request
.extensions()
.get::<Scope>()
.map(|scope| scope.request_id.clone());
match axum::Json::<T>::from_request(request, state).await {
Ok(axum::Json(value)) => Ok(Json(value)),
Err(rejection) => {
let mut problem = match &rejection {
axum::extract::rejection::JsonRejection::BytesRejection(rejection)
if rejection.status() == StatusCode::PAYLOAD_TOO_LARGE =>
{
Problem::request_too_large()
}
_ => Problem::validation_failed(rejection.body_text()),
};
if let Some(instance) = instance {
problem = problem.instance(&instance);
}
Err(problem)
}
}
}
}
impl<T: serde::Serialize> IntoResponse for Json<T> {
fn into_response(self) -> AxumResponse {
axum::Json(self.0).into_response()
}
}
pub struct Form<T>(pub T);
impl<T, S> FromRequest<S> for Form<T>
where
T: DeserializeOwned,
S: Send + Sync,
{
type Rejection = Problem;
async fn from_request(request: HttpRequest<Body>, state: &S) -> Result<Self, Self::Rejection> {
let instance = request
.extensions()
.get::<Scope>()
.map(|scope| scope.request_id.clone());
match axum::Form::<T>::from_request(request, state).await {
Ok(axum::Form(value)) => Ok(Form(value)),
Err(rejection) => {
let mut problem = match &rejection {
axum::extract::rejection::FormRejection::BytesRejection(rejection)
if rejection.status() == StatusCode::PAYLOAD_TOO_LARGE =>
{
Problem::request_too_large()
}
_ => Problem::validation_failed(rejection.body_text()),
};
if let Some(instance) = instance {
problem = problem.instance(&instance);
}
Err(problem)
}
}
}
}
pub fn rate_limited(decision: &Decision) -> AxumResponse {
let problem = Problem::new(&crate::problems::SLUGS.rate_limited);
let mut response = problem.into_response();
let pause = decision
.retry_after
.or_else(|| decision.quota.as_ref().map(|quota| quota.reset));
if let Some(pause) = pause
&& let Ok(value) = HeaderValue::from_str(&delta_seconds(pause))
{
response.headers_mut().insert(header::RETRY_AFTER, value);
}
if let Some(quota) = decision.quota.as_ref() {
for (name, value) in quota_headers(quota) {
if let Ok(value) = HeaderValue::from_str(&value) {
response
.headers_mut()
.insert(header::HeaderName::from_static(name), value);
}
}
}
response
}
pub fn allowance_exhausted(meter: &str, outcome: &Exhausted) -> AxumResponse {
let mut response = allowance_exhausted_problem(meter, outcome).into_response();
if let Ok(value) = HeaderValue::from_str(&delta_seconds(outcome.retry_after)) {
response.headers_mut().insert(header::RETRY_AFTER, value);
}
response
}
fn allowance_exhausted_problem(meter: &str, outcome: &Exhausted) -> Problem {
Problem::new(&crate::problems::SLUGS.allowance_exhausted)
.with_extension("meter", meter)
.with_extension("limit", outcome.limit)
.with_extension("used", outcome.used)
.with_extension("period_end", crate::usage::rfc3339(outcome.period_end))
}
fn quota_headers(quota: &Quota) -> [(&'static str, String); 3] {
[
("ratelimit-limit", quota.limit.to_string()),
("ratelimit-remaining", quota.remaining.to_string()),
("ratelimit-reset", delta_seconds(quota.reset)),
]
}
fn delta_seconds(pause: Duration) -> String {
let millis = u64::try_from(pause.as_millis()).unwrap_or(u64::MAX);
millis.div_ceil(1_000).max(1).to_string()
}
pub(crate) fn cors_layer(origins: &[String]) -> tower_http::cors::CorsLayer {
use tower_http::cors::{AllowOrigin, CorsLayer};
let allowed: Vec<HeaderValue> = origins
.iter()
.filter_map(|origin| HeaderValue::from_str(origin).ok())
.collect();
CorsLayer::new()
.allow_origin(AllowOrigin::list(allowed))
.allow_methods([
axum::http::Method::GET,
axum::http::Method::POST,
axum::http::Method::PUT,
axum::http::Method::DELETE,
axum::http::Method::OPTIONS,
])
.allow_headers([header::CONTENT_TYPE, header::AUTHORIZATION])
.max_age(CORS_PREFLIGHT_MAX_AGE)
}
#[cfg(test)]
mod tests {
use super::*;
use axum::http::StatusCode;
fn header(response: &AxumResponse, name: &str) -> Option<String> {
response
.headers()
.get(name)
.and_then(|value| value.to_str().ok().map(str::to_owned))
}
#[test]
fn only_a_handlers_same_origin_policy_survives_the_no_store_stamp() {
let stamped = |policy: Option<&'static str>| {
let mut headers = axum::http::HeaderMap::new();
if let Some(policy) = policy {
headers.insert(header::REFERRER_POLICY, HeaderValue::from_static(policy));
}
insert_no_store_headers(&mut headers);
let values: Vec<_> = headers
.get_all(header::REFERRER_POLICY)
.iter()
.map(|value| value.to_str().unwrap_or_default().to_owned())
.collect();
values.join(",")
};
assert_eq!(stamped(None), "no-referrer");
assert_eq!(stamped(Some("same-origin")), "same-origin");
assert_eq!(stamped(Some("Same-Origin")), "Same-Origin");
for weaker in [
"unsafe-url",
"origin",
"strict-origin-when-cross-origin",
"no-referrer-when-downgrade",
"same-origin, unsafe-url",
] {
assert_eq!(stamped(Some(weaker)), "no-referrer", "{weaker} survived");
}
let mut headers = axum::http::HeaderMap::new();
headers.append(
header::REFERRER_POLICY,
HeaderValue::from_static("same-origin"),
);
headers.append(
header::REFERRER_POLICY,
HeaderValue::from_static("unsafe-url"),
);
insert_no_store_headers(&mut headers);
assert_eq!(headers.get_all(header::REFERRER_POLICY).iter().count(), 1);
assert_eq!(headers.get(header::REFERRER_POLICY).unwrap(), "no-referrer");
}
#[test]
fn a_pause_becomes_retry_after_and_nothing_else() {
let response = rate_limited(&Decision {
ok: false,
retry_after: Some(Duration::from_secs(7)),
quota: None,
});
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(header(&response, "retry-after").as_deref(), Some("7"));
assert!(header(&response, "ratelimit-limit").is_none());
}
#[test]
fn a_quota_rides_the_ratelimit_headers() {
let response = rate_limited(&Decision {
ok: false,
retry_after: None,
quota: Some(Quota {
limit: 30,
remaining: 0,
reset: Duration::from_secs(43),
}),
});
assert_eq!(header(&response, "retry-after").as_deref(), Some("43"));
assert_eq!(header(&response, "ratelimit-limit").as_deref(), Some("30"));
assert_eq!(
header(&response, "ratelimit-remaining").as_deref(),
Some("0")
);
assert_eq!(header(&response, "ratelimit-reset").as_deref(), Some("43"));
}
#[test]
fn the_quotas_reset_backfills_a_missing_retry_after() {
let response = rate_limited(&Decision {
ok: false,
retry_after: None,
quota: Some(Quota {
limit: 5,
remaining: 2,
reset: Duration::from_millis(1500),
}),
});
assert_eq!(header(&response, "retry-after").as_deref(), Some("2"));
assert_eq!(header(&response, "ratelimit-reset").as_deref(), Some("2"));
}
#[test]
fn an_answerless_decision_is_a_bare_429() {
let response = rate_limited(&Decision {
ok: false,
retry_after: None,
quota: None,
});
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert!(header(&response, "retry-after").is_none());
assert!(header(&response, "ratelimit-limit").is_none());
}
fn exhausted() -> Exhausted {
Exhausted {
used: 100,
limit: 100,
period_end: time::OffsetDateTime::from_unix_timestamp(1_800_000_000).expect("in range"),
retry_after: Duration::from_secs(120),
}
}
#[test]
fn an_exhausted_allowance_is_a_429_with_a_retry_after() {
let response = allowance_exhausted("mail_sent", &exhausted());
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(header(&response, "retry-after").as_deref(), Some("120"));
}
struct BodyThatNeverArrives;
impl futures_core::Stream for BodyThatNeverArrives {
type Item = Result<bytes::Bytes, std::io::Error>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
_: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
std::task::Poll::Ready(Some(Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"connection reset by peer",
))))
}
}
fn request_with_a_dead_body(content_type: &'static str) -> Request {
Request::builder()
.method("POST")
.uri("/v1/thing")
.header(header::CONTENT_TYPE, content_type)
.body(Body::from_stream(BodyThatNeverArrives))
.expect("request builds")
}
#[pollster::test]
async fn a_body_that_never_arrives_is_a_400_not_a_413() {
let Err(problem) = Json::<serde_json::Value>::from_request(
request_with_a_dead_body("application/json"),
&(),
)
.await
else {
panic!("the body never arrives");
};
assert_eq!(
problem.status,
StatusCode::BAD_REQUEST,
"a failed body read is not request-too-large: {}",
problem.slug
);
assert_eq!(problem.slug, crate::problems::SLUGS.validation_failed.slug);
}
#[pollster::test]
async fn a_form_body_that_never_arrives_is_a_400_not_a_413() {
let Err(problem) = Form::<std::collections::HashMap<String, String>>::from_request(
request_with_a_dead_body("application/x-www-form-urlencoded"),
&(),
)
.await
else {
panic!("the body never arrives");
};
assert_eq!(
problem.status,
StatusCode::BAD_REQUEST,
"a failed body read is not request-too-large: {}",
problem.slug
);
}
#[test]
fn an_exhausted_allowance_carries_its_meter_limit_used_and_period_end() {
let body = allowance_exhausted_problem("mail_sent", &exhausted())
.body("https://api.test.example/problems/");
assert_eq!(
body["type"],
"https://api.test.example/problems/usage/allowance-exhausted"
);
assert_eq!(body["meter"], "mail_sent");
assert_eq!(body["limit"], 100);
assert_eq!(body["used"], 100);
assert_eq!(body["period_end"], "2027-01-15T08:00:00Z");
assert_eq!(body["status"], 429);
}
}