Skip to main content

crabllm_proxy/
auth.rs

1use crate::AppState;
2use axum::{
3    Json,
4    extract::{Request, State},
5    http::StatusCode,
6    middleware::Next,
7    response::{IntoResponse, Response},
8};
9use crabllm_core::{ApiError, Provider, Storage};
10
11/// Wrapper for the authenticated key name, inserted into request extensions.
12#[derive(Clone, Debug)]
13pub struct KeyName(pub Option<String>);
14
15/// Auth middleware: validates Bearer token against configured virtual keys.
16/// Skips auth only when no admin_token is configured AND key_map is empty.
17/// Inserts `KeyName` into request extensions for downstream handlers.
18pub async fn auth<S: Storage + 'static, P: Provider + 'static>(
19    State(state): State<AppState<S, P>>,
20    mut request: Request,
21    next: Next,
22) -> Response {
23    // Skip auth when key management is disabled and no keys exist.
24    if state.config.admin_token.is_none()
25        && state
26            .key_map
27            .read()
28            .unwrap_or_else(|e| e.into_inner())
29            .is_empty()
30    {
31        request.extensions_mut().insert(KeyName(None));
32        return next.run(request).await;
33    }
34
35    // Accept either OpenAI-style `Authorization: Bearer <key>` or Anthropic-style
36    // `x-api-key: <key>`. Both map to the same virtual-key lookup.
37    let headers = request.headers();
38    let bearer = headers
39        .get("authorization")
40        .and_then(|v| v.to_str().ok())
41        .and_then(|h| h.strip_prefix("Bearer "));
42    let x_api_key = headers.get("x-api-key").and_then(|v| v.to_str().ok());
43
44    let token = match bearer.or(x_api_key) {
45        Some(t) => t,
46        None => {
47            return (
48                StatusCode::UNAUTHORIZED,
49                Json(ApiError::new(
50                    "missing Authorization or x-api-key header",
51                    "authentication_error",
52                )),
53            )
54                .into_response();
55        }
56    };
57
58    let key_name = state
59        .key_map
60        .read()
61        .unwrap_or_else(|e| e.into_inner())
62        .get(token)
63        .cloned();
64
65    let Some(key_name) = key_name else {
66        return (
67            StatusCode::UNAUTHORIZED,
68            Json(ApiError::new("invalid API key", "authentication_error")),
69        )
70            .into_response();
71    };
72
73    request.extensions_mut().insert(KeyName(Some(key_name)));
74
75    next.run(request).await
76}