mentedb_server/
auth.rs

1//! JWT authentication middleware and token management.
2
3use std::sync::Arc;
4
5use axum::Json;
6use axum::body::Body;
7use axum::extract::State;
8use axum::http::{Request, StatusCode};
9use axum::middleware::Next;
10use axum::response::{IntoResponse, Response};
11use http_body_util::BodyExt;
12use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
13use serde::{Deserialize, Serialize};
14use serde_json::json;
15
16use crate::error::ApiError;
17use crate::state::AppState;
18
19#[derive(Debug, Clone)]
20pub struct AuthenticatedAgent {
21    pub agent_id: String,
22    pub admin: bool,
23}
24
25/// JWT claims embedded in every token.
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct Claims {
28    /// The agent ID this token was issued for.
29    pub agent_id: String,
30    /// Whether this token grants admin privileges.
31    #[serde(default)]
32    pub admin: bool,
33    /// Token expiration as a Unix timestamp.
34    pub exp: usize,
35}
36
37/// Request body for the token generation endpoint.
38#[derive(Deserialize)]
39pub struct TokenRequest {
40    /// The agent ID to issue a token for.
41    pub agent_id: String,
42    /// How many hours until the token expires (default: 24).
43    #[serde(default = "default_expiry_hours")]
44    pub expiry_hours: u64,
45}
46
47fn default_expiry_hours() -> u64 {
48    24
49}
50
51/// Create a signed JWT for the given agent.
52pub fn create_token(secret: &str, agent_id: &str, admin: bool, expiry_hours: u64) -> String {
53    let exp = jsonwebtoken::get_current_timestamp() as usize + (expiry_hours as usize * 3600);
54    let claims = Claims {
55        agent_id: agent_id.to_string(),
56        admin,
57        exp,
58    };
59    encode(
60        &Header::default(),
61        &claims,
62        &EncodingKey::from_secret(secret.as_bytes()),
63    )
64    .expect("JWT encoding should not fail with valid inputs")
65}
66
67/// Validate a JWT and return the embedded claims.
68pub fn validate_token(secret: &str, token: &str) -> Result<Claims, ApiError> {
69    decode::<Claims>(
70        token,
71        &DecodingKey::from_secret(secret.as_bytes()),
72        &Validation::default(),
73    )
74    .map(|data| data.claims)
75    .map_err(|e| ApiError::Unauthorized(format!("invalid token: {e}")))
76}
77
78/// Handler: POST /v1/auth/token: generate a new JWT.
79fn extract_admin_key(request: &Request<Body>) -> Option<String> {
80    if let Some(v) = request
81        .headers()
82        .get("x-api-key")
83        .and_then(|v| v.to_str().ok())
84    {
85        return Some(v.to_string());
86    }
87    request
88        .headers()
89        .get("authorization")
90        .and_then(|v| v.to_str().ok())
91        .and_then(|v| v.strip_prefix("Bearer "))
92        .map(String::from)
93}
94
95pub async fn generate_token(
96    State(state): State<Arc<AppState>>,
97    request: Request<Body>,
98) -> Result<impl IntoResponse, ApiError> {
99    let secret = state.jwt_secret.as_deref().ok_or_else(|| {
100        ApiError::BadRequest("auth is disabled (no jwt-secret configured)".into())
101    })?;
102    match &state.admin_key {
103        None => {
104            return Err(ApiError::Forbidden(
105                "token endpoint disabled: no admin key configured".into(),
106            ));
107        }
108        Some(expected) => {
109            let provided = extract_admin_key(&request)
110                .ok_or_else(|| ApiError::Unauthorized("admin key required".into()))?;
111            if provided != *expected {
112                return Err(ApiError::Forbidden("invalid admin key".into()));
113            }
114        }
115    }
116    let body_bytes = request
117        .into_body()
118        .collect()
119        .await
120        .map_err(|e| ApiError::BadRequest(format!("failed to read body: {e}")))?
121        .to_bytes();
122    let req: TokenRequest = serde_json::from_slice(&body_bytes)
123        .map_err(|e| ApiError::BadRequest(format!("invalid JSON: {e}")))?;
124    let token = create_token(secret, &req.agent_id, false, req.expiry_hours);
125    Ok(Json(json!({ "token": token, "agent_id": req.agent_id })))
126}
127
128/// Middleware that enforces JWT auth when a secret is configured.
129pub async fn auth_middleware(
130    State(state): State<Arc<AppState>>,
131    mut request: Request<Body>,
132    next: Next,
133) -> Response {
134    let secret = match &state.jwt_secret {
135        Some(s) => s.clone(),
136        // No secret configured means development mode; skip auth.
137        None => return next.run(request).await,
138    };
139
140    // Allow health and token endpoints without auth.
141    let path = request.uri().path();
142    if path == "/v1/health" || path == "/v1/auth/token" {
143        return next.run(request).await;
144    }
145
146    let auth_header = request
147        .headers()
148        .get("authorization")
149        .and_then(|v| v.to_str().ok())
150        .map(String::from);
151
152    let token = match auth_header.as_deref() {
153        Some(h) if h.starts_with("Bearer ") => &h[7..],
154        _ => {
155            return (
156                StatusCode::UNAUTHORIZED,
157                Json(json!({ "error": "missing or invalid Authorization header" })),
158            )
159                .into_response();
160        }
161    };
162
163    match validate_token(&secret, token) {
164        Ok(claims) => {
165            request.extensions_mut().insert(AuthenticatedAgent {
166                agent_id: claims.agent_id,
167                admin: claims.admin,
168            });
169            next.run(request).await
170        }
171        Err(e) => e.into_response(),
172    }
173}
mentedb_server/auth.rs

mentedb_server/
auth.rs
