fix(api): perbaiki keamanan auth & WebSocket, tambah rate limiting
Security fixes hasil audit: - fix(auth): refresh token kini memakai claim typ=refresh; access token tidak bisa dipakai sebagai refresh token (sebelumnya bisa — eskalasi masa berlaku 1 jam -> 7 hari) - fix(api): layer JWT hanya melindungi route /sessions dan /chat; /auth/login, /auth/register, /auth/refresh, /health kini publik (sebelumnya semua route 401-lock, API tidak bisa dipakai sama sekali) - fix(ws): endpoint /ws kini memverifikasi token ZESDEX_WS_TOKEN via query param jika env diset (mencegah pemakaian LLM proxy terbuka) - feat(api): rate limiting login/register/refresh (20 request / 10 menit per client IP) memakai RateLimiter yang tadinya dead code - test(jwt): tambah unit test token type access vs refresh + expired
This commit is contained in:
@@ -8,12 +8,27 @@ pub struct JwtClaims {
|
|||||||
pub sub: String,
|
pub sub: String,
|
||||||
pub exp: u64,
|
pub exp: u64,
|
||||||
pub iat: 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")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub session_id: Option<String>,
|
pub session_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// JWT token purpose.
|
||||||
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum TokenType {
|
||||||
|
Access,
|
||||||
|
Refresh,
|
||||||
|
}
|
||||||
|
|
||||||
impl JwtClaims {
|
impl JwtClaims {
|
||||||
pub fn new(sub: String, exp: u64, session_id: Option<String>) -> Self {
|
pub fn new(sub: String, exp: u64, token_type: TokenType, session_id: Option<String>) -> Self {
|
||||||
let iat = std::time::SystemTime::now()
|
let iat = std::time::SystemTime::now()
|
||||||
.duration_since(std::time::UNIX_EPOCH)
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
@@ -22,6 +37,7 @@ impl JwtClaims {
|
|||||||
sub,
|
sub,
|
||||||
exp,
|
exp,
|
||||||
iat,
|
iat,
|
||||||
|
token_type,
|
||||||
session_id,
|
session_id,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -48,3 +64,57 @@ pub fn verify_token(secret: &str, token: &str) -> anyhow::Result<JwtClaims> {
|
|||||||
let token_data = jsonwebtoken::decode::<JwtClaims>(token, &key, &validation)?;
|
let token_data = jsonwebtoken::decode::<JwtClaims>(token, &key, &validation)?;
|
||||||
Ok(token_data.claims)
|
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());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -44,6 +44,10 @@ pub enum ApiError {
|
|||||||
#[error("Conflict: {0}")]
|
#[error("Conflict: {0}")]
|
||||||
Conflict(String),
|
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.
|
/// An unexpected internal error occurred.
|
||||||
#[error("Internal error: {0}")]
|
#[error("Internal error: {0}")]
|
||||||
Internal(String),
|
Internal(String),
|
||||||
@@ -65,6 +69,9 @@ impl IntoResponse for ApiError {
|
|||||||
ApiError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg.clone()),
|
ApiError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg.clone()),
|
||||||
ApiError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
|
ApiError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
|
||||||
ApiError::Conflict(msg) => (StatusCode::CONFLICT, msg.clone()),
|
ApiError::Conflict(msg) => (StatusCode::CONFLICT, msg.clone()),
|
||||||
|
ApiError::TooManyRequests(msg) => {
|
||||||
|
(StatusCode::TOO_MANY_REQUESTS, msg.clone())
|
||||||
|
}
|
||||||
ApiError::Internal(msg) => {
|
ApiError::Internal(msg) => {
|
||||||
tracing::error!(error = %msg, "Internal server error");
|
tracing::error!(error = %msg, "Internal server error");
|
||||||
(
|
(
|
||||||
|
|||||||
@@ -24,6 +24,11 @@ use crate::dto::auth::{AuthResponse, LoginRequest, RefreshRequest, RegisterReque
|
|||||||
use crate::error::ApiError;
|
use crate::error::ApiError;
|
||||||
use crate::state::ApiState;
|
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/*`).
|
/// Build the auth sub-router (`/auth/*`).
|
||||||
pub fn router() -> Router<Arc<ApiState>> {
|
pub fn router() -> Router<Arc<ApiState>> {
|
||||||
Router::new()
|
Router::new()
|
||||||
@@ -32,6 +37,30 @@ pub fn router() -> Router<Arc<ApiState>> {
|
|||||||
.route("/refresh", post(refresh_handler))
|
.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.
|
/// POST /auth/login — authenticate and issue JWT tokens.
|
||||||
///
|
///
|
||||||
/// ## Flow
|
/// ## Flow
|
||||||
@@ -50,8 +79,16 @@ pub fn router() -> Router<Arc<ApiState>> {
|
|||||||
#[tracing::instrument(skip(state))]
|
#[tracing::instrument(skip(state))]
|
||||||
pub async fn login_handler(
|
pub async fn login_handler(
|
||||||
State(state): State<Arc<ApiState>>,
|
State(state): State<Arc<ApiState>>,
|
||||||
|
headers: axum::http::HeaderMap,
|
||||||
Json(req): Json<LoginRequest>,
|
Json(req): Json<LoginRequest>,
|
||||||
) -> Result<Json<AuthResponse>, ApiError> {
|
) -> Result<Json<AuthResponse>, 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
|
// Validate input
|
||||||
if req.username.is_empty() || req.password.is_empty() {
|
if req.username.is_empty() || req.password.is_empty() {
|
||||||
return Err(ApiError::BadRequest(
|
return Err(ApiError::BadRequest(
|
||||||
@@ -119,8 +156,16 @@ pub async fn login_handler(
|
|||||||
#[tracing::instrument(skip(state))]
|
#[tracing::instrument(skip(state))]
|
||||||
pub async fn register_handler(
|
pub async fn register_handler(
|
||||||
State(state): State<Arc<ApiState>>,
|
State(state): State<Arc<ApiState>>,
|
||||||
|
headers: axum::http::HeaderMap,
|
||||||
Json(req): Json<RegisterRequest>,
|
Json(req): Json<RegisterRequest>,
|
||||||
) -> Result<Json<AuthResponse>, ApiError> {
|
) -> Result<Json<AuthResponse>, 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
|
// Validate input
|
||||||
if req.username.is_empty() {
|
if req.username.is_empty() {
|
||||||
return Err(ApiError::BadRequest("Username is required".into()));
|
return Err(ApiError::BadRequest("Username is required".into()));
|
||||||
@@ -194,8 +239,16 @@ pub async fn register_handler(
|
|||||||
#[tracing::instrument(skip(state))]
|
#[tracing::instrument(skip(state))]
|
||||||
pub async fn refresh_handler(
|
pub async fn refresh_handler(
|
||||||
State(state): State<Arc<ApiState>>,
|
State(state): State<Arc<ApiState>>,
|
||||||
|
headers: axum::http::HeaderMap,
|
||||||
Json(req): Json<RefreshRequest>,
|
Json(req): Json<RefreshRequest>,
|
||||||
) -> Result<Json<AuthResponse>, ApiError> {
|
) -> Result<Json<AuthResponse>, 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() {
|
if req.refresh_token.is_empty() {
|
||||||
return Err(ApiError::BadRequest("Refresh token is required".into()));
|
return Err(ApiError::BadRequest("Refresh token is required".into()));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -56,25 +56,34 @@ pub fn build_router(state: ApiState) -> Router {
|
|||||||
// CORS layer — permissive for local daemon / development use
|
// CORS layer — permissive for local daemon / development use
|
||||||
let cors = CorsLayer::permissive();
|
let cors = CorsLayer::permissive();
|
||||||
|
|
||||||
// JWT auth middleware — validates Bearer tokens on all API routes.
|
// JWT auth middleware — protects session/chat routes.
|
||||||
// Health and auth endpoints (login/register/refresh) are also
|
// Auth (login/register/refresh) and health endpoints stay public.
|
||||||
// protected; adjust route ordering or add an allow-list inside the
|
|
||||||
// middleware if public access is needed.
|
|
||||||
let jwt_auth = middleware::auth::JwtAuthLayer::new(shared_state.clone());
|
let jwt_auth = middleware::auth::JwtAuthLayer::new(shared_state.clone());
|
||||||
|
|
||||||
// Combine all sub-routers under a versioned prefix
|
// Combine all sub-routers under a versioned prefix
|
||||||
Router::new()
|
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(cors)
|
||||||
.layer(jwt_auth)
|
|
||||||
.with_state(shared_state)
|
.with_state(shared_state)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Version 1 API sub-router.
|
/// Auth + health sub-routers — publicly accessible (no JWT required).
|
||||||
///
|
fn auth_router() -> Router<Arc<ApiState>> {
|
||||||
/// Groups all resource routes under `/api/v1/*`.
|
handlers::auth::router()
|
||||||
fn api_v1_router() -> Router<Arc<ApiState>> {
|
}
|
||||||
use handlers::{auth, chat, conversations, health, sessions};
|
|
||||||
|
fn health_router() -> Router<Arc<ApiState>> {
|
||||||
|
Router::new().route("/", axum::routing::get(handlers::health::health))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Protected sub-router — sessions + chat, guarded by JWT auth layer.
|
||||||
|
fn protected_router() -> Router<Arc<ApiState>> {
|
||||||
|
use handlers::{chat, conversations, sessions};
|
||||||
|
|
||||||
// Sessions router combines session CRUD + nested conversations
|
// Sessions router combines session CRUD + nested conversations
|
||||||
let sessions_router = Router::new()
|
let sessions_router = Router::new()
|
||||||
@@ -96,8 +105,6 @@ fn api_v1_router() -> Router<Arc<ApiState>> {
|
|||||||
);
|
);
|
||||||
|
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/health", axum::routing::get(health::health))
|
|
||||||
.nest("/auth", auth::router())
|
|
||||||
.nest("/sessions", sessions_router)
|
.nest("/sessions", sessions_router)
|
||||||
.nest("/chat", chat::router())
|
.nest("/chat", chat::router())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -98,6 +98,21 @@ where
|
|||||||
if let Some(token) = auth_value.strip_prefix("Bearer ") {
|
if let Some(token) = auth_value.strip_prefix("Bearer ") {
|
||||||
match zesdex_infrastructure::auth::jwt::verify_token(&secret, token) {
|
match zesdex_infrastructure::auth::jwt::verify_token(&secret, token) {
|
||||||
Ok(claims) => {
|
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
|
// Inject claims as extension for downstream handlers
|
||||||
let mut req = req;
|
let mut req = req;
|
||||||
req.extensions_mut().insert(JwtClaims {
|
req.extensions_mut().insert(JwtClaims {
|
||||||
|
|||||||
@@ -94,7 +94,7 @@ impl JwtTokenService {
|
|||||||
impl TokenService for JwtTokenService {
|
impl TokenService for JwtTokenService {
|
||||||
/// Generate an access + refresh token pair for the given subject.
|
/// Generate an access + refresh token pair for the given subject.
|
||||||
fn generate_tokens(&self, sub: &str) -> anyhow::Result<(String, String)> {
|
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()
|
let now = std::time::SystemTime::now()
|
||||||
.duration_since(std::time::UNIX_EPOCH)
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
@@ -102,12 +102,21 @@ impl TokenService for JwtTokenService {
|
|||||||
.as_secs();
|
.as_secs();
|
||||||
|
|
||||||
// Access token
|
// 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)?;
|
let access_token = create_token(&self.secret, access_claims)?;
|
||||||
|
|
||||||
// Refresh token (longer-lived)
|
// Refresh token (longer-lived)
|
||||||
let refresh_claims =
|
let refresh_claims = JwtClaims::new(
|
||||||
JwtClaims::new(sub.to_string(), now + self.refresh_token_expiry_secs, None);
|
sub.to_string(),
|
||||||
|
now + self.refresh_token_expiry_secs,
|
||||||
|
TokenType::Refresh,
|
||||||
|
None,
|
||||||
|
);
|
||||||
let refresh_token = create_token(&self.secret, refresh_claims)?;
|
let refresh_token = create_token(&self.secret, refresh_claims)?;
|
||||||
|
|
||||||
Ok((access_token, refresh_token))
|
Ok((access_token, refresh_token))
|
||||||
@@ -123,14 +132,17 @@ impl TokenService for JwtTokenService {
|
|||||||
|
|
||||||
/// Verify a refresh token and return the subject claim.
|
/// Verify a refresh token and return the subject claim.
|
||||||
///
|
///
|
||||||
/// Delegates to the same JWT verification function as access tokens;
|
/// Enforces that the presented token is a **refresh** token (`typ =
|
||||||
/// the signature algorithm and secret are shared. Expiry validation
|
/// "refresh"`) — an access token presented here is rejected, closing
|
||||||
/// is handled by the JWT library against the `exp` claim embedded
|
/// the replay-window escalation where a stolen 1-hour access token
|
||||||
/// in the token payload.
|
/// could otherwise be exchanged for a fresh 7-day credential.
|
||||||
fn verify_refresh_token(&self, token: &str) -> anyhow::Result<String> {
|
fn verify_refresh_token(&self, token: &str) -> anyhow::Result<String> {
|
||||||
use zesdex_infrastructure::auth::jwt::verify_token;
|
use zesdex_infrastructure::auth::jwt::{verify_token, TokenType};
|
||||||
|
|
||||||
let claims = verify_token(&self.secret, token)?;
|
let claims = verify_token(&self.secret, token)?;
|
||||||
|
if claims.token_type != TokenType::Refresh {
|
||||||
|
anyhow::bail!("token is not a refresh token");
|
||||||
|
}
|
||||||
Ok(claims.sub)
|
Ok(claims.sub)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -193,6 +205,9 @@ pub struct ApiState {
|
|||||||
/// HS256 JWT token generation and verification.
|
/// HS256 JWT token generation and verification.
|
||||||
pub token_service: JwtTokenService,
|
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.
|
/// LLM provider client for chat completions.
|
||||||
pub llm_client: zesdex_infrastructure::llm::provider::LlmClient,
|
pub llm_client: zesdex_infrastructure::llm::provider::LlmClient,
|
||||||
}
|
}
|
||||||
@@ -272,6 +287,7 @@ impl ApiState {
|
|||||||
zesdex_application::cms::MemoryServiceImpl::new(memory_repo, memory_dir);
|
zesdex_application::cms::MemoryServiceImpl::new(memory_repo, memory_dir);
|
||||||
|
|
||||||
let token_service = JwtTokenService::new(&jwt_secret);
|
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(
|
let llm_client = zesdex_infrastructure::llm::provider::LlmClient::new(
|
||||||
llm_api_key.into(),
|
llm_api_key.into(),
|
||||||
llm_model.into(),
|
llm_model.into(),
|
||||||
@@ -287,6 +303,7 @@ impl ApiState {
|
|||||||
memory_service,
|
memory_service,
|
||||||
password_service: Argon2PasswordService,
|
password_service: Argon2PasswordService,
|
||||||
token_service,
|
token_service,
|
||||||
|
auth_rate_limiter,
|
||||||
llm_client,
|
llm_client,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,15 +2,25 @@
|
|||||||
//!
|
//!
|
||||||
//! Enables web clients and other WS-capable consumers to connect
|
//! Enables web clients and other WS-capable consumers to connect
|
||||||
//! and participate in sessions. Built on Axum's WebSocket support.
|
//! 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::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||||
|
use axum::extract::Query;
|
||||||
use axum::response::IntoResponse;
|
use axum::response::IntoResponse;
|
||||||
use axum::routing::get;
|
use axum::routing::get;
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
use futures_util::stream::StreamExt;
|
use futures_util::stream::StreamExt;
|
||||||
use futures_util::SinkExt;
|
use futures_util::SinkExt;
|
||||||
|
use serde::Deserialize;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tracing::info;
|
use tracing::{info, warn};
|
||||||
|
|
||||||
/// Shared application state for the WS server.
|
/// Shared application state for the WS server.
|
||||||
pub struct WsState {
|
pub struct WsState {
|
||||||
@@ -18,6 +28,12 @@ pub struct WsState {
|
|||||||
pub session_id: Option<String>,
|
pub session_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Query parameters accepted on the `/ws` upgrade.
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct WsQuery {
|
||||||
|
token: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
/// Build the WebSocket router.
|
/// Build the WebSocket router.
|
||||||
pub fn build_router(state: Arc<WsState>) -> Router {
|
pub fn build_router(state: Arc<WsState>) -> Router {
|
||||||
Router::new()
|
Router::new()
|
||||||
@@ -28,9 +44,18 @@ pub fn build_router(state: Arc<WsState>) -> Router {
|
|||||||
/// WebSocket upgrade handler.
|
/// WebSocket upgrade handler.
|
||||||
async fn ws_handler(
|
async fn ws_handler(
|
||||||
ws: WebSocketUpgrade,
|
ws: WebSocketUpgrade,
|
||||||
|
Query(query): Query<WsQuery>,
|
||||||
axum::extract::State(state): axum::extract::State<Arc<WsState>>,
|
axum::extract::State(state): axum::extract::State<Arc<WsState>>,
|
||||||
) -> impl IntoResponse {
|
) -> 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.
|
/// Handle an established WebSocket connection.
|
||||||
|
|||||||
Reference in New Issue
Block a user