Skip to main content

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
95/// Authorize an admin request from request headers against the configured admin
96/// key. Errors if no admin key is configured (admin endpoints are disabled) or the
97/// provided `x-api-key` / bearer does not match. Used by the /v1/admin endpoints.
98pub fn admin_authorized(
99    headers: &axum::http::HeaderMap,
100    admin_key: &Option<String>,
101) -> Result<(), ApiError> {
102    let expected = admin_key
103        .as_deref()
104        .ok_or_else(|| ApiError::Unauthorized("admin endpoints disabled: no admin key".into()))?;
105    let provided = headers
106        .get("x-api-key")
107        .and_then(|v| v.to_str().ok())
108        .or_else(|| {
109            headers
110                .get("authorization")
111                .and_then(|v| v.to_str().ok())
112                .and_then(|v| v.strip_prefix("Bearer "))
113        })
114        .ok_or_else(|| ApiError::Unauthorized("admin key required".into()))?;
115    if provided != expected {
116        return Err(ApiError::Forbidden("invalid admin key".into()));
117    }
118    Ok(())
119}
120
121pub async fn generate_token(
122    State(state): State<Arc<AppState>>,
123    request: Request<Body>,
124) -> Result<impl IntoResponse, ApiError> {
125    let secret = state.jwt_secret.as_deref().ok_or_else(|| {
126        ApiError::BadRequest("auth is disabled (no jwt-secret configured)".into())
127    })?;
128    match &state.admin_key {
129        None => {
130            return Err(ApiError::Forbidden(
131                "token endpoint disabled: no admin key configured".into(),
132            ));
133        }
134        Some(expected) => {
135            let provided = extract_admin_key(&request)
136                .ok_or_else(|| ApiError::Unauthorized("admin key required".into()))?;
137            if provided != *expected {
138                return Err(ApiError::Forbidden("invalid admin key".into()));
139            }
140        }
141    }
142    let body_bytes = request
143        .into_body()
144        .collect()
145        .await
146        .map_err(|e| ApiError::BadRequest(format!("failed to read body: {e}")))?
147        .to_bytes();
148    let req: TokenRequest = serde_json::from_slice(&body_bytes)
149        .map_err(|e| ApiError::BadRequest(format!("invalid JSON: {e}")))?;
150    let token = create_token(secret, &req.agent_id, false, req.expiry_hours);
151    Ok(Json(json!({ "token": token, "agent_id": req.agent_id })))
152}
153
154/// Middleware that enforces JWT auth when a secret is configured.
155pub async fn auth_middleware(
156    State(state): State<Arc<AppState>>,
157    mut request: Request<Body>,
158    next: Next,
159) -> Response {
160    let secret = match &state.jwt_secret {
161        Some(s) => s.clone(),
162        // No secret configured means development mode; skip auth.
163        None => return next.run(request).await,
164    };
165
166    // Allow health, token, metrics, and peer-to-peer gossip endpoints without auth.
167    let path = request.uri().path();
168    // Admin endpoints are gated by the admin key inside the handler, not the JWT,
169    // so the bundled console (which sends x-api-key) can reach them.
170    if path == "/v1/health"
171        || path == "/v1/auth/token"
172        || path == "/metrics"
173        || path == "/console"
174        || path.starts_with("/v1/admin")
175        || path == crate::cluster::GOSSIP_PATH
176    {
177        return next.run(request).await;
178    }
179
180    let auth_header = request
181        .headers()
182        .get("authorization")
183        .and_then(|v| v.to_str().ok())
184        .map(String::from);
185
186    let token = match auth_header.as_deref() {
187        Some(h) if h.starts_with("Bearer ") => &h[7..],
188        _ => {
189            return (
190                StatusCode::UNAUTHORIZED,
191                Json(json!({ "error": "missing or invalid Authorization header" })),
192            )
193                .into_response();
194        }
195    };
196
197    match validate_token(&secret, token) {
198        Ok(claims) => {
199            request.extensions_mut().insert(AuthenticatedAgent {
200                agent_id: claims.agent_id,
201                admin: claims.admin,
202            });
203            next.run(request).await
204        }
205        Err(e) => e.into_response(),
206    }
207}