Skip to main content

victauri_core/
middleware.rs

1//! Shared axum middleware for Victauri's localhost HTTP servers.
2//!
3//! Gated behind the `middleware` feature flag.  Provides thin middleware
4//! wrappers around the pure-logic security primitives in [`super::security`].
5
6use std::sync::Arc;
7
8use axum::extract::Request;
9use axum::http::StatusCode;
10use axum::middleware::Next;
11use axum::response::Response;
12
13use crate::security::{self, RateLimiter, constant_time_eq, is_allowed_origin, is_localhost_host};
14
15const BEARER_PREFIX_LEN: usize = "Bearer ".len();
16
17/// Shared authentication state holding the optional Bearer token for the MCP
18/// server.
19#[derive(Clone)]
20pub struct AuthState {
21    /// The expected Bearer token, or `None` if authentication is disabled.
22    pub token: Option<String>,
23}
24
25/// Axum middleware that validates the `Authorization: Bearer <token>` header
26/// against [`AuthState`].
27///
28/// Case-insensitive prefix matching per RFC 7235.  Constant-time token
29/// comparison via [`constant_time_eq`].
30///
31/// # Errors
32///
33/// Returns [`StatusCode::UNAUTHORIZED`] if the token is missing or invalid.
34pub async fn require_auth(
35    axum::extract::State(auth): axum::extract::State<Arc<AuthState>>,
36    request: Request,
37    next: Next,
38) -> Result<Response, StatusCode> {
39    let Some(expected) = &auth.token else {
40        return Ok(next.run(request).await);
41    };
42
43    if bearer_matches(request.headers(), expected) {
44        Ok(next.run(request).await)
45    } else {
46        tracing::warn!("Victauri: rejected request — invalid or missing auth token");
47        Err(StatusCode::UNAUTHORIZED)
48    }
49}
50
51/// Whether `headers` carry `Authorization: Bearer <expected>` (scheme matched
52/// case-insensitively per RFC 7235, token compared in constant time).
53#[must_use]
54pub fn bearer_matches(headers: &axum::http::HeaderMap, expected: &str) -> bool {
55    headers
56        .get("authorization")
57        .and_then(|v| v.to_str().ok())
58        .filter(|v| {
59            v.get(..BEARER_PREFIX_LEN)
60                .is_some_and(|scheme| scheme.eq_ignore_ascii_case("bearer "))
61        })
62        .is_some_and(|v| constant_time_eq(&v.as_bytes()[BEARER_PREFIX_LEN..], expected.as_bytes()))
63}
64
65/// Create a rate limiter with the default capacity of
66/// [`DEFAULT_RATE_LIMIT`](security::DEFAULT_RATE_LIMIT) requests per second.
67#[must_use]
68pub fn default_rate_limiter() -> Arc<RateLimiter> {
69    Arc::new(RateLimiter::new(security::DEFAULT_RATE_LIMIT))
70}
71
72/// Axum middleware that rejects requests with 429 when the token bucket is
73/// exhausted.
74///
75/// # Errors
76///
77/// Returns [`StatusCode::TOO_MANY_REQUESTS`] with `Retry-After: 1` header when
78/// the rate limit is exceeded.
79pub async fn rate_limit(
80    axum::extract::State(limiter): axum::extract::State<Arc<RateLimiter>>,
81    request: Request,
82    next: Next,
83) -> Result<
84    Response,
85    (
86        StatusCode,
87        [(axum::http::HeaderName, axum::http::HeaderValue); 1],
88    ),
89> {
90    if limiter.try_acquire() {
91        Ok(next.run(request).await)
92    } else {
93        Err((
94            StatusCode::TOO_MANY_REQUESTS,
95            [(
96                axum::http::header::RETRY_AFTER,
97                axum::http::HeaderValue::from_static("1"),
98            )],
99        ))
100    }
101}
102
103/// Axum middleware that blocks DNS rebinding attacks.
104///
105/// Rejects any request where the `Host` header is not a localhost address.
106///
107/// # Errors
108///
109/// Returns [`StatusCode::FORBIDDEN`] if the `Host` header is not `localhost`,
110/// `127.0.0.1`, or `::1`.
111pub async fn dns_rebinding_guard(request: Request, next: Next) -> Result<Response, StatusCode> {
112    let host = request
113        .headers()
114        .get("host")
115        .and_then(|v| v.to_str().ok())
116        .unwrap_or("");
117    if !is_localhost_host(host) {
118        tracing::warn!("DNS rebinding attempt blocked: Host={host}");
119        return Err(StatusCode::FORBIDDEN);
120    }
121    Ok(next.run(request).await)
122}
123
124/// Axum middleware that blocks requests a web page made a browser send.
125///
126/// Runs before the rate limiter, so a refused request never spends a token (a page looping
127/// `<img src="http://127.0.0.1:7373/health">` sends no `Origin` and used to drain the shared
128/// bucket until the agent was answered 429).
129///
130/// # Errors
131///
132/// Returns [`StatusCode::FORBIDDEN`] if the request carries a `Sec-Fetch-Site` header other
133/// than `none` (every modern browser marks page-initiated requests; `none` is a user
134/// navigation), or an `Origin` header that is unreadable or not a localhost origin.
135pub async fn origin_guard(request: Request, next: Next) -> Result<Response, StatusCode> {
136    if let Some(site) = request.headers().get("sec-fetch-site")
137        && site.as_bytes() != b"none"
138    {
139        tracing::warn!("Browser-initiated request blocked: Sec-Fetch-Site={site:?}");
140        return Err(StatusCode::FORBIDDEN);
141    }
142    if let Some(origin) = request.headers().get("origin") {
143        // An Origin that is not visible ASCII cannot be judged, so it is refused rather than
144        // treated as absent (absent means "not a browser").
145        match origin.to_str() {
146            Ok(origin) if is_allowed_origin(origin) => {}
147            _ => {
148                tracing::warn!("Cross-origin request blocked: Origin={origin:?}");
149                return Err(StatusCode::FORBIDDEN);
150            }
151        }
152    }
153    Ok(next.run(request).await)
154}
155
156/// Axum middleware that sets security-hardening response headers on every
157/// response.
158pub async fn security_headers(request: Request, next: Next) -> Response {
159    let mut response = next.run(request).await;
160    let headers = response.headers_mut();
161    headers.insert(
162        axum::http::header::X_CONTENT_TYPE_OPTIONS,
163        axum::http::HeaderValue::from_static("nosniff"),
164    );
165    headers.insert(
166        axum::http::header::CACHE_CONTROL,
167        axum::http::HeaderValue::from_static("no-store"),
168    );
169    headers.insert(
170        axum::http::header::HeaderName::from_static("x-frame-options"),
171        axum::http::HeaderValue::from_static("DENY"),
172    );
173    // No `Access-Control-Allow-Origin` at all: `null` is not "deny" — it GRANTS the opaque
174    // `null` origin (sandboxed iframes, `data:` / `file:` documents).
175    headers.insert(
176        axum::http::header::HeaderName::from_static("content-security-policy"),
177        axum::http::HeaderValue::from_static("default-src 'none'"),
178    );
179    response
180}