use crate::error::ApiErrorResponse;
use crate::router::{MATCHED_KB, MATCHED_KB_OBJECT, MATCHED_KB_SEARCH, MATCHED_KBS};
use crate::state::AppState;
use axum::body::Body;
use axum::extract::{MatchedPath, State};
use axum::http::{Method, Request, header::AUTHORIZATION};
use axum::middleware::Next;
use axum::response::Response;
use notedthat_core::{Principal, extract_bearer_from_header, verify_bearer_token};
use tower_http::request_id::RequestId;
const ANONYMOUS_REACHABLE: &[(&Method, &str)] = &[
(&Method::GET, MATCHED_KBS),
(&Method::HEAD, MATCHED_KBS),
(&Method::GET, MATCHED_KB),
(&Method::HEAD, MATCHED_KB),
(&Method::GET, MATCHED_KB_OBJECT),
(&Method::HEAD, MATCHED_KB_OBJECT),
(&Method::POST, MATCHED_KB_SEARCH),
];
pub async fn auth_middleware(
State(state): State<AppState>,
mut req: Request<Body>,
next: Next,
) -> Result<Response, ApiErrorResponse> {
let request_id = extract_request_id(&req);
let principal = resolve_principal(req.headers(), &state.bearer_token)
.map_err(|CredentialRefused| ApiErrorResponse::unauthorized(request_id.clone()))?;
if principal == Principal::Anyone && !anonymous_may_reach(&req) {
return Err(ApiErrorResponse::unauthorized(request_id));
}
req.extensions_mut().insert(principal);
Ok(next.run(req).await)
}
fn anonymous_may_reach<B>(req: &Request<B>) -> bool {
let Some(matched) = req.extensions().get::<MatchedPath>() else {
return false;
};
let matched = matched.as_str();
ANONYMOUS_REACHABLE
.iter()
.any(|(method, route)| *method == req.method() && *route == matched)
}
pub fn resolve_principal(
headers: &axum::http::HeaderMap,
expected_token: &str,
) -> Result<Principal, CredentialRefused> {
let mut values = headers.get_all(AUTHORIZATION).iter();
let Some(header) = values.next() else {
return Ok(Principal::Anyone);
};
if values.next().is_some() {
return Err(CredentialRefused);
}
header
.to_str()
.ok()
.and_then(extract_bearer_from_header)
.filter(|token| verify_bearer_token(token, expected_token))
.map(|_| Principal::SignedIn)
.ok_or(CredentialRefused)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CredentialRefused;
pub fn principal<B>(req: &Request<B>) -> Principal {
req.extensions()
.get::<Principal>()
.copied()
.unwrap_or(Principal::Anyone)
}
pub use notedthat_core::is_internal_path;
pub fn extract_request_id<B>(req: &Request<B>) -> String {
req.extensions()
.get::<RequestId>()
.and_then(|r| r.header_value().to_str().ok())
.map_or_else(
|| {
tracing::warn!("request_id missing from Extensions — generating fallback");
uuid::Uuid::now_v7().to_string()
},
str::to_string,
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing::InMemoryStorage;
use axum::middleware::from_fn_with_state;
use axum::response::IntoResponse;
use axum::routing::get;
use axum::{Router, body::Body, http::StatusCode};
use std::collections::BTreeMap;
use std::sync::Arc;
use tower::util::ServiceExt;
fn test_state(token: &str) -> AppState {
let (indexer_tx, _rx) = tokio::sync::mpsc::channel(1024);
AppState {
storage: Arc::new(InMemoryStorage::default()),
declared_kbs: Arc::new(BTreeMap::new()),
access_policies: Arc::new(BTreeMap::new()),
bearer_token: Arc::new(token.to_string()),
max_body_size: 16 * 1024 * 1024,
max_patchable_size: 16 * 1024 * 1024,
indexer_tx,
searcher: Arc::new(crate::testing::NoopSearcher),
}
}
fn app(token: &str) -> Router {
let state = test_state(token);
Router::new()
.route("/protected", get(|| async { "secret".into_response() }))
.layer(from_fn_with_state(state.clone(), auth_middleware))
.with_state(state)
}
#[tokio::test]
async fn no_path_is_exempt_once_the_request_reaches_this_layer() {
for uri in ["/healthz", "/readyz", "/llms.txt", "/protected"] {
let resp = app("my-token")
.oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(
resp.status(),
StatusCode::UNAUTHORIZED,
"{uri} must not be exempted by the auth layer itself"
);
}
}
#[tokio::test]
async fn test_rejects_missing_auth() {
let resp = app("my-token")
.oneshot(
Request::builder()
.uri("/protected")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn test_rejects_wrong_token() {
let resp = app("real-token")
.oneshot(
Request::builder()
.uri("/protected")
.header("authorization", "Bearer wrong-token")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn test_accepts_correct_token() {
let resp = app("my-token")
.oneshot(
Request::builder()
.uri("/protected")
.header("authorization", "Bearer my-token")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_accepts_lowercase_bearer_scheme() {
let resp = app("my-token")
.oneshot(
Request::builder()
.uri("/protected")
.header("authorization", "bearer my-token")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
}