diff --git a/server/src/api/auth_logout.rs b/server/src/api/auth_logout.rs index 6a9ede3..822c3a9 100644 --- a/server/src/api/auth_logout.rs +++ b/server/src/api/auth_logout.rs @@ -60,10 +60,27 @@ pub async fn logout( // Extract token from headers let token = extract_token_from_headers(request.headers())?; - // Add token to blacklist + // Add the login JWT to the blacklist. state.token_blacklist.revoke(&token).await; - info!("User {} logged out (token revoked)", user.username); + // Also revoke any outstanding session-scoped VIEWER tokens this user minted + // (CRITICAL #2). The login-JWT blacklist alone leaves a viewer token minted + // before logout valid for the rest of its 5-minute TTL, keeping a live + // viewer/remote-control plane open after logout. `user.user_id` is the `sub` + // the viewer tokens were registered under (the same claim stamped into them). + // The viewer WS already blacklist-checks the exact token string, so adding + // them here is sufficient — no WS change needed. take_for_user drains and + // clears the registry entry; expired tokens are pruned (not returned). + let viewer_tokens = state.viewer_tokens.take_for_user(&user.user_id); + let revoked_viewer_count = viewer_tokens.len(); + for viewer_token in viewer_tokens { + state.token_blacklist.revoke(&viewer_token).await; + } + + info!( + "User {} logged out (login token revoked, {} viewer token(s) revoked)", + user.username, revoked_viewer_count + ); Ok(Json(LogoutResponse { message: "Logged out successfully".to_string(), @@ -196,3 +213,74 @@ pub async fn cleanup_blacklist( remaining_count: remaining, })) } + +#[cfg(test)] +mod tests { + use crate::auth::{JwtConfig, TokenBlacklist, ViewerAccess, ViewerTokenRegistry}; + use std::time::Duration; + use uuid::Uuid; + + /// End-to-end (component-level) proof of CRITICAL #2: a viewer token minted + /// under a user, then revoked at logout via the registry drain, is rejected + /// by the SAME blacklist check the viewer WS runs (`is_revoked`). + /// + /// This exercises the real mint → register → logout-drain → blacklist path + /// without the HTTP/DB plumbing of the full handler. Uses local component + /// instances only (no process-global env), so it is parallel-safe. + #[tokio::test] + async fn logout_revokes_minted_viewer_token() { + let jwt = JwtConfig::new("test-secret-at-least-32-chars-long!!".to_string(), 24); + let registry = ViewerTokenRegistry::new(); + let blacklist = TokenBlacklist::new(); + + let user_sub = Uuid::new_v4().to_string(); + let session_id = Uuid::new_v4(); + let tenant_id = Uuid::new_v4(); + + // Mint a viewer token and register it under the user (as mint_viewer_token does). + let viewer_token = jwt + .create_viewer_token(&user_sub, session_id, tenant_id, ViewerAccess::Control) + .unwrap(); + registry.register(&user_sub, &viewer_token, Duration::from_secs(300)); + + // The viewer WS check (is_revoked) passes BEFORE logout — token is live. + assert!(!blacklist.is_revoked(&viewer_token).await); + + // Logout drains the user's viewer tokens into the blacklist (handler logic). + for tok in registry.take_for_user(&user_sub) { + blacklist.revoke(&tok).await; + } + + // After logout the same WS check now REJECTS the viewer token. + assert!(blacklist.is_revoked(&viewer_token).await); + // The token also remains a structurally valid viewer JWT (not expired), + // proving revocation — not natural expiry — is what blocks it. + assert!(jwt.validate_viewer_token(&viewer_token).is_ok()); + } + + /// A different user's logout must NOT revoke this user's viewer token. + #[tokio::test] + async fn logout_does_not_revoke_other_users_viewer_token() { + let jwt = JwtConfig::new("test-secret-at-least-32-chars-long!!".to_string(), 24); + let registry = ViewerTokenRegistry::new(); + let blacklist = TokenBlacklist::new(); + + let user_a = Uuid::new_v4().to_string(); + let user_b = Uuid::new_v4().to_string(); + let session_id = Uuid::new_v4(); + let tenant_id = Uuid::new_v4(); + + let token_b = jwt + .create_viewer_token(&user_b, session_id, tenant_id, ViewerAccess::ViewOnly) + .unwrap(); + registry.register(&user_b, &token_b, Duration::from_secs(300)); + + // user_a logs out — drains only user_a's (empty) token set. + for tok in registry.take_for_user(&user_a) { + blacklist.revoke(&tok).await; + } + + // user_b's token is untouched. + assert!(!blacklist.is_revoked(&token_b).await); + } +} diff --git a/server/src/api/sessions.rs b/server/src/api/sessions.rs index af58f10..6536482 100644 --- a/server/src/api/sessions.rs +++ b/server/src/api/sessions.rs @@ -160,6 +160,18 @@ pub async fn mint_viewer_token( ) })?; + // Register the minted token under the authenticated user's `sub` so logout + // can revoke it (CRITICAL #2): a viewer token minted before logout would + // otherwise stay valid for the rest of its 5-minute TTL. The registry prunes + // expired entries on insert, so it cannot grow unbounded. The token string + // is held transiently in memory (already handled server-side) and never + // logged. `user.user_id` is the `sub` claim stamped into the viewer token. + state.viewer_tokens.register( + &user.user_id, + &token, + std::time::Duration::from_secs(crate::auth::jwt::VIEWER_TOKEN_TTL_SECS as u64), + ); + tracing::info!( "User {} minted a {} viewer token for session {} (agent {})", user.username, diff --git a/server/src/auth/mod.rs b/server/src/auth/mod.rs index da9ab96..8cebc3d 100644 --- a/server/src/auth/mod.rs +++ b/server/src/auth/mod.rs @@ -7,10 +7,12 @@ pub mod agent_keys; pub mod jwt; pub mod password; pub mod token_blacklist; +pub mod viewer_token_registry; pub use jwt::{Claims, JwtConfig, ViewerAccess}; pub use password::{generate_random_password, hash_password, verify_password}; pub use token_blacklist::TokenBlacklist; +pub use viewer_token_registry::ViewerTokenRegistry; use axum::{ extract::FromRequestParts, diff --git a/server/src/auth/viewer_token_registry.rs b/server/src/auth/viewer_token_registry.rs new file mode 100644 index 0000000..8c05c83 --- /dev/null +++ b/server/src/auth/viewer_token_registry.rs @@ -0,0 +1,186 @@ +//! Per-user registry of outstanding viewer tokens. +//! +//! Minted viewer tokens ([`super::jwt::ViewerClaims`]) are short-lived +//! (`VIEWER_TOKEN_TTL_SECS`, 5 min) and session-scoped. The login-JWT blacklist +//! revokes the LOGIN token on logout, but a viewer token minted before logout +//! stays valid for the rest of its TTL — letting a just-logged-out user keep a +//! live remote-control plane until natural expiry. +//! +//! This registry closes that gap with the least-invasive mechanism: when a +//! viewer token is minted it is registered here under the minting user's `sub`; +//! on logout, the user's outstanding viewer tokens are taken from the registry +//! and added to the SAME exact-string blacklist the viewer WebSocket already +//! checks (`relay::viewer_ws_handler` -> `TokenBlacklist::is_revoked`). No +//! change to the WS path is required. +//! +//! Design notes: +//! - The registry holds bearer-token strings transiently in memory. The server +//! already handles these tokens; they are never logged or otherwise exposed +//! by this module. +//! - Entries are pruned-on-access (insert and take) by expiry, so the map can +//! never grow unbounded even if a user never logs out — expired viewer tokens +//! self-evict on the next operation that touches the map. +//! - Locking is a poison-tolerant `std::sync::Mutex` (matching the rate-limiter +//! pattern: `lock().unwrap_or_else(|e| e.into_inner())`); a panic while the +//! lock is held must not wedge logout/mint. No `.unwrap()` on lock paths. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +/// A single registered viewer token: the token string and the instant at which +/// it expires (mint time + TTL). +#[derive(Clone)] +struct Entry { + token: String, + expires_at: Instant, +} + +/// Per-user registry of outstanding (un-expired, un-revoked) viewer tokens. +/// +/// Cheap to clone (shares the inner map via `Arc`), so it can live in +/// `AppState` and be cloned with it like the other shared state. +#[derive(Clone, Default)] +pub struct ViewerTokenRegistry { + /// Map of user `sub` -> that user's outstanding viewer tokens. + inner: Arc>>>, +} + +impl ViewerTokenRegistry { + /// Create an empty registry. + pub fn new() -> Self { + Self::default() + } + + /// Register a freshly minted viewer token under its owning user (`sub`), + /// expiring `ttl` from now. + /// + /// Prunes expired entries for this user on insert so the per-user vector and + /// the overall map stay bounded. The token string is stored verbatim (it is + /// the exact string the viewer WS blacklist-checks); it is never logged. + pub fn register(&self, sub: &str, token: &str, ttl: Duration) { + let now = Instant::now(); + let expires_at = now + ttl; + let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + + let entries = map.entry(sub.to_string()).or_default(); + // Prune-on-insert: drop this user's already-expired tokens. + entries.retain(|e| e.expires_at > now); + entries.push(Entry { + token: token.to_string(), + expires_at, + }); + } + + /// Remove and return all currently-valid viewer tokens for `sub`. + /// + /// Expired tokens are pruned (not returned) — they are already invalid at + /// the JWT layer, so there is no value in blacklisting them. The user's + /// entry is cleared from the map regardless (taking everything out). Returns + /// an empty vec if the user has no registered tokens. + pub fn take_for_user(&self, sub: &str) -> Vec { + let now = Instant::now(); + let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + + match map.remove(sub) { + Some(entries) => entries + .into_iter() + .filter(|e| e.expires_at > now) + .map(|e| e.token) + .collect(), + None => Vec::new(), + } + } + + /// Number of users currently tracked (test/observability helper). Does not + /// prune; reflects the raw map size. + #[cfg(test)] + fn user_count(&self) -> usize { + let map = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + map.len() + } + + /// Number of registered (un-pruned) tokens for a user (test helper). + #[cfg(test)] + fn token_count(&self, sub: &str) -> usize { + let map = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + map.get(sub).map(|v| v.len()).unwrap_or(0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn register_then_take_returns_and_clears() { + let reg = ViewerTokenRegistry::new(); + reg.register("user-a", "tok-1", Duration::from_secs(300)); + reg.register("user-a", "tok-2", Duration::from_secs(300)); + + assert_eq!(reg.token_count("user-a"), 2); + + let mut taken = reg.take_for_user("user-a"); + taken.sort(); + assert_eq!(taken, vec!["tok-1".to_string(), "tok-2".to_string()]); + + // Taking clears the user's entry. + assert_eq!(reg.token_count("user-a"), 0); + assert_eq!(reg.user_count(), 0); + assert!(reg.take_for_user("user-a").is_empty()); + } + + #[test] + fn take_for_unknown_user_is_empty() { + let reg = ViewerTokenRegistry::new(); + assert!(reg.take_for_user("nobody").is_empty()); + } + + #[test] + fn tokens_are_scoped_per_user() { + let reg = ViewerTokenRegistry::new(); + reg.register("user-a", "a-tok", Duration::from_secs(300)); + reg.register("user-b", "b-tok", Duration::from_secs(300)); + + assert_eq!(reg.take_for_user("user-a"), vec!["a-tok".to_string()]); + // user-b is untouched. + assert_eq!(reg.take_for_user("user-b"), vec!["b-tok".to_string()]); + } + + #[test] + fn expired_tokens_are_not_returned_by_take() { + let reg = ViewerTokenRegistry::new(); + // Zero TTL -> already expired by the time take runs. + reg.register("user-a", "expired", Duration::from_secs(0)); + reg.register("user-a", "live", Duration::from_secs(300)); + + let taken = reg.take_for_user("user-a"); + assert_eq!(taken, vec!["live".to_string()]); + } + + #[test] + fn expired_tokens_are_pruned_on_insert() { + let reg = ViewerTokenRegistry::new(); + // Register an already-expired token, then a live one for the same user. + reg.register("user-a", "expired", Duration::from_secs(0)); + // First insert leaves 1 entry; the second insert prunes the expired one + // before pushing, so only the live token remains. + reg.register("user-a", "live", Duration::from_secs(300)); + + assert_eq!(reg.token_count("user-a"), 1); + assert_eq!(reg.take_for_user("user-a"), vec!["live".to_string()]); + } + + #[test] + fn multiple_logins_accumulate_then_all_taken() { + // A user with several concurrent logins mints several viewer tokens; + // logout must revoke ALL of them at once. + let reg = ViewerTokenRegistry::new(); + for i in 0..5 { + reg.register("user-a", &format!("tok-{i}"), Duration::from_secs(300)); + } + assert_eq!(reg.token_count("user-a"), 5); + assert_eq!(reg.take_for_user("user-a").len(), 5); + assert_eq!(reg.user_count(), 0); + } +} diff --git a/server/src/main.rs b/server/src/main.rs index 94d0981..374181f 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -37,7 +37,10 @@ use tower_http::trace::TraceLayer; use tracing::{info, Level}; use tracing_subscriber::FmtSubscriber; -use auth::{generate_random_password, hash_password, AuthenticatedUser, JwtConfig, TokenBlacklist}; +use auth::{ + generate_random_password, hash_password, AuthenticatedUser, JwtConfig, TokenBlacklist, + ViewerTokenRegistry, +}; /// Root of the static asset tree, relative to the server's working directory. /// Holds the agent `downloads/` tree AND the v2 SPA build under `app/`. @@ -66,6 +69,13 @@ pub struct AppState { db: Option, pub jwt_config: Arc, pub token_blacklist: TokenBlacklist, + /// Per-user registry of outstanding session-scoped viewer tokens. Minting a + /// viewer token registers it here under the minting user's `sub`; logout + /// drains the user's registered viewer tokens into `token_blacklist` so a + /// just-logged-out user cannot keep a live viewer/remote-control plane until + /// the token's natural 5-minute expiry. The viewer WS already blacklist- + /// checks the exact token string, so no WS change is needed. + pub viewer_tokens: ViewerTokenRegistry, /// Optional API key for persistent agents (env: AGENT_API_KEY) pub agent_api_key: Option, /// Prometheus metrics @@ -291,6 +301,7 @@ async fn main() -> Result<()> { // Create application state let token_blacklist = TokenBlacklist::new(); + let viewer_tokens = ViewerTokenRegistry::new(); let state = AppState { sessions, @@ -298,6 +309,7 @@ async fn main() -> Result<()> { db: database, jwt_config, token_blacklist, + viewer_tokens, agent_api_key, metrics, registry, diff --git a/server/src/relay/mod.rs b/server/src/relay/mod.rs index 68e8a13..4cc69c1 100644 --- a/server/src/relay/mod.rs +++ b/server/src/relay/mod.rs @@ -772,8 +772,11 @@ async fn handle_agent_connection( let _ = frame_tx.send(data.to_vec()); } Some(proto::message::Payload::ChatMessage(chat)) => { - // Broadcast chat message to all viewers - info!("Chat from client: {}", chat.content); + // Broadcast chat message to all viewers. Do NOT log + // the chat body: support-session chat can carry + // secrets/PII. Log only the (non-sensitive) length + // so relay activity stays observable. + info!("Chat from client ({} chars)", chat.content.len()); let _ = frame_tx.send(data.to_vec()); } Some(proto::message::Payload::AgentStatus(status)) => { @@ -1369,8 +1372,10 @@ async fn handle_viewer_connection( } Some(proto::message::Payload::ChatMessage(chat)) => { // Forward chat message to agent (not throttled — - // not an injected-input vector). Bounded send. - info!("Chat from technician: {}", chat.content); + // not an injected-input vector). Bounded send. Do + // NOT log the chat body (secrets/PII); log only the + // (non-sensitive) length so the relay is observable. + info!("Chat from technician ({} chars)", chat.content.len()); let _ = input_tx.try_send(data.to_vec()); } _ => {}