diff --git a/apps/infrastructure/src/auth/jwt.rs b/apps/infrastructure/src/auth/jwt.rs index 85580b7..182f001 100644 --- a/apps/infrastructure/src/auth/jwt.rs +++ b/apps/infrastructure/src/auth/jwt.rs @@ -8,12 +8,27 @@ pub struct JwtClaims { pub sub: String, pub exp: u64, pub iat: u64, + /// Token purpose: `"access"` or `"refresh"`. + /// + /// Prevents an access token from being replayed as a refresh token + /// (which would otherwise extend a short-lived credential into the + /// 7-day refresh window). + #[serde(rename = "typ")] + pub token_type: TokenType, #[serde(skip_serializing_if = "Option::is_none")] pub session_id: Option, } +/// JWT token purpose. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum TokenType { + Access, + Refresh, +} + impl JwtClaims { - pub fn new(sub: String, exp: u64, session_id: Option) -> Self { + pub fn new(sub: String, exp: u64, token_type: TokenType, session_id: Option) -> Self { let iat = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() @@ -22,6 +37,7 @@ impl JwtClaims { sub, exp, iat, + token_type, session_id, } } @@ -48,3 +64,57 @@ pub fn verify_token(secret: &str, token: &str) -> anyhow::Result { let token_data = jsonwebtoken::decode::(token, &key, &validation)?; Ok(token_data.claims) } + +#[cfg(test)] +mod tests { + use super::*; + + fn claims(exp_secs_from_now: u64, token_type: TokenType) -> JwtClaims { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + JwtClaims::new( + "user-1".to_string(), + now + exp_secs_from_now, + token_type, + None, + ) + } + + #[test] + fn access_and_refresh_tokens_roundtrip() { + let secret = "test-secret"; + let access = create_token(secret, claims(3600, TokenType::Access)).unwrap(); + let refresh = create_token(secret, claims(604800, TokenType::Refresh)).unwrap(); + + let acc = verify_token(secret, &access).unwrap(); + assert_eq!(acc.token_type, TokenType::Access); + + let refr = verify_token(secret, &refresh).unwrap(); + assert_eq!(refr.token_type, TokenType::Refresh); + } + + #[test] + fn token_type_is_distinct() { + let secret = "test-secret"; + let access = create_token(secret, claims(3600, TokenType::Access)).unwrap(); + let claims = verify_token(secret, &access).unwrap(); + assert_ne!(claims.token_type, TokenType::Refresh); + } + + #[test] + fn expired_token_is_rejected() { + let secret = "test-secret"; + // exp well in the past (beyond the library's default 60s leeway) → + // verification must fail. + let past = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + .saturating_sub(120); + let expired = JwtClaims::new("user-1".to_string(), past, TokenType::Access, None); + let token = create_token(secret, expired).unwrap(); + assert!(verify_token(secret, &token).is_err()); + } +} diff --git a/apps/interfaces/api/src/error.rs b/apps/interfaces/api/src/error.rs index a79b9f4..6ae1a93 100644 --- a/apps/interfaces/api/src/error.rs +++ b/apps/interfaces/api/src/error.rs @@ -44,6 +44,10 @@ pub enum ApiError { #[error("Conflict: {0}")] Conflict(String), + /// The client has sent too many requests in a given time window. + #[error("Too many requests: {0}")] + TooManyRequests(String), + /// An unexpected internal error occurred. #[error("Internal error: {0}")] Internal(String), @@ -65,6 +69,9 @@ impl IntoResponse for ApiError { ApiError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg.clone()), ApiError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()), ApiError::Conflict(msg) => (StatusCode::CONFLICT, msg.clone()), + ApiError::TooManyRequests(msg) => { + (StatusCode::TOO_MANY_REQUESTS, msg.clone()) + } ApiError::Internal(msg) => { tracing::error!(error = %msg, "Internal server error"); ( diff --git a/apps/interfaces/api/src/handlers/auth.rs b/apps/interfaces/api/src/handlers/auth.rs index bac6057..b7ce363 100644 --- a/apps/interfaces/api/src/handlers/auth.rs +++ b/apps/interfaces/api/src/handlers/auth.rs @@ -24,6 +24,11 @@ use crate::dto::auth::{AuthResponse, LoginRequest, RefreshRequest, RegisterReque use crate::error::ApiError; use crate::state::ApiState; +/// Login/register brute-force protection: 20 attempts per 10-minute window +/// per client IP. +const AUTH_RATE_LIMIT_MAX: u32 = 20; +const AUTH_RATE_LIMIT_WINDOW_SECS: u64 = 600; + /// Build the auth sub-router (`/auth/*`). pub fn router() -> Router> { Router::new() @@ -32,6 +37,30 @@ pub fn router() -> Router> { .route("/refresh", post(refresh_handler)) } +/// Extract a coarse client identity from the request headers (IP via +/// X-Forwarded-For fallback). Used as the rate-limit key. +fn client_id(headers: &axum::http::HeaderMap) -> String { + headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.split(',').next()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "unknown".to_string()) +} + +/// Enforce the auth rate limit, returning `true` if the request is allowed. +fn rate_limited(state: &ApiState, headers: &axum::http::HeaderMap) -> bool { + !state + .auth_rate_limiter + .check_rate_limit( + &client_id(headers), + AUTH_RATE_LIMIT_MAX, + AUTH_RATE_LIMIT_WINDOW_SECS, + ) + .unwrap_or(true) +} + /// POST /auth/login — authenticate and issue JWT tokens. /// /// ## Flow @@ -50,8 +79,16 @@ pub fn router() -> Router> { #[tracing::instrument(skip(state))] pub async fn login_handler( State(state): State>, + headers: axum::http::HeaderMap, Json(req): Json, ) -> Result, ApiError> { + // Rate-limit login attempts (brute-force protection). + if rate_limited(&state, &headers) { + return Err(ApiError::TooManyRequests( + "Too many login attempts, try again later".into(), + )); + } + // Validate input if req.username.is_empty() || req.password.is_empty() { return Err(ApiError::BadRequest( @@ -119,8 +156,16 @@ pub async fn login_handler( #[tracing::instrument(skip(state))] pub async fn register_handler( State(state): State>, + headers: axum::http::HeaderMap, Json(req): Json, ) -> Result, ApiError> { + // Rate-limit registration (abuse protection). + if rate_limited(&state, &headers) { + return Err(ApiError::TooManyRequests( + "Too many registration attempts, try again later".into(), + )); + } + // Validate input if req.username.is_empty() { return Err(ApiError::BadRequest("Username is required".into())); @@ -194,8 +239,16 @@ pub async fn register_handler( #[tracing::instrument(skip(state))] pub async fn refresh_handler( State(state): State>, + headers: axum::http::HeaderMap, Json(req): Json, ) -> Result, ApiError> { + // Rate-limit refresh attempts. + if rate_limited(&state, &headers) { + return Err(ApiError::TooManyRequests( + "Too many requests, try again later".into(), + )); + } + if req.refresh_token.is_empty() { return Err(ApiError::BadRequest("Refresh token is required".into())); } diff --git a/apps/interfaces/api/src/lib.rs b/apps/interfaces/api/src/lib.rs index f95107e..d9492f2 100644 --- a/apps/interfaces/api/src/lib.rs +++ b/apps/interfaces/api/src/lib.rs @@ -56,25 +56,34 @@ pub fn build_router(state: ApiState) -> Router { // CORS layer — permissive for local daemon / development use let cors = CorsLayer::permissive(); - // JWT auth middleware — validates Bearer tokens on all API routes. - // Health and auth endpoints (login/register/refresh) are also - // protected; adjust route ordering or add an allow-list inside the - // middleware if public access is needed. + // JWT auth middleware — protects session/chat routes. + // Auth (login/register/refresh) and health endpoints stay public. let jwt_auth = middleware::auth::JwtAuthLayer::new(shared_state.clone()); // Combine all sub-routers under a versioned prefix Router::new() - .nest("/api/v1", api_v1_router()) + .nest("/api/v1/auth", auth_router()) + .nest("/api/v1/health", health_router()) + .nest( + "/api/v1", + protected_router().layer(jwt_auth), + ) .layer(cors) - .layer(jwt_auth) .with_state(shared_state) } -/// Version 1 API sub-router. -/// -/// Groups all resource routes under `/api/v1/*`. -fn api_v1_router() -> Router> { - use handlers::{auth, chat, conversations, health, sessions}; +/// Auth + health sub-routers — publicly accessible (no JWT required). +fn auth_router() -> Router> { + handlers::auth::router() +} + +fn health_router() -> Router> { + Router::new().route("/", axum::routing::get(handlers::health::health)) +} + +/// Protected sub-router — sessions + chat, guarded by JWT auth layer. +fn protected_router() -> Router> { + use handlers::{chat, conversations, sessions}; // Sessions router combines session CRUD + nested conversations let sessions_router = Router::new() @@ -96,8 +105,6 @@ fn api_v1_router() -> Router> { ); Router::new() - .route("/health", axum::routing::get(health::health)) - .nest("/auth", auth::router()) .nest("/sessions", sessions_router) .nest("/chat", chat::router()) } diff --git a/apps/interfaces/api/src/middleware/auth.rs b/apps/interfaces/api/src/middleware/auth.rs index 4bd4ffc..5c3ecca 100644 --- a/apps/interfaces/api/src/middleware/auth.rs +++ b/apps/interfaces/api/src/middleware/auth.rs @@ -98,6 +98,21 @@ where if let Some(token) = auth_value.strip_prefix("Bearer ") { match zesdex_infrastructure::auth::jwt::verify_token(&secret, token) { Ok(claims) => { + // Reject refresh tokens on protected routes — only access + // tokens are acceptable here. + if claims.token_type + != zesdex_infrastructure::auth::jwt::TokenType::Access + { + let response = ( + StatusCode::UNAUTHORIZED, + Json(json!({ + "error": "Invalid token", + "detail": "refresh tokens are not accepted on protected routes" + })), + ) + .into_response(); + return Box::pin(async move { Ok(response) }); + } // Inject claims as extension for downstream handlers let mut req = req; req.extensions_mut().insert(JwtClaims { diff --git a/apps/interfaces/api/src/state.rs b/apps/interfaces/api/src/state.rs index 1b08091..3194285 100644 --- a/apps/interfaces/api/src/state.rs +++ b/apps/interfaces/api/src/state.rs @@ -94,7 +94,7 @@ impl JwtTokenService { impl TokenService for JwtTokenService { /// Generate an access + refresh token pair for the given subject. fn generate_tokens(&self, sub: &str) -> anyhow::Result<(String, String)> { - use zesdex_infrastructure::auth::jwt::{create_token, JwtClaims}; + use zesdex_infrastructure::auth::jwt::{create_token, JwtClaims, TokenType}; let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -102,12 +102,21 @@ impl TokenService for JwtTokenService { .as_secs(); // Access token - let access_claims = JwtClaims::new(sub.to_string(), now + self.access_token_expiry_secs, None); + let access_claims = JwtClaims::new( + sub.to_string(), + now + self.access_token_expiry_secs, + TokenType::Access, + None, + ); let access_token = create_token(&self.secret, access_claims)?; // Refresh token (longer-lived) - let refresh_claims = - JwtClaims::new(sub.to_string(), now + self.refresh_token_expiry_secs, None); + let refresh_claims = JwtClaims::new( + sub.to_string(), + now + self.refresh_token_expiry_secs, + TokenType::Refresh, + None, + ); let refresh_token = create_token(&self.secret, refresh_claims)?; Ok((access_token, refresh_token)) @@ -123,14 +132,17 @@ impl TokenService for JwtTokenService { /// Verify a refresh token and return the subject claim. /// - /// Delegates to the same JWT verification function as access tokens; - /// the signature algorithm and secret are shared. Expiry validation - /// is handled by the JWT library against the `exp` claim embedded - /// in the token payload. + /// Enforces that the presented token is a **refresh** token (`typ = + /// "refresh"`) — an access token presented here is rejected, closing + /// the replay-window escalation where a stolen 1-hour access token + /// could otherwise be exchanged for a fresh 7-day credential. fn verify_refresh_token(&self, token: &str) -> anyhow::Result { - use zesdex_infrastructure::auth::jwt::verify_token; + use zesdex_infrastructure::auth::jwt::{verify_token, TokenType}; let claims = verify_token(&self.secret, token)?; + if claims.token_type != TokenType::Refresh { + anyhow::bail!("token is not a refresh token"); + } Ok(claims.sub) } } @@ -193,6 +205,9 @@ pub struct ApiState { /// HS256 JWT token generation and verification. pub token_service: JwtTokenService, + /// Shared sliding-window limiter for auth endpoints (login/register/refresh). + pub auth_rate_limiter: zesdex_infrastructure::middleware::rate_limit::RateLimiter, + /// LLM provider client for chat completions. pub llm_client: zesdex_infrastructure::llm::provider::LlmClient, } @@ -272,6 +287,7 @@ impl ApiState { zesdex_application::cms::MemoryServiceImpl::new(memory_repo, memory_dir); let token_service = JwtTokenService::new(&jwt_secret); + let auth_rate_limiter = zesdex_infrastructure::middleware::rate_limit::RateLimiter::new(); let llm_client = zesdex_infrastructure::llm::provider::LlmClient::new( llm_api_key.into(), llm_model.into(), @@ -287,6 +303,7 @@ impl ApiState { memory_service, password_service: Argon2PasswordService, token_service, + auth_rate_limiter, llm_client, } } diff --git a/apps/interfaces/ws/src/lib.rs b/apps/interfaces/ws/src/lib.rs index e405897..47bec4d 100644 --- a/apps/interfaces/ws/src/lib.rs +++ b/apps/interfaces/ws/src/lib.rs @@ -2,15 +2,25 @@ //! //! Enables web clients and other WS-capable consumers to connect //! and participate in sessions. Built on Axum's WebSocket support. +//! +//! # Security +//! +//! The WS endpoint accepts an optional `?token=` query parameter. When a +//! `ZEESDEX_WS_TOKEN` env var is set, connections MUST present a matching +//! token — otherwise the connection is rejected. This prevents the endpoint +//! from being used as an open LLM proxy (anyone who can reach the port would +//! otherwise run prompts at the server's API cost). use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; +use axum::extract::Query; use axum::response::IntoResponse; use axum::routing::get; use axum::Router; use futures_util::stream::StreamExt; use futures_util::SinkExt; +use serde::Deserialize; use std::sync::Arc; -use tracing::info; +use tracing::{info, warn}; /// Shared application state for the WS server. pub struct WsState { @@ -18,6 +28,12 @@ pub struct WsState { pub session_id: Option, } +/// Query parameters accepted on the `/ws` upgrade. +#[derive(Debug, Deserialize)] +struct WsQuery { + token: Option, +} + /// Build the WebSocket router. pub fn build_router(state: Arc) -> Router { Router::new() @@ -28,9 +44,18 @@ pub fn build_router(state: Arc) -> Router { /// WebSocket upgrade handler. async fn ws_handler( ws: WebSocketUpgrade, + Query(query): Query, axum::extract::State(state): axum::extract::State>, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| handle_socket(socket, state)) + let configured = std::env::var("ZESDEX_WS_TOKEN").ok().filter(|s| !s.is_empty()); + match configured { + Some(expected) if query.token.as_deref() != Some(expected.as_str()) => { + warn!("rejecting WS connection: missing/invalid token"); + // 401 Unauthorized — client did not present the required token. + (axum::http::StatusCode::UNAUTHORIZED, "missing or invalid token").into_response() + } + _ => ws.on_upgrade(move |socket| handle_socket(socket, state)), + } } /// Handle an established WebSocket connection.