Explore the GuruConnect server codebase at C:/Users/MikeSwanson/Claude/guru-connect/server/ to understand:

1. Current session management in src/session/mod.rs - how sessions are stored and managed
2. The overall server architecture in src/main.rs - how state is shared
3. Any existing database or persistence patterns
4. The Cargo.toml to see existing dependencies

Provide detailed analysis of the current architecture and identify where database integration would need to happen.


I'll explore the GuruConnect server codebase to understand the current architecture and session management. Let me start by examining the structure and key files.


Bash

find "C:/Users/MikeSwanson/Claude/guru-connect/server" -type f -name "*.rs" -o -name "*.toml" | head -20

List Rust source files and config files in server


     1→[package]
     2→name = "guruconnect-server"
     3→version = "0.1.0"
     4→edition = "2021"
     5→authors = ["AZ Computer Guru"]
     6→description = "GuruConnect Remote Desktop Relay Server"
     7→
     8→[dependencies]
     9→# Async runtime
    10→tokio = { version = "1", features = ["full", "sync", "time", "rt-multi-thread", "macros"] }
    11→
    12→# Web framework
    13→axum = { version = "0.7", features = ["ws", "macros"] }
    14→tower = "0.5"
    15→tower-http = { version = "0.6", features = ["cors", "trace", "compression-gzip", "fs"] }
    16→
    17→# WebSocket
    18→futures-util = "0.3"
    19→
    20→# Database
    21→sqlx = { version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "chrono", "json"] }
    22→
    23→# Protocol (protobuf)
    24→prost = "0.13"
    25→prost-types = "0.13"
    26→bytes = "1"
    27→
    28→# Serialization
    29→serde = { version = "1", features = ["derive"] }
    30→serde_json = "1"
    31→
    32→# Logging
    33→tracing = "0.1"
    34→tracing-subscriber = { version = "0.3", features = ["env-filter"] }
    35→
    36→# Error handling
    37→anyhow = "1"
    38→thiserror = "1"
    39→
    40→# Configuration
    41→toml = "0.8"
    42→
    43→# Auth
    44→jsonwebtoken = "9"
    45→argon2 = "0.5"
    46→
    47→# Crypto
    48→ring = "0.17"
    49→
    50→# UUID
    51→uuid = { version = "1", features = ["v4", "serde"] }
    52→
    53→# Time
    54→chrono = { version = "0.4", features = ["serde"] }
    55→rand = "0.8"
    56→
    57→[build-dependencies]
    58→prost-build = "0.13"
    59→
    60→[profile.release]
    61→lto = true
    62→codegen-units = 1
    63→strip = true
    64→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

[package]
name = "guruconnect-server"
version = "0.1.0"
edition = "2021"
authors = ["AZ Computer Guru"]
description = "GuruConnect Remote Desktop Relay Server"

[dependencies]
# Async runtime
tokio = { version = "1", features = ["full", "sync", "time", "rt-multi-thread", "macros"] }

# Web framework
axum = { version = "0.7", features = ["ws", "macros"] }
tower = "0.5"
tower-http = { version = "0.6", features = ["cors", "trace", "compression-gzip", "fs"] }

# WebSocket
futures-util = "0.3"

# Database
sqlx = { version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "chrono", "json"] }

# Protocol (protobuf)
prost = "0.13"
prost-types = "0.13"
bytes = "1"

# Serialization
serde = { version = "1", features = ["derive"] }
serde_json = "1"

# Logging
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

# Error handling
anyhow = "1"
thiserror = "1"

# Configuration
toml = "0.8"

# Auth
jsonwebtoken = "9"
argon2 = "0.5"

# Crypto
ring = "0.17"

# UUID
uuid = { version = "1", features = ["v4", "serde"] }

# Time
chrono = { version = "0.4", features = ["serde"] }
rand = "0.8"

[build-dependencies]
prost-build = "0.13"

[profile.release]
lto = true
codegen-units = 1
strip = true


Now let me read the key files to understand the architecture:


     1→//! GuruConnect Server - WebSocket Relay Server
     2→//!
     3→//! Handles connections from both agents and dashboard viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→mod config;
     7→mod relay;
     8→mod session;
     9→mod auth;
    10→mod api;
    11→mod db;
    12→mod support_codes;
    13→
    14→pub mod proto {
    15→    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
    16→}
    17→
    18→use anyhow::Result;
    19→use axum::{
    20→    Router,
    21→    routing::{get, post},
    22→    extract::{Path, State, Json},
    23→    response::{Html, IntoResponse},
    24→    http::StatusCode,
    25→};
    26→use std::net::SocketAddr;
    27→use tower_http::cors::{Any, CorsLayer};
    28→use tower_http::trace::TraceLayer;
    29→use tower_http::services::ServeDir;
    30→use tracing::{info, Level};
    31→use tracing_subscriber::FmtSubscriber;
    32→use serde::Deserialize;
    33→
    34→use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
    35→
    36→/// Application state
    37→#[derive(Clone)]
    38→pub struct AppState {
    39→    sessions: session::SessionManager,
    40→    support_codes: SupportCodeManager,
    41→}
    42→
    43→#[tokio::main]
    44→async fn main() -> Result<()> {
    45→    // Initialize logging
    46→    let _subscriber = FmtSubscriber::builder()
    47→        .with_max_level(Level::INFO)
    48→        .with_target(true)
    49→        .init();
    50→
    51→    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));
    52→
    53→    // Load configuration
    54→    let config = config::Config::load()?;
    55→
    56→    // Use port 3002 for GuruConnect
    57→    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    58→    info!("Loaded configuration, listening on {}", listen_addr);
    59→
    60→    // Create application state
    61→    let state = AppState {
    62→        sessions: session::SessionManager::new(),
    63→        support_codes: SupportCodeManager::new(),
    64→    };
    65→
    66→    // Build router
    67→    let app = Router::new()
    68→        // Health check
    69→        .route("/health", get(health))
    70→
    71→        // Portal API - Support codes
    72→        .route("/api/codes", post(create_code))
    73→        .route("/api/codes", get(list_codes))
    74→        .route("/api/codes/:code/validate", get(validate_code))
    75→        .route("/api/codes/:code/cancel", post(cancel_code))
    76→
    77→        // WebSocket endpoints
    78→        .route("/ws/agent", get(relay::agent_ws_handler))
    79→        .route("/ws/viewer", get(relay::viewer_ws_handler))
    80→
    81→        // REST API - Sessions
    82→        .route("/api/sessions", get(list_sessions))
    83→        .route("/api/sessions/:id", get(get_session))
    84→        .route("/api/sessions/:id", axum::routing::delete(disconnect_session))
    85→
    86→        // HTML page routes (clean URLs)
    87→        .route("/login", get(serve_login))
    88→        .route("/dashboard", get(serve_dashboard))
    89→
    90→        // State
    91→        .with_state(state)
    92→
    93→        // Serve static files for portal (fallback)
    94→        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))
    95→
    96→        // Middleware
    97→        .layer(TraceLayer::new_for_http())
    98→        .layer(
    99→            CorsLayer::new()
   100→                .allow_origin(Any)
   101→                .allow_methods(Any)
   102→                .allow_headers(Any),
   103→        );
   104→
   105→    // Start server
   106→    let addr: SocketAddr = listen_addr.parse()?;
   107→    let listener = tokio::net::TcpListener::bind(addr).await?;
   108→
   109→    info!("Server listening on {}", addr);
   110→
   111→    axum::serve(listener, app).await?;
   112→
   113→    Ok(())
   114→}
   115→
   116→async fn health() -> &'static str {
   117→    "OK"
   118→}
   119→
   120→// Support code API handlers
   121→
   122→async fn create_code(
   123→    State(state): State<AppState>,
   124→    Json(request): Json<CreateCodeRequest>,
   125→) -> Json<SupportCode> {
   126→    let code = state.support_codes.create_code(request).await;
   127→    info!("Created support code: {}", code.code);
   128→    Json(code)
   129→}
   130→
   131→async fn list_codes(
   132→    State(state): State<AppState>,
   133→) -> Json<Vec<SupportCode>> {
   134→    Json(state.support_codes.list_active_codes().await)
   135→}
   136→
   137→#[derive(Deserialize)]
   138→struct ValidateParams {
   139→    code: String,
   140→}
   141→
   142→async fn validate_code(
   143→    State(state): State<AppState>,
   144→    Path(code): Path<String>,
   145→) -> Json<CodeValidation> {
   146→    Json(state.support_codes.validate_code(&code).await)
   147→}
   148→
   149→async fn cancel_code(
   150→    State(state): State<AppState>,
   151→    Path(code): Path<String>,
   152→) -> impl IntoResponse {
   153→    if state.support_codes.cancel_code(&code).await {
   154→        (StatusCode::OK, "Code cancelled")
   155→    } else {
   156→        (StatusCode::BAD_REQUEST, "Cannot cancel code")
   157→    }
   158→}
   159→
   160→// Session API handlers (updated to use AppState)
   161→
   162→async fn list_sessions(
   163→    State(state): State<AppState>,
   164→) -> Json<Vec<api::SessionInfo>> {
   165→    let sessions = state.sessions.list_sessions().await;
   166→    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
   167→}
   168→
   169→async fn get_session(
   170→    State(state): State<AppState>,
   171→    Path(id): Path<String>,
   172→) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
   173→    let session_id = uuid::Uuid::parse_str(&id)
   174→        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;
   175→
   176→    let session = state.sessions.get_session(session_id).await
   177→        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;
   178→
   179→    Ok(Json(api::SessionInfo::from(session)))
   180→}
   181→
   182→async fn disconnect_session(
   183→    State(state): State<AppState>,
   184→    Path(id): Path<String>,
   185→) -> impl IntoResponse {
   186→    let session_id = match uuid::Uuid::parse_str(&id) {
   187→        Ok(id) => id,
   188→        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
   189→    };
   190→
   191→    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
   192→        info!("Session {} disconnected by admin", session_id);
   193→        (StatusCode::OK, "Session disconnected")
   194→    } else {
   195→        (StatusCode::NOT_FOUND, "Session not found")
   196→    }
   197→}
   198→
   199→// Static page handlers
   200→async fn serve_login() -> impl IntoResponse {
   201→    match tokio::fs::read_to_string("static/login.html").await {
   202→        Ok(content) => Html(content).into_response(),
   203→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   204→    }
   205→}
   206→
   207→async fn serve_dashboard() -> impl IntoResponse {
   208→    match tokio::fs::read_to_string("static/dashboard.html").await {
   209→        Ok(content) => Html(content).into_response(),
   210→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   211→    }
   212→}
   213→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! GuruConnect Server - WebSocket Relay Server
//!
//! Handles connections from both agents and dashboard viewers,
//! relaying video frames and input events between them.

mod config;
mod relay;
mod session;
mod auth;
mod api;
mod db;
mod support_codes;

pub mod proto {
    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
}

use anyhow::Result;
use axum::{
    Router,
    routing::{get, post},
    extract::{Path, State, Json},
    response::{Html, IntoResponse},
    http::StatusCode,
};
use std::net::SocketAddr;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
use tower_http::services::ServeDir;
use tracing::{info, Level};
use tracing_subscriber::FmtSubscriber;
use serde::Deserialize;

use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};

/// Application state
#[derive(Clone)]
pub struct AppState {
    sessions: session::SessionManager,
    support_codes: SupportCodeManager,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    let _subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .with_target(true)
        .init();

    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));

    // Load configuration
    let config = config::Config::load()?;

    // Use port 3002 for GuruConnect
    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    info!("Loaded configuration, listening on {}", listen_addr);

    // Create application state
    let state = AppState {
        sessions: session::SessionManager::new(),
        support_codes: SupportCodeManager::new(),
    };

    // Build router
    let app = Router::new()
        // Health check
        .route("/health", get(health))

        // Portal API - Support codes
        .route("/api/codes", post(create_code))
        .route("/api/codes", get(list_codes))
        .route("/api/codes/:code/validate", get(validate_code))
        .route("/api/codes/:code/cancel", post(cancel_code))

        // WebSocket endpoints
        .route("/ws/agent", get(relay::agent_ws_handler))
        .route("/ws/viewer", get(relay::viewer_ws_handler))

        // REST API - Sessions
        .route("/api/sessions", get(list_sessions))
        .route("/api/sessions/:id", get(get_session))
        .route("/api/sessions/:id", axum::routing::delete(disconnect_session))

        // HTML page routes (clean URLs)
        .route("/login", get(serve_login))
        .route("/dashboard", get(serve_dashboard))

        // State
        .with_state(state)

        // Serve static files for portal (fallback)
        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))

        // Middleware
        .layer(TraceLayer::new_for_http())
        .layer(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                .allow_headers(Any),
        );

    // Start server
    let addr: SocketAddr = listen_addr.parse()?;
    let listener = tokio::net::TcpListener::bind(addr).await?;

    info!("Server listening on {}", addr);

    axum::serve(listener, app).await?;

    Ok(())
}

async fn health() -> &'static str {
    "OK"
}

// Support code API handlers

async fn create_code(
    State(state): State<AppState>,
    Json(request): Json<CreateCodeRequest>,
) -> Json<SupportCode> {
    let code = state.support_codes.create_code(request).await;
    info!("Created support code: {}", code.code);
    Json(code)
}

async fn list_codes(
    State(state): State<AppState>,
) -> Json<Vec<SupportCode>> {
    Json(state.support_codes.list_active_codes().await)
}

#[derive(Deserialize)]
struct ValidateParams {
    code: String,
}

async fn validate_code(
    State(state): State<AppState>,
    Path(code): Path<String>,
) -> Json<CodeValidation> {
    Json(state.support_codes.validate_code(&code).await)
}

async fn cancel_code(
    State(state): State<AppState>,
    Path(code): Path<String>,
) -> impl IntoResponse {
    if state.support_codes.cancel_code(&code).await {
        (StatusCode::OK, "Code cancelled")
    } else {
        (StatusCode::BAD_REQUEST, "Cannot cancel code")
    }
}

// Session API handlers (updated to use AppState)

async fn list_sessions(
    State(state): State<AppState>,
) -> Json<Vec<api::SessionInfo>> {
    let sessions = state.sessions.list_sessions().await;
    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
}

async fn get_session(
    State(state): State<AppState>,
    Path(id): Path<String>,
) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
    let session_id = uuid::Uuid::parse_str(&id)
        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;

    let session = state.sessions.get_session(session_id).await
        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;

    Ok(Json(api::SessionInfo::from(session)))
}

async fn disconnect_session(
    State(state): State<AppState>,
    Path(id): Path<String>,
) -> impl IntoResponse {
    let session_id = match uuid::Uuid::parse_str(&id) {
        Ok(id) => id,
        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
    };

    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
        info!("Session {} disconnected by admin", session_id);
        (StatusCode::OK, "Session disconnected")
    } else {
        (StatusCode::NOT_FOUND, "Session not found")
    }
}

// Static page handlers
async fn serve_login() -> impl IntoResponse {
    match tokio::fs::read_to_string("static/login.html").await {
        Ok(content) => Html(content).into_response(),
        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
    }
}

async fn serve_dashboard() -> impl IntoResponse {
    match tokio::fs::read_to_string("static/dashboard.html").await {
        Ok(content) => Html(content).into_response(),
        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
    }
}


     1→//! Session management for GuruConnect
     2→//!
     3→//! Manages active remote desktop sessions, tracking which agents
     4→//! are connected and which viewers are watching them.
     5→
     6→use std::collections::HashMap;
     7→use std::sync::Arc;
     8→use std::time::Instant;
     9→use tokio::sync::{broadcast, RwLock};
    10→use uuid::Uuid;
    11→
    12→/// Unique identifier for a session
    13→pub type SessionId = Uuid;
    14→
    15→/// Unique identifier for an agent
    16→pub type AgentId = String;
    17→
    18→/// Unique identifier for a viewer
    19→pub type ViewerId = String;
    20→
    21→/// Information about a connected viewer/technician
    22→#[derive(Debug, Clone)]
    23→pub struct ViewerInfo {
    24→    pub id: ViewerId,
    25→    pub name: String,
    26→    pub connected_at: chrono::DateTime<chrono::Utc>,
    27→}
    28→
    29→/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
    30→const HEARTBEAT_TIMEOUT_SECS: u64 = 90;
    31→
    32→/// Session state
    33→#[derive(Debug, Clone)]
    34→pub struct Session {
    35→    pub id: SessionId,
    36→    pub agent_id: AgentId,
    37→    pub agent_name: String,
    38→    pub started_at: chrono::DateTime<chrono::Utc>,
    39→    pub viewer_count: usize,
    40→    pub viewers: Vec<ViewerInfo>,  // List of connected technicians
    41→    pub is_streaming: bool,
    42→    pub is_online: bool,  // Whether agent is currently connected
    43→    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    44→    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    45→    // Agent status info
    46→    pub os_version: Option<String>,
    47→    pub is_elevated: bool,
    48→    pub uptime_secs: i64,
    49→    pub display_count: i32,
    50→}
    51→
    52→/// Channel for sending frames from agent to viewers
    53→pub type FrameSender = broadcast::Sender<Vec<u8>>;
    54→pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;
    55→
    56→/// Channel for sending input events from viewer to agent
    57→pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
    58→pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;
    59→
    60→/// Internal session data with channels
    61→struct SessionData {
    62→    info: Session,
    63→    /// Channel for video frames (agent -> viewers)
    64→    frame_tx: FrameSender,
    65→    /// Channel for input events (viewer -> agent)
    66→    input_tx: InputSender,
    67→    input_rx: Option<InputReceiver>,
    68→    /// Map of connected viewers (id -> info)
    69→    viewers: HashMap<ViewerId, ViewerInfo>,
    70→    /// Instant for heartbeat tracking
    71→    last_heartbeat_instant: Instant,
    72→}
    73→
    74→/// Manages all active sessions
    75→#[derive(Clone)]
    76→pub struct SessionManager {
    77→    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    78→    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
    79→}
    80→
    81→impl SessionManager {
    82→    pub fn new() -> Self {
    83→        Self {
    84→            sessions: Arc::new(RwLock::new(HashMap::new())),
    85→            agents: Arc::new(RwLock::new(HashMap::new())),
    86→        }
    87→    }
    88→
    89→    /// Register a new agent and create a session
    90→    /// If agent was previously connected (offline session exists), reuse that session
    91→    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
    92→        // Check if this agent already has an offline session (reconnecting)
    93→        {
    94→            let agents = self.agents.read().await;
    95→            if let Some(&existing_session_id) = agents.get(&agent_id) {
    96→                let mut sessions = self.sessions.write().await;
    97→                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
    98→                    if !session_data.info.is_online {
    99→                        // Reuse existing session - mark as online and create new channels
   100→                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);
   101→
   102→                        let (frame_tx, _) = broadcast::channel(16);
   103→                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);
   104→
   105→                        session_data.info.is_online = true;
   106→                        session_data.info.last_heartbeat = chrono::Utc::now();
   107→                        session_data.info.agent_name = agent_name; // Update name in case it changed
   108→                        session_data.frame_tx = frame_tx.clone();
   109→                        session_data.input_tx = input_tx;
   110→                        session_data.last_heartbeat_instant = Instant::now();
   111→
   112→                        return (existing_session_id, frame_tx, input_rx);
   113→                    }
   114→                }
   115→            }
   116→        }
   117→
   118→        // Create new session
   119→        let session_id = Uuid::new_v4();
   120→
   121→        // Create channels
   122→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
   123→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
   124→
   125→        let now = chrono::Utc::now();
   126→        let session = Session {
   127→            id: session_id,
   128→            agent_id: agent_id.clone(),
   129→            agent_name,
   130→            started_at: now,
   131→            viewer_count: 0,
   132→            viewers: Vec::new(),
   133→            is_streaming: false,
   134→            is_online: true,
   135→            is_persistent,
   136→            last_heartbeat: now,
   137→            os_version: None,
   138→            is_elevated: false,
   139→            uptime_secs: 0,
   140→            display_count: 1,
   141→        };
   142→
   143→        let session_data = SessionData {
   144→            info: session,
   145→            frame_tx: frame_tx.clone(),
   146→            input_tx,
   147→            input_rx: None,
   148→            viewers: HashMap::new(),
   149→            last_heartbeat_instant: Instant::now(),
   150→        };
   151→
   152→        let mut sessions = self.sessions.write().await;
   153→        sessions.insert(session_id, session_data);
   154→
   155→        let mut agents = self.agents.write().await;
   156→        agents.insert(agent_id, session_id);
   157→
   158→        (session_id, frame_tx, input_rx)
   159→    }
   160→
   161→    /// Update agent status from heartbeat or status message
   162→    pub async fn update_agent_status(
   163→        &self,
   164→        session_id: SessionId,
   165→        os_version: Option<String>,
   166→        is_elevated: bool,
   167→        uptime_secs: i64,
   168→        display_count: i32,
   169→        is_streaming: bool,
   170→    ) {
   171→        let mut sessions = self.sessions.write().await;
   172→        if let Some(session_data) = sessions.get_mut(&session_id) {
   173→            session_data.info.last_heartbeat = chrono::Utc::now();
   174→            session_data.last_heartbeat_instant = Instant::now();
   175→            session_data.info.is_streaming = is_streaming;
   176→            if let Some(os) = os_version {
   177→                session_data.info.os_version = Some(os);
   178→            }
   179→            session_data.info.is_elevated = is_elevated;
   180→            session_data.info.uptime_secs = uptime_secs;
   181→            session_data.info.display_count = display_count;
   182→        }
   183→    }
   184→
   185→    /// Update heartbeat timestamp
   186→    pub async fn update_heartbeat(&self, session_id: SessionId) {
   187→        let mut sessions = self.sessions.write().await;
   188→        if let Some(session_data) = sessions.get_mut(&session_id) {
   189→            session_data.info.last_heartbeat = chrono::Utc::now();
   190→            session_data.last_heartbeat_instant = Instant::now();
   191→        }
   192→    }
   193→
   194→    /// Check if a session has timed out (no heartbeat for too long)
   195→    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
   196→        let sessions = self.sessions.read().await;
   197→        if let Some(session_data) = sessions.get(&session_id) {
   198→            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
   199→        } else {
   200→            true // Non-existent sessions are considered timed out
   201→        }
   202→    }
   203→
   204→    /// Get sessions that have timed out
   205→    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
   206→        let sessions = self.sessions.read().await;
   207→        sessions
   208→            .iter()
   209→            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
   210→            .map(|(id, _)| *id)
   211→            .collect()
   212→    }
   213→
   214→    /// Get a session by agent ID
   215→    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
   216→        let agents = self.agents.read().await;
   217→        let session_id = agents.get(agent_id)?;
   218→
   219→        let sessions = self.sessions.read().await;
   220→        sessions.get(session_id).map(|s| s.info.clone())
   221→    }
   222→
   223→    /// Get a session by session ID
   224→    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
   225→        let sessions = self.sessions.read().await;
   226→        sessions.get(&session_id).map(|s| s.info.clone())
   227→    }
   228→
   229→    /// Join a session as a viewer, returns channels and sends StartStream to agent
   230→    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId, viewer_name: String) -> Option<(FrameReceiver, InputSender)> {
   231→        let mut sessions = self.sessions.write().await;
   232→        let session_data = sessions.get_mut(&session_id)?;
   233→
   234→        let was_empty = session_data.viewers.is_empty();
   235→
   236→        // Add viewer info
   237→        let viewer_info = ViewerInfo {
   238→            id: viewer_id.clone(),
   239→            name: viewer_name.clone(),
   240→            connected_at: chrono::Utc::now(),
   241→        };
   242→        session_data.viewers.insert(viewer_id.clone(), viewer_info);
   243→
   244→        // Update session info
   245→        session_data.info.viewer_count = session_data.viewers.len();
   246→        session_data.info.viewers = session_data.viewers.values().cloned().collect();
   247→
   248→        let frame_rx = session_data.frame_tx.subscribe();
   249→        let input_tx = session_data.input_tx.clone();
   250→
   251→        // If this is the first viewer, send StartStream to agent
   252→        if was_empty {
   253→            tracing::info!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
   254→            Self::send_start_stream_internal(session_data, &viewer_id).await;
   255→        } else {
   256→            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
   257→        }
   258→
   259→        Some((frame_rx, input_tx))
   260→    }
   261→
   262→    /// Internal helper to send StartStream message
   263→    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
   264→        use crate::proto;
   265→        use prost::Message;
   266→
   267→        let start_stream = proto::Message {
   268→            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
   269→                viewer_id: viewer_id.to_string(),
   270→                display_id: 0, // Primary display
   271→            })),
   272→        };
   273→
   274→        let mut buf = Vec::new();
   275→        if start_stream.encode(&mut buf).is_ok() {
   276→            let _ = session_data.input_tx.send(buf).await;
   277→        }
   278→    }
   279→
   280→    /// Leave a session as a viewer, sends StopStream if no viewers left
   281→    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
   282→        let mut sessions = self.sessions.write().await;
   283→        if let Some(session_data) = sessions.get_mut(&session_id) {
   284→            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
   285→            session_data.viewers.remove(viewer_id);
   286→            session_data.info.viewer_count = session_data.viewers.len();
   287→            session_data.info.viewers = session_data.viewers.values().cloned().collect();
   288→
   289→            // If no more viewers, send StopStream to agent
   290→            if session_data.viewers.is_empty() {
   291→                tracing::info!("Last viewer {} ({}) left session {}, sending StopStream",
   292→                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
   293→                Self::send_stop_stream_internal(session_data, viewer_id).await;
   294→            } else {
   295→                tracing::info!("Viewer {} ({}) left session {}",
   296→                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
   297→            }
   298→        }
   299→    }
   300→
   301→    /// Internal helper to send StopStream message
   302→    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
   303→        use crate::proto;
   304→        use prost::Message;
   305→
   306→        let stop_stream = proto::Message {
   307→            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
   308→                viewer_id: viewer_id.to_string(),
   309→            })),
   310→        };
   311→
   312→        let mut buf = Vec::new();
   313→        if stop_stream.encode(&mut buf).is_ok() {
   314→            let _ = session_data.input_tx.send(buf).await;
   315→        }
   316→    }
   317→
   318→    /// Mark agent as disconnected
   319→    /// For persistent agents: keep session but mark as offline
   320→    /// For support sessions: remove session entirely
   321→    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
   322→        let mut sessions = self.sessions.write().await;
   323→        if let Some(session_data) = sessions.get_mut(&session_id) {
   324→            if session_data.info.is_persistent {
   325→                // Persistent agent - keep session but mark as offline
   326→                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
   327→                    session_data.info.agent_id, session_id);
   328→                session_data.info.is_online = false;
   329→                session_data.info.is_streaming = false;
   330→                session_data.info.viewer_count = 0;
   331→                session_data.info.viewers.clear();
   332→                session_data.viewers.clear();
   333→            } else {
   334→                // Support session - remove entirely
   335→                let agent_id = session_data.info.agent_id.clone();
   336→                sessions.remove(&session_id);
   337→                drop(sessions); // Release sessions lock before acquiring agents lock
   338→                let mut agents = self.agents.write().await;
   339→                agents.remove(&agent_id);
   340→                tracing::info!("Support session {} removed", session_id);
   341→            }
   342→        }
   343→    }
   344→
   345→    /// Remove a session entirely (for cleanup)
   346→    pub async fn remove_session(&self, session_id: SessionId) {
   347→        let mut sessions = self.sessions.write().await;
   348→        if let Some(session_data) = sessions.remove(&session_id) {
   349→            drop(sessions);
   350→            let mut agents = self.agents.write().await;
   351→            agents.remove(&session_data.info.agent_id);
   352→        }
   353→    }
   354→
   355→    /// Disconnect a session by sending a disconnect message to the agent
   356→    /// Returns true if the message was sent successfully
   357→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   358→        let sessions = self.sessions.read().await;
   359→        if let Some(session_data) = sessions.get(&session_id) {
   360→            // Create disconnect message
   361→            use crate::proto;
   362→            use prost::Message;
   363→
   364→            let disconnect_msg = proto::Message {
   365→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   366→                    reason: reason.to_string(),
   367→                })),
   368→            };
   369→
   370→            let mut buf = Vec::new();
   371→            if disconnect_msg.encode(&mut buf).is_ok() {
   372→                // Send via input channel (will be forwarded to agent's WebSocket)
   373→                if session_data.input_tx.send(buf).await.is_ok() {
   374→                    return true;
   375→                }
   376→            }
   377→        }
   378→        false
   379→    }
   380→
   381→    /// List all active sessions
   382→    pub async fn list_sessions(&self) -> Vec<Session> {
   383→        let sessions = self.sessions.read().await;
   384→        sessions.values().map(|s| s.info.clone()).collect()
   385→    }
   386→}
   387→
   388→impl Default for SessionManager {
   389→    fn default() -> Self {
   390→        Self::new()
   391→    }
   392→}
   393→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! Session management for GuruConnect
//!
//! Manages active remote desktop sessions, tracking which agents
//! are connected and which viewers are watching them.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;

/// Unique identifier for a session
pub type SessionId = Uuid;

/// Unique identifier for an agent
pub type AgentId = String;

/// Unique identifier for a viewer
pub type ViewerId = String;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
const HEARTBEAT_TIMEOUT_SECS: u64 = 90;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
    pub viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    // Agent status info
    pub os_version: Option<String>,
    pub is_elevated: bool,
    pub uptime_secs: i64,
    pub display_count: i32,
}

/// Channel for sending frames from agent to viewers
pub type FrameSender = broadcast::Sender<Vec<u8>>;
pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

/// Channel for sending input events from viewer to agent
pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

/// Internal session data with channels
struct SessionData {
    info: Session,
    /// Channel for video frames (agent -> viewers)
    frame_tx: FrameSender,
    /// Channel for input events (viewer -> agent)
    input_tx: InputSender,
    input_rx: Option<InputReceiver>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// Instant for heartbeat tracking
    last_heartbeat_instant: Instant,
}

/// Manages all active sessions
#[derive(Clone)]
pub struct SessionManager {
    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
}

impl SessionManager {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            agents: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new agent and create a session
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        let session_id = Uuid::new_v4();

        // Create channels
        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

        let now = chrono::Utc::now();
        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            last_heartbeat: now,
            os_version: None,
            is_elevated: false,
            uptime_secs: 0,
            display_count: 1,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None,
            viewers: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };

        let mut sessions = self.sessions.write().await;
        sessions.insert(session_id, session_data);

        let mut agents = self.agents.write().await;
        agents.insert(agent_id, session_id);

        (session_id, frame_tx, input_rx)
    }

    /// Update agent status from heartbeat or status message
    pub async fn update_agent_status(
        &self,
        session_id: SessionId,
        os_version: Option<String>,
        is_elevated: bool,
        uptime_secs: i64,
        display_count: i32,
        is_streaming: bool,
    ) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
            session_data.info.is_streaming = is_streaming;
            if let Some(os) = os_version {
                session_data.info.os_version = Some(os);
            }
            session_data.info.is_elevated = is_elevated;
            session_data.info.uptime_secs = uptime_secs;
            session_data.info.display_count = display_count;
        }
    }

    /// Update heartbeat timestamp
    pub async fn update_heartbeat(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
        }
    }

    /// Check if a session has timed out (no heartbeat for too long)
    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
        } else {
            true // Non-existent sessions are considered timed out
        }
    }

    /// Get sessions that have timed out
    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
        let sessions = self.sessions.read().await;
        sessions
            .iter()
            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
            .map(|(id, _)| *id)
            .collect()
    }

    /// Get a session by agent ID
    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
        let agents = self.agents.read().await;
        let session_id = agents.get(agent_id)?;

        let sessions = self.sessions.read().await;
        sessions.get(session_id).map(|s| s.info.clone())
    }

    /// Get a session by session ID
    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
        let sessions = self.sessions.read().await;
        sessions.get(&session_id).map(|s| s.info.clone())
    }

    /// Join a session as a viewer, returns channels and sends StartStream to agent
    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId, viewer_name: String) -> Option<(FrameReceiver, InputSender)> {
        let mut sessions = self.sessions.write().await;
        let session_data = sessions.get_mut(&session_id)?;

        let was_empty = session_data.viewers.is_empty();

        // Add viewer info
        let viewer_info = ViewerInfo {
            id: viewer_id.clone(),
            name: viewer_name.clone(),
            connected_at: chrono::Utc::now(),
        };
        session_data.viewers.insert(viewer_id.clone(), viewer_info);

        // Update session info
        session_data.info.viewer_count = session_data.viewers.len();
        session_data.info.viewers = session_data.viewers.values().cloned().collect();

        let frame_rx = session_data.frame_tx.subscribe();
        let input_tx = session_data.input_tx.clone();

        // If this is the first viewer, send StartStream to agent
        if was_empty {
            tracing::info!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
            Self::send_start_stream_internal(session_data, &viewer_id).await;
        } else {
            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
        }

        Some((frame_rx, input_tx))
    }

    /// Internal helper to send StartStream message
    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let start_stream = proto::Message {
            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
                viewer_id: viewer_id.to_string(),
                display_id: 0, // Primary display
            })),
        };

        let mut buf = Vec::new();
        if start_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Leave a session as a viewer, sends StopStream if no viewers left
    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
            session_data.viewers.remove(viewer_id);
            session_data.info.viewer_count = session_data.viewers.len();
            session_data.info.viewers = session_data.viewers.values().cloned().collect();

            // If no more viewers, send StopStream to agent
            if session_data.viewers.is_empty() {
                tracing::info!("Last viewer {} ({}) left session {}, sending StopStream",
                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
                Self::send_stop_stream_internal(session_data, viewer_id).await;
            } else {
                tracing::info!("Viewer {} ({}) left session {}",
                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
            }
        }
    }

    /// Internal helper to send StopStream message
    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let stop_stream = proto::Message {
            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
                viewer_id: viewer_id.to_string(),
            })),
        };

        let mut buf = Vec::new();
        if stop_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.info.viewers.clear();
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            drop(sessions);
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Disconnect a session by sending a disconnect message to the agent
    /// Returns true if the message was sent successfully
    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            // Create disconnect message
            use crate::proto;
            use prost::Message;

            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if disconnect_msg.encode(&mut buf).is_ok() {
                // Send via input channel (will be forwarded to agent's WebSocket)
                if session_data.input_tx.send(buf).await.is_ok() {
                    return true;
                }
            }
        }
        false
    }

    /// List all active sessions
    pub async fn list_sessions(&self) -> Vec<Session> {
        let sessions = self.sessions.read().await;
        sessions.values().map(|s| s.info.clone()).collect()
    }
}

impl Default for SessionManager {
    fn default() -> Self {
        Self::new()
    }
}


     1→//! Database module
     2→//!
     3→//! Handles session logging and persistence.
     4→//! Optional for MVP - sessions are kept in memory only.
     5→
     6→use anyhow::Result;
     7→
     8→/// Database connection pool (placeholder)
     9→#[derive(Clone)]
    10→pub struct Database {
    11→    // TODO: Add sqlx pool when PostgreSQL is needed
    12→    _placeholder: (),
    13→}
    14→
    15→impl Database {
    16→    /// Initialize database connection
    17→    pub async fn init(_database_url: &str) -> Result<Self> {
    18→        // TODO: Initialize PostgreSQL connection pool
    19→        Ok(Self { _placeholder: () })
    20→    }
    21→}
    22→
    23→/// Session event for audit logging
    24→#[derive(Debug)]
    25→pub struct SessionEvent {
    26→    pub session_id: String,
    27→    pub event_type: SessionEventType,
    28→    pub details: Option<String>,
    29→}
    30→
    31→#[derive(Debug)]
    32→pub enum SessionEventType {
    33→    Started,
    34→    ViewerJoined,
    35→    ViewerLeft,
    36→    Ended,
    37→}
    38→
    39→impl Database {
    40→    /// Log a session event (placeholder)
    41→    pub async fn log_session_event(&self, _event: SessionEvent) -> Result<()> {
    42→        // TODO: Insert into connect_session_events table
    43→        Ok(())
    44→    }
    45→}
    46→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! Database module
//!
//! Handles session logging and persistence.
//! Optional for MVP - sessions are kept in memory only.

use anyhow::Result;

/// Database connection pool (placeholder)
#[derive(Clone)]
pub struct Database {
    // TODO: Add sqlx pool when PostgreSQL is needed
    _placeholder: (),
}

impl Database {
    /// Initialize database connection
    pub async fn init(_database_url: &str) -> Result<Self> {
        // TODO: Initialize PostgreSQL connection pool
        Ok(Self { _placeholder: () })
    }
}

/// Session event for audit logging
#[derive(Debug)]
pub struct SessionEvent {
    pub session_id: String,
    pub event_type: SessionEventType,
    pub details: Option<String>,
}

#[derive(Debug)]
pub enum SessionEventType {
    Started,
    ViewerJoined,
    ViewerLeft,
    Ended,
}

impl Database {
    /// Log a session event (placeholder)
    pub async fn log_session_event(&self, _event: SessionEvent) -> Result<()> {
        // TODO: Insert into connect_session_events table
        Ok(())
    }
}


     1→//! Server configuration
     2→
     3→use anyhow::Result;
     4→use serde::Deserialize;
     5→use std::env;
     6→
     7→#[derive(Debug, Clone, Deserialize)]
     8→pub struct Config {
     9→    /// Address to listen on (e.g., "0.0.0.0:8080")
    10→    pub listen_addr: String,
    11→
    12→    /// Database URL (optional for MVP)
    13→    pub database_url: Option<String>,
    14→
    15→    /// JWT secret for authentication
    16→    pub jwt_secret: Option<String>,
    17→
    18→    /// Enable debug logging
    19→    pub debug: bool,
    20→}
    21→
    22→impl Config {
    23→    /// Load configuration from environment variables
    24→    pub fn load() -> Result<Self> {
    25→        Ok(Self {
    26→            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
    27→            database_url: env::var("DATABASE_URL").ok(),
    28→            jwt_secret: env::var("JWT_SECRET").ok(),
    29→            debug: env::var("DEBUG")
    30→                .map(|v| v == "1" || v.to_lowercase() == "true")
    31→                .unwrap_or(false),
    32→        })
    33→    }
    34→}
    35→
    36→impl Default for Config {
    37→    fn default() -> Self {
    38→        Self {
    39→            listen_addr: "0.0.0.0:8080".to_string(),
    40→            database_url: None,
    41→            jwt_secret: None,
    42→            debug: false,
    43→        }
    44→    }
    45→}
    46→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! Server configuration

use anyhow::Result;
use serde::Deserialize;
use std::env;

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional for MVP)
    pub database_url: Option<String>,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            jwt_secret: None,
            debug: false,
        }
    }
}


Now let me read the auth and API modules to complete the picture:


     1→//! Authentication module
     2→//!
     3→//! Handles JWT validation for dashboard users and API key
     4→//! validation for agents.
     5→
     6→use axum::{
     7→    extract::FromRequestParts,
     8→    http::{request::Parts, StatusCode},
     9→};
    10→
    11→/// Authenticated user from JWT
    12→#[derive(Debug, Clone)]
    13→pub struct AuthenticatedUser {
    14→    pub user_id: String,
    15→    pub email: String,
    16→    pub roles: Vec<String>,
    17→}
    18→
    19→/// Authenticated agent from API key
    20→#[derive(Debug, Clone)]
    21→pub struct AuthenticatedAgent {
    22→    pub agent_id: String,
    23→    pub org_id: String,
    24→}
    25→
    26→/// Extract authenticated user from request (placeholder for MVP)
    27→#[axum::async_trait]
    28→impl<S> FromRequestParts<S> for AuthenticatedUser
    29→where
    30→    S: Send + Sync,
    31→{
    32→    type Rejection = (StatusCode, &'static str);
    33→
    34→    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
    35→        // TODO: Implement JWT validation
    36→        // For MVP, accept any request
    37→
    38→        // Look for Authorization header
    39→        let _auth_header = parts
    40→            .headers
    41→            .get("Authorization")
    42→            .and_then(|v| v.to_str().ok());
    43→
    44→        // Placeholder - in production, validate JWT
    45→        Ok(AuthenticatedUser {
    46→            user_id: "mvp-user".to_string(),
    47→            email: "mvp@example.com".to_string(),
    48→            roles: vec!["admin".to_string()],
    49→        })
    50→    }
    51→}
    52→
    53→/// Validate an agent API key (placeholder for MVP)
    54→pub fn validate_agent_key(_api_key: &str) -> Option<AuthenticatedAgent> {
    55→    // TODO: Implement actual API key validation
    56→    // For MVP, accept any key
    57→    Some(AuthenticatedAgent {
    58→        agent_id: "mvp-agent".to_string(),
    59→        org_id: "mvp-org".to_string(),
    60→    })
    61→}
    62→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! Authentication module
//!
//! Handles JWT validation for dashboard users and API key
//! validation for agents.

use axum::{
    extract::FromRequestParts,
    http::{request::Parts, StatusCode},
};

/// Authenticated user from JWT
#[derive(Debug, Clone)]
pub struct AuthenticatedUser {
    pub user_id: String,
    pub email: String,
    pub roles: Vec<String>,
}

/// Authenticated agent from API key
#[derive(Debug, Clone)]
pub struct AuthenticatedAgent {
    pub agent_id: String,
    pub org_id: String,
}

/// Extract authenticated user from request (placeholder for MVP)
#[axum::async_trait]
impl<S> FromRequestParts<S> for AuthenticatedUser
where
    S: Send + Sync,
{
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        // TODO: Implement JWT validation
        // For MVP, accept any request

        // Look for Authorization header
        let _auth_header = parts
            .headers
            .get("Authorization")
            .and_then(|v| v.to_str().ok());

        // Placeholder - in production, validate JWT
        Ok(AuthenticatedUser {
            user_id: "mvp-user".to_string(),
            email: "mvp@example.com".to_string(),
            roles: vec!["admin".to_string()],
        })
    }
}

/// Validate an agent API key (placeholder for MVP)
pub fn validate_agent_key(_api_key: &str) -> Option<AuthenticatedAgent> {
    // TODO: Implement actual API key validation
    // For MVP, accept any key
    Some(AuthenticatedAgent {
        agent_id: "mvp-agent".to_string(),
        org_id: "mvp-org".to_string(),
    })
}


     1→//! REST API endpoints
     2→
     3→use axum::{
     4→    extract::{Path, State},
     5→    Json,
     6→};
     7→use serde::Serialize;
     8→use uuid::Uuid;
     9→
    10→use crate::session::SessionManager;
    11→
    12→/// Viewer info returned by API
    13→#[derive(Debug, Serialize)]
    14→pub struct ViewerInfoApi {
    15→    pub id: String,
    16→    pub name: String,
    17→    pub connected_at: String,
    18→}
    19→
    20→impl From<crate::session::ViewerInfo> for ViewerInfoApi {
    21→    fn from(v: crate::session::ViewerInfo) -> Self {
    22→        Self {
    23→            id: v.id,
    24→            name: v.name,
    25→            connected_at: v.connected_at.to_rfc3339(),
    26→        }
    27→    }
    28→}
    29→
    30→/// Session info returned by API
    31→#[derive(Debug, Serialize)]
    32→pub struct SessionInfo {
    33→    pub id: String,
    34→    pub agent_id: String,
    35→    pub agent_name: String,
    36→    pub started_at: String,
    37→    pub viewer_count: usize,
    38→    pub viewers: Vec<ViewerInfoApi>,
    39→    pub is_streaming: bool,
    40→    pub is_online: bool,
    41→    pub is_persistent: bool,
    42→    pub last_heartbeat: String,
    43→    pub os_version: Option<String>,
    44→    pub is_elevated: bool,
    45→    pub uptime_secs: i64,
    46→    pub display_count: i32,
    47→}
    48→
    49→impl From<crate::session::Session> for SessionInfo {
    50→    fn from(s: crate::session::Session) -> Self {
    51→        Self {
    52→            id: s.id.to_string(),
    53→            agent_id: s.agent_id,
    54→            agent_name: s.agent_name,
    55→            started_at: s.started_at.to_rfc3339(),
    56→            viewer_count: s.viewer_count,
    57→            viewers: s.viewers.into_iter().map(ViewerInfoApi::from).collect(),
    58→            is_streaming: s.is_streaming,
    59→            is_online: s.is_online,
    60→            is_persistent: s.is_persistent,
    61→            last_heartbeat: s.last_heartbeat.to_rfc3339(),
    62→            os_version: s.os_version,
    63→            is_elevated: s.is_elevated,
    64→            uptime_secs: s.uptime_secs,
    65→            display_count: s.display_count,
    66→        }
    67→    }
    68→}
    69→
    70→/// List all active sessions
    71→pub async fn list_sessions(
    72→    State(sessions): State<SessionManager>,
    73→) -> Json<Vec<SessionInfo>> {
    74→    let sessions = sessions.list_sessions().await;
    75→    Json(sessions.into_iter().map(SessionInfo::from).collect())
    76→}
    77→
    78→/// Get a specific session by ID
    79→pub async fn get_session(
    80→    State(sessions): State<SessionManager>,
    81→    Path(id): Path<String>,
    82→) -> Result<Json<SessionInfo>, (axum::http::StatusCode, &'static str)> {
    83→    let session_id = Uuid::parse_str(&id)
    84→        .map_err(|_| (axum::http::StatusCode::BAD_REQUEST, "Invalid session ID"))?;
    85→
    86→    let session = sessions.get_session(session_id).await
    87→        .ok_or((axum::http::StatusCode::NOT_FOUND, "Session not found"))?;
    88→
    89→    Ok(Json(SessionInfo::from(session)))
    90→}
    91→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! REST API endpoints

use axum::{
    extract::{Path, State},
    Json,
};
use serde::Serialize;
use uuid::Uuid;

use crate::session::SessionManager;

/// Viewer info returned by API
#[derive(Debug, Serialize)]
pub struct ViewerInfoApi {
    pub id: String,
    pub name: String,
    pub connected_at: String,
}

impl From<crate::session::ViewerInfo> for ViewerInfoApi {
    fn from(v: crate::session::ViewerInfo) -> Self {
        Self {
            id: v.id,
            name: v.name,
            connected_at: v.connected_at.to_rfc3339(),
        }
    }
}

/// Session info returned by API
#[derive(Debug, Serialize)]
pub struct SessionInfo {
    pub id: String,
    pub agent_id: String,
    pub agent_name: String,
    pub started_at: String,
    pub viewer_count: usize,
    pub viewers: Vec<ViewerInfoApi>,
    pub is_streaming: bool,
    pub is_online: bool,
    pub is_persistent: bool,
    pub last_heartbeat: String,
    pub os_version: Option<String>,
    pub is_elevated: bool,
    pub uptime_secs: i64,
    pub display_count: i32,
}

impl From<crate::session::Session> for SessionInfo {
    fn from(s: crate::session::Session) -> Self {
        Self {
            id: s.id.to_string(),
            agent_id: s.agent_id,
            agent_name: s.agent_name,
            started_at: s.started_at.to_rfc3339(),
            viewer_count: s.viewer_count,
            viewers: s.viewers.into_iter().map(ViewerInfoApi::from).collect(),
            is_streaming: s.is_streaming,
            is_online: s.is_online,
            is_persistent: s.is_persistent,
            last_heartbeat: s.last_heartbeat.to_rfc3339(),
            os_version: s.os_version,
            is_elevated: s.is_elevated,
            uptime_secs: s.uptime_secs,
            display_count: s.display_count,
        }
    }
}

/// List all active sessions
pub async fn list_sessions(
    State(sessions): State<SessionManager>,
) -> Json<Vec<SessionInfo>> {
    let sessions = sessions.list_sessions().await;
    Json(sessions.into_iter().map(SessionInfo::from).collect())
}

/// Get a specific session by ID
pub async fn get_session(
    State(sessions): State<SessionManager>,
    Path(id): Path<String>,
) -> Result<Json<SessionInfo>, (axum::http::StatusCode, &'static str)> {
    let session_id = Uuid::parse_str(&id)
        .map_err(|_| (axum::http::StatusCode::BAD_REQUEST, "Invalid session ID"))?;

    let session = sessions.get_session(session_id).await
        .ok_or((axum::http::StatusCode::NOT_FOUND, "Session not found"))?;

    Ok(Json(SessionInfo::from(session)))
}


     1→//! WebSocket relay handlers
     2→//!
     3→//! Handles WebSocket connections from agents and viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→use axum::{
     7→    extract::{
     8→        ws::{Message, WebSocket, WebSocketUpgrade},
     9→        Query, State,
    10→    },
    11→    response::IntoResponse,
    12→};
    13→use futures_util::{SinkExt, StreamExt};
    14→use prost::Message as ProstMessage;
    15→use serde::Deserialize;
    16→use tracing::{error, info, warn};
    17→use uuid::Uuid;
    18→
    19→use crate::proto;
    20→use crate::session::SessionManager;
    21→use crate::AppState;
    22→
    23→#[derive(Debug, Deserialize)]
    24→pub struct AgentParams {
    25→    agent_id: String,
    26→    #[serde(default)]
    27→    agent_name: Option<String>,
    28→    #[serde(default)]
    29→    support_code: Option<String>,
    30→    #[serde(default)]
    31→    hostname: Option<String>,
    32→}
    33→
    34→#[derive(Debug, Deserialize)]
    35→pub struct ViewerParams {
    36→    session_id: String,
    37→    #[serde(default = "default_viewer_name")]
    38→    viewer_name: String,
    39→}
    40→
    41→fn default_viewer_name() -> String {
    42→    "Technician".to_string()
    43→}
    44→
    45→/// WebSocket handler for agent connections
    46→pub async fn agent_ws_handler(
    47→    ws: WebSocketUpgrade,
    48→    State(state): State<AppState>,
    49→    Query(params): Query<AgentParams>,
    50→) -> impl IntoResponse {
    51→    let agent_id = params.agent_id;
    52→    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    53→    let support_code = params.support_code;
    54→    let sessions = state.sessions.clone();
    55→    let support_codes = state.support_codes.clone();
    56→
    57→    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
    58→}
    59→
    60→/// WebSocket handler for viewer connections
    61→pub async fn viewer_ws_handler(
    62→    ws: WebSocketUpgrade,
    63→    State(state): State<AppState>,
    64→    Query(params): Query<ViewerParams>,
    65→) -> impl IntoResponse {
    66→    let session_id = params.session_id;
    67→    let viewer_name = params.viewer_name;
    68→    let sessions = state.sessions.clone();
    69→
    70→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id, viewer_name))
    71→}
    72→
    73→/// Handle an agent WebSocket connection
    74→async fn handle_agent_connection(
    75→    socket: WebSocket,
    76→    sessions: SessionManager,
    77→    support_codes: crate::support_codes::SupportCodeManager,
    78→    agent_id: String,
    79→    agent_name: String,
    80→    support_code: Option<String>,
    81→) {
    82→    info!("Agent connected: {} ({})", agent_name, agent_id);
    83→
    84→    let (mut ws_sender, mut ws_receiver) = socket.split();
    85→
    86→    // If a support code was provided, check if it's valid
    87→    if let Some(ref code) = support_code {
    88→        // Check if the code is cancelled or invalid
    89→        if support_codes.is_cancelled(code).await {
    90→            warn!("Agent tried to connect with cancelled code: {}", code);
    91→            // Send disconnect message to agent
    92→            let disconnect_msg = proto::Message {
    93→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
    94→                    reason: "Support session was cancelled by technician".to_string(),
    95→                })),
    96→            };
    97→            let mut buf = Vec::new();
    98→            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
    99→                let _ = ws_sender.send(Message::Binary(buf.into())).await;
   100→            }
   101→            let _ = ws_sender.close().await;
   102→            return;
   103→        }
   104→    }
   105→
   106→    // Register the agent and get channels
   107→    // Persistent agents (no support code) keep their session when disconnected
   108→    let is_persistent = support_code.is_none();
   109→    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;
   110→
   111→    info!("Session created: {} (agent in idle mode)", session_id);
   112→
   113→    // If a support code was provided, mark it as connected
   114→    if let Some(ref code) = support_code {
   115→        info!("Linking support code {} to session {}", code, session_id);
   116→        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
   117→        support_codes.link_session(code, session_id).await;
   118→    }
   119→
   120→    // Use Arc<Mutex> for sender so we can use it from multiple places
   121→    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
   122→    let ws_sender_input = ws_sender.clone();
   123→    let ws_sender_cancel = ws_sender.clone();
   124→
   125→    // Task to forward input events from viewers to agent
   126→    let input_forward = tokio::spawn(async move {
   127→        while let Some(input_data) = input_rx.recv().await {
   128→            let mut sender = ws_sender_input.lock().await;
   129→            if sender.send(Message::Binary(input_data.into())).await.is_err() {
   130→                break;
   131→            }
   132→        }
   133→    });
   134→
   135→    let sessions_cleanup = sessions.clone();
   136→    let sessions_status = sessions.clone();
   137→    let support_codes_cleanup = support_codes.clone();
   138→    let support_code_cleanup = support_code.clone();
   139→    let support_code_check = support_code.clone();
   140→    let support_codes_check = support_codes.clone();
   141→
   142→    // Task to check for cancellation every 2 seconds
   143→    let cancel_check = tokio::spawn(async move {
   144→        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
   145→        loop {
   146→            interval.tick().await;
   147→            if let Some(ref code) = support_code_check {
   148→                if support_codes_check.is_cancelled(code).await {
   149→                    info!("Support code {} was cancelled, disconnecting agent", code);
   150→                    // Send disconnect message
   151→                    let disconnect_msg = proto::Message {
   152→                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   153→                            reason: "Support session was cancelled by technician".to_string(),
   154→                        })),
   155→                    };
   156→                    let mut buf = Vec::new();
   157→                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
   158→                        let mut sender = ws_sender_cancel.lock().await;
   159→                        let _ = sender.send(Message::Binary(buf.into())).await;
   160→                        let _ = sender.close().await;
   161→                    }
   162→                    break;
   163→                }
   164→            }
   165→        }
   166→    });
   167→
   168→    // Main loop: receive messages from agent
   169→    while let Some(msg) = ws_receiver.next().await {
   170→        match msg {
   171→            Ok(Message::Binary(data)) => {
   172→                // Try to decode as protobuf message
   173→                match proto::Message::decode(data.as_ref()) {
   174→                    Ok(proto_msg) => {
   175→                        match &proto_msg.payload {
   176→                            Some(proto::message::Payload::VideoFrame(_)) => {
   177→                                // Broadcast frame to all viewers (only sent when streaming)
   178→                                let _ = frame_tx.send(data.to_vec());
   179→                            }
   180→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   181→                                // Broadcast chat message to all viewers
   182→                                info!("Chat from client: {}", chat.content);
   183→                                let _ = frame_tx.send(data.to_vec());
   184→                            }
   185→                            Some(proto::message::Payload::AgentStatus(status)) => {
   186→                                // Update session with agent status
   187→                                sessions_status.update_agent_status(
   188→                                    session_id,
   189→                                    Some(status.os_version.clone()),
   190→                                    status.is_elevated,
   191→                                    status.uptime_secs,
   192→                                    status.display_count,
   193→                                    status.is_streaming,
   194→                                ).await;
   195→                                info!("Agent status update: {} - streaming={}, uptime={}s",
   196→                                    status.hostname, status.is_streaming, status.uptime_secs);
   197→                            }
   198→                            Some(proto::message::Payload::Heartbeat(_)) => {
   199→                                // Update heartbeat timestamp
   200→                                sessions_status.update_heartbeat(session_id).await;
   201→                            }
   202→                            Some(proto::message::Payload::HeartbeatAck(_)) => {
   203→                                // Agent acknowledged our heartbeat
   204→                                sessions_status.update_heartbeat(session_id).await;
   205→                            }
   206→                            _ => {}
   207→                        }
   208→                    }
   209→                    Err(e) => {
   210→                        warn!("Failed to decode agent message: {}", e);
   211→                    }
   212→                }
   213→            }
   214→            Ok(Message::Close(_)) => {
   215→                info!("Agent disconnected: {}", agent_id);
   216→                break;
   217→            }
   218→            Ok(Message::Ping(data)) => {
   219→                // Pong is handled automatically by axum
   220→                let _ = data;
   221→            }
   222→            Ok(_) => {}
   223→            Err(e) => {
   224→                error!("WebSocket error from agent {}: {}", agent_id, e);
   225→                break;
   226→            }
   227→        }
   228→    }
   229→
   230→    // Cleanup
   231→    input_forward.abort();
   232→    cancel_check.abort();
   233→    // Mark agent as disconnected (persistent agents stay in list as offline)
   234→    sessions_cleanup.mark_agent_disconnected(session_id).await;
   235→
   236→    // Mark support code as completed if one was used (unless cancelled)
   237→    if let Some(ref code) = support_code_cleanup {
   238→        if !support_codes_cleanup.is_cancelled(code).await {
   239→            support_codes_cleanup.mark_completed(code).await;
   240→            info!("Support code {} marked as completed", code);
   241→        }
   242→    }
   243→
   244→    info!("Session {} ended", session_id);
   245→}
   246→
   247→/// Handle a viewer WebSocket connection
   248→async fn handle_viewer_connection(
   249→    socket: WebSocket,
   250→    sessions: SessionManager,
   251→    session_id_str: String,
   252→    viewer_name: String,
   253→) {
   254→    // Parse session ID
   255→    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
   256→        Ok(id) => id,
   257→        Err(_) => {
   258→            warn!("Invalid session ID: {}", session_id_str);
   259→            return;
   260→        }
   261→    };
   262→
   263→    // Generate unique viewer ID
   264→    let viewer_id = Uuid::new_v4().to_string();
   265→
   266→    // Join the session (this sends StartStream to agent if first viewer)
   267→    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone(), viewer_name.clone()).await {
   268→        Some(channels) => channels,
   269→        None => {
   270→            warn!("Session not found: {}", session_id);
   271→            return;
   272→        }
   273→    };
   274→
   275→    info!("Viewer {} ({}) joined session: {}", viewer_name, viewer_id, session_id);
   276→
   277→    let (mut ws_sender, mut ws_receiver) = socket.split();
   278→
   279→    // Task to forward frames from agent to this viewer
   280→    let frame_forward = tokio::spawn(async move {
   281→        while let Ok(frame_data) = frame_rx.recv().await {
   282→            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
   283→                break;
   284→            }
   285→        }
   286→    });
   287→
   288→    let sessions_cleanup = sessions.clone();
   289→    let viewer_id_cleanup = viewer_id.clone();
   290→
   291→    // Main loop: receive input from viewer and forward to agent
   292→    while let Some(msg) = ws_receiver.next().await {
   293→        match msg {
   294→            Ok(Message::Binary(data)) => {
   295→                // Try to decode as protobuf message
   296→                match proto::Message::decode(data.as_ref()) {
   297→                    Ok(proto_msg) => {
   298→                        match &proto_msg.payload {
   299→                            Some(proto::message::Payload::MouseEvent(_)) |
   300→                            Some(proto::message::Payload::KeyEvent(_)) |
   301→                            Some(proto::message::Payload::SpecialKey(_)) => {
   302→                                // Forward input to agent
   303→                                let _ = input_tx.send(data.to_vec()).await;
   304→                            }
   305→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   306→                                // Forward chat message to agent
   307→                                info!("Chat from technician: {}", chat.content);
   308→                                let _ = input_tx.send(data.to_vec()).await;
   309→                            }
   310→                            _ => {}
   311→                        }
   312→                    }
   313→                    Err(e) => {
   314→                        warn!("Failed to decode viewer message: {}", e);
   315→                    }
   316→                }
   317→            }
   318→            Ok(Message::Close(_)) => {
   319→                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
   320→                break;
   321→            }
   322→            Ok(_) => {}
   323→            Err(e) => {
   324→                error!("WebSocket error from viewer {}: {}", viewer_id, e);
   325→                break;
   326→            }
   327→        }
   328→    }
   329→
   330→    // Cleanup (this sends StopStream to agent if last viewer)
   331→    frame_forward.abort();
   332→    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
   333→    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
   334→}
   335→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! WebSocket relay handlers
//!
//! Handles WebSocket connections from agents and viewers,
//! relaying video frames and input events between them.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use serde::Deserialize;
use tracing::{error, info, warn};
use uuid::Uuid;

use crate::proto;
use crate::session::SessionManager;
use crate::AppState;

#[derive(Debug, Deserialize)]
pub struct AgentParams {
    agent_id: String,
    #[serde(default)]
    agent_name: Option<String>,
    #[serde(default)]
    support_code: Option<String>,
    #[serde(default)]
    hostname: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
}

fn default_viewer_name() -> String {
    "Technician".to_string()
}

/// WebSocket handler for agent connections
pub async fn agent_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<AgentParams>,
) -> impl IntoResponse {
    let agent_id = params.agent_id;
    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    let support_code = params.support_code;
    let sessions = state.sessions.clone();
    let support_codes = state.support_codes.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
}

/// WebSocket handler for viewer connections
pub async fn viewer_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<ViewerParams>,
) -> impl IntoResponse {
    let session_id = params.session_id;
    let viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();

    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id, viewer_name))
}

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {
    info!("Agent connected: {} ({})", agent_name, agent_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // If a support code was provided, check if it's valid
    if let Some(ref code) = support_code {
        // Check if the code is cancelled or invalid
        if support_codes.is_cancelled(code).await {
            warn!("Agent tried to connect with cancelled code: {}", code);
            // Send disconnect message to agent
            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: "Support session was cancelled by technician".to_string(),
                })),
            };
            let mut buf = Vec::new();
            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                let _ = ws_sender.send(Message::Binary(buf.into())).await;
            }
            let _ = ws_sender.close().await;
            return;
        }
    }

    // Register the agent and get channels
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;

    info!("Session created: {} (agent in idle mode)", session_id);

    // If a support code was provided, mark it as connected
    if let Some(ref code) = support_code {
        info!("Linking support code {} to session {}", code, session_id);
        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
        support_codes.link_session(code, session_id).await;
    }

    // Use Arc<Mutex> for sender so we can use it from multiple places
    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
    let ws_sender_input = ws_sender.clone();
    let ws_sender_cancel = ws_sender.clone();

    // Task to forward input events from viewers to agent
    let input_forward = tokio::spawn(async move {
        while let Some(input_data) = input_rx.recv().await {
            let mut sender = ws_sender_input.lock().await;
            if sender.send(Message::Binary(input_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let sessions_status = sessions.clone();
    let support_codes_cleanup = support_codes.clone();
    let support_code_cleanup = support_code.clone();
    let support_code_check = support_code.clone();
    let support_codes_check = support_codes.clone();

    // Task to check for cancellation every 2 seconds
    let cancel_check = tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
        loop {
            interval.tick().await;
            if let Some(ref code) = support_code_check {
                if support_codes_check.is_cancelled(code).await {
                    info!("Support code {} was cancelled, disconnecting agent", code);
                    // Send disconnect message
                    let disconnect_msg = proto::Message {
                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                            reason: "Support session was cancelled by technician".to_string(),
                        })),
                    };
                    let mut buf = Vec::new();
                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                        let mut sender = ws_sender_cancel.lock().await;
                        let _ = sender.send(Message::Binary(buf.into())).await;
                        let _ = sender.close().await;
                    }
                    break;
                }
            }
        }
    });

    // Main loop: receive messages from agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::VideoFrame(_)) => {
                                // Broadcast frame to all viewers (only sent when streaming)
                                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);
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::AgentStatus(status)) => {
                                // Update session with agent status
                                sessions_status.update_agent_status(
                                    session_id,
                                    Some(status.os_version.clone()),
                                    status.is_elevated,
                                    status.uptime_secs,
                                    status.display_count,
                                    status.is_streaming,
                                ).await;
                                info!("Agent status update: {} - streaming={}, uptime={}s",
                                    status.hostname, status.is_streaming, status.uptime_secs);
                            }
                            Some(proto::message::Payload::Heartbeat(_)) => {
                                // Update heartbeat timestamp
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            Some(proto::message::Payload::HeartbeatAck(_)) => {
                                // Agent acknowledged our heartbeat
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode agent message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent disconnected: {}", agent_id);
                break;
            }
            Ok(Message::Ping(data)) => {
                // Pong is handled automatically by axum
                let _ = data;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from agent {}: {}", agent_id, e);
                break;
            }
        }
    }

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(session_id).await;

    // Mark support code as completed if one was used (unless cancelled)
    if let Some(ref code) = support_code_cleanup {
        if !support_codes_cleanup.is_cancelled(code).await {
            support_codes_cleanup.mark_completed(code).await;
            info!("Support code {} marked as completed", code);
        }
    }

    info!("Session {} ended", session_id);
}

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    session_id_str: String,
    viewer_name: String,
) {
    // Parse session ID
    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
        Ok(id) => id,
        Err(_) => {
            warn!("Invalid session ID: {}", session_id_str);
            return;
        }
    };

    // Generate unique viewer ID
    let viewer_id = Uuid::new_v4().to_string();

    // Join the session (this sends StartStream to agent if first viewer)
    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, viewer_id, session_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // Task to forward frames from agent to this viewer
    let frame_forward = tokio::spawn(async move {
        while let Ok(frame_data) = frame_rx.recv().await {
            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let viewer_id_cleanup = viewer_id.clone();

    // Main loop: receive input from viewer and forward to agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::MouseEvent(_)) |
                            Some(proto::message::Payload::KeyEvent(_)) |
                            Some(proto::message::Payload::SpecialKey(_)) => {
                                // Forward input to agent
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Forward chat message to agent
                                info!("Chat from technician: {}", chat.content);
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode viewer message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
                break;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from viewer {}: {}", viewer_id, e);
                break;
            }
        }
    }

    // Cleanup (this sends StopStream to agent if last viewer)
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}