1use 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#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct Claims {
28 pub agent_id: String,
30 #[serde(default)]
32 pub admin: bool,
33 pub exp: usize,
35}
36
37#[derive(Deserialize)]
39pub struct TokenRequest {
40 pub agent_id: String,
42 #[serde(default = "default_expiry_hours")]
44 pub expiry_hours: u64,
45}
46
47fn default_expiry_hours() -> u64 {
48 24
49}
50
51pub 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
67pub 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
78fn 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 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
154pub 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 None => return next.run(request).await,
164 };
165
166 let path = request.uri().path();
168 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}