The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\relay\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
     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→    http::StatusCode,
    13→};
    14→use futures_util::{SinkExt, StreamExt};
    15→use prost::Message as ProstMessage;
    16→use serde::Deserialize;
    17→use tracing::{error, info, warn};
    18→use uuid::Uuid;
    19→
    20→use crate::proto;
    21→use crate::session::SessionManager;
    22→use crate::db::{self, Database};
    23→use crate::AppState;
    24→
    25→#[derive(Debug, Deserialize)]
    26→pub struct AgentParams {
    27→    agent_id: String,

//! 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,
    http::StatusCode,
};
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::db::{self, Database};
use crate::AppState;

//! 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::db::{self, Database};
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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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();
    let db = state.db.clone();

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    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);

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

    // 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;

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_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;

    // Database: end session and mark machine offline
    if let Some(ref db) = db {
        // End the session record
        let _ = db::sessions::end_session(db.pool(), session_id, "ended").await;

        // Mark machine as offline
        let _ = db::machines::mark_machine_offline(db.pool(), &agent_id).await;

        // Log session ended event
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::SESSION_ENDED,
            None, None, None, None,
        ).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;

            // Database: mark code as completed
            if let Some(ref db) = db {
                let _ = db::support_codes::mark_code_completed(db.pool(), 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,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

    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();
    let viewer_name_cleanup = viewer_name.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;

    // Database: log viewer left event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_LEFT,
            Some(&viewer_id_cleanup),
            Some(&viewer_name_cleanup),
            None, None,
        ).await;
    }

    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}


Conversation compacted


This session is being continued from a previous conversation that ran out of context. The conversation is summarized below:
Analysis:
Let me chronologically analyze this conversation:

1. **Session Start Context**: This is a continuation from a previous session that worked on GuruConnect user management (JWT authentication, Argon2 password hashing, role-based access control).

2. **First User Message**: "Thanks. Connect button doesn't offer native viewer option for testing."
   - User wanted native viewer option in the dashboard's Connect button
   - I added a modal with Web Viewer and Native Viewer options
   - Initially showed command to copy/paste - user was unhappy

3. **Second User Message**: "This doesn't seem like the implementation we discussed"
   - User showed screenshot of the basic command-line approach
   - User wanted a proper protocol handler implementation

4. **Third User Message**: "yes implement the protocol handler - and if it's not found (not installed) it should offer to run the viewer and hopefully retry the protocol handler"
   - I tried to create a separate viewer crate but user corrected me

5. **Fourth User Message**: "Viewer is integrated into the main exe now, double check your notes"
   - Found viewer in agent/src/viewer/
   - main.rs already had `install` and `launch` commands for protocol handler
   - install.rs had protocol handler registration

6. **Fifth User Message**: "build the agent for windows"
   - Built with `cargo build -p guruconnect --release`
   - Uploaded to server at /downloads/guruconnect.exe

7. **Sixth User Message**: "The web launcher should pass the 'install' option. It does not."
   - Modified main.rs to auto-detect if protocol not registered and run install mode
   - Added `is_protocol_handler_registered()` function to install.rs

8. **Seventh User Message**: "btw: this error after launching protocol handler" (showed "Invalid URL: Missing session ID")
   - Problem: URL parser treated "view" as host, not path in `guruconnect://view/SESSION_ID`
   - Fixed by reading action from `url.host_str()` instead of path

9. **Eighth User Message**: "We'll come back to that. Let's deal with general workflow problems..."
   - User identified SECURITY CONCERN: machines appearing in dashboard without authentication
   - This shouldn't be possible after adding login requirements

10. **Ninth User Message**: "yes, fix it"
    - Started securing endpoints
    - Added `AuthenticatedUser` to REST API handlers
    - Working on WebSocket handlers when summary was requested

Key security findings:
- auth_layer only injected JWT config, didn't enforce auth
- Unprotected: /api/sessions, /api/codes, /api/machines, /ws/agent, /ws/viewer
- Anyone could list machines or connect as agent

Files modified for security:
- main.rs: Added AuthenticatedUser to create_code, list_codes, cancel_code, list_sessions, get_session, disconnect_session, list_machines, get_machine, get_machine_history, delete_machine
- relay/mod.rs: Started adding StatusCode import for WebSocket auth

Summary:
1. Primary Request and Intent:
   - User initially wanted native viewer option in the Connect button on dashboard
   - User wanted proper `guruconnect://` protocol handler implementation with fallback to download if not installed
   - User identified a CRITICAL SECURITY CONCERN: machines appearing in dashboard despite login requirements - this is the current priority
   - User explicitly requested to fix the security vulnerabilities allowing unauthenticated access

2. Key Technical Concepts:
   - Custom protocol handler registration on Windows (registry keys in HKCR or HKCU\Software\Classes)
   - URL parsing quirks: `guruconnect://view/SESSION_ID` has "view" as host, not path
   - JWT authentication with Axum extractors (AuthenticatedUser, AdminUser)
   - WebSocket authentication for both agents and viewers
   - The auth_layer middleware only injects JWT config - doesn't enforce auth
   - Authentication enforcement requires using extractor parameters in handlers

3. Files and Code Sections:

   - **server/src/main.rs** (CRITICAL - Security Fix)
     - Added `AuthenticatedUser` import and requirement to all REST API handlers
     - Handlers secured: create_code, list_codes, cancel_code, list_sessions, get_session, disconnect_session, list_machines, get_machine, get_machine_history, delete_machine
     ```rust
     use auth::{JwtConfig, hash_password, generate_random_password, AuthenticatedUser};

     async fn list_sessions(
         _user: AuthenticatedUser,  // Require authentication
         State(state): State<AppState>,
     ) -> Json<Vec<api::SessionInfo>> {
     ```

   - **server/src/relay/mod.rs** (IN PROGRESS - WebSocket Security)
     - Added StatusCode import for rejecting unauthorized connections
     - Need to add JWT validation to viewer_ws_handler
     - Need to add support_code or API key validation to agent_ws_handler
     ```rust
     use axum::{
         extract::{
             ws::{Message, WebSocket, WebSocketUpgrade},
             Query, State,
         },
         response::IntoResponse,
         http::StatusCode,  // Added for auth rejection
     };
     ```

   - **agent/src/main.rs** (Protocol Handler Auto-Install)
     - Modified to auto-install if protocol handler not registered
     ```rust
     None => {
         if let Some(code) = cli.support_code {
             run_agent_mode(Some(code))
         } else {
             if !install::is_protocol_handler_registered() {
                 info!("Protocol handler not registered, running installer");
                 run_install(false)
             } else {
                 run_agent_mode(None)
             }
         }
     }
     ```

   - **agent/src/install.rs** (Protocol Handler Detection & URL Parsing Fix)
     - Added `is_protocol_handler_registered()` function
     - Fixed URL parsing to read action from host, session from path
     ```rust
     pub fn is_protocol_handler_registered() -> bool {
         // Check HKCR\guruconnect and HKCU\Software\Classes\guruconnect
     }

     pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
         // The "action" (view/connect) is parsed as the host
         let action = url.host_str()
             .ok_or_else(|| anyhow!("Missing action in URL"))?;
         // The session ID is the first path segment
         let path = url.path().trim_start_matches('/');
     }
     ```

   - **server/static/dashboard.html** (Native Viewer UI)
     - Added connect options modal with Web/Native viewer choices
     - Protocol handler launch with fallback to download
     ```javascript
     function launchNativeViewer() {
         const protocolUrl = `guruconnect://view/${connectSessionId}?server=${serverUrl}`;
         // Uses iframe to launch, detects blur for success, shows download if fails
     }
     ```

4. Errors and Fixes:
   - **"Invalid URL: Missing session ID"**: URL parser treated "view" as host in `guruconnect://view/SESSION_ID`. Fixed by reading action from `url.host_str()` instead of path.
   - **User feedback on initial native viewer**: "This doesn't seem like the implementation we discussed" - user didn't want copy/paste command, wanted protocol handler
   - **User feedback on viewer location**: "Viewer is integrated into the main exe now, double check your notes" - viewer is in agent/src/viewer/, not separate crate
   - **User feedback on auto-install**: "The web launcher should pass the 'install' option. It does not." - modified exe to auto-detect and install if protocol not registered

5. Problem Solving:
   - Solved: Native viewer protocol handler implementation with auto-install
   - Solved: URL parsing for custom protocol
   - In Progress: Security vulnerabilities in API endpoints
   - Completed: REST API endpoints now require AuthenticatedUser
   - In Progress: WebSocket handlers still need authentication

6. All User Messages:
   - "Thanks. Connect button doesn't offer native viewer option for testing."
   - "This doesn't seem like the implementation we discussed" (with screenshot)
   - "yes implement the protocol handler - and if it's not found (not installed) it should offer to run the viewer and hopefully retry the protocol handler"
   - "Viewer is integrated into the main exe now, double check your notes"
   - "build the agent for windows"
   - "The web launcher should pass the 'install' option. It does not."
   - "btw: this error after launching protocol handler" (showed Invalid URL error)
   - "We'll come back to that. Let's deal with general workflow problems. We instituted login requirements last night, but this morning there are a number of machines listed in the dashboard. This shouldn't have been possible. We need to determine how these devices managed to show up in the list. It seems like a security concern to me."
   - "yes, fix it"

7. Pending Tasks (from todo list):
   - [COMPLETED] Add AuthenticatedUser to REST API endpoints (sessions, codes, machines)
   - [IN PROGRESS] Add JWT auth to viewer WebSocket
   - [PENDING] Add agent authentication (API key or support code required)
   - [PENDING] Rebuild and deploy server

8. Current Work:
   Fixing security vulnerabilities. Just completed adding `AuthenticatedUser` to all REST API endpoints in main.rs. Started modifying relay/mod.rs to add authentication to WebSocket handlers - added StatusCode import but have not yet implemented the actual auth checks for viewer_ws_handler and agent_ws_handler.

9. Optional Next Step:
   Continue securing WebSocket handlers in relay/mod.rs:
   - Add JWT token validation to viewer_ws_handler (require token in query params)
   - Add authentication to agent_ws_handler (require valid support_code OR API key)

   User's explicit request: "yes, fix it" in response to the security concern about unauthenticated access. I was in the middle of editing relay/mod.rs when the summary was requested..
Please continue the conversation from where we left it off without asking the user any further questions. Continue with the last task that you were asked to work on.


The user wants me to continue where I left off. Looking at the summary:

1. I was in the middle of securing the WebSocket handlers in relay/mod.rs
2. I had already completed securing the REST API endpoints with AuthenticatedUser
3. The todo list shows:
   - [COMPLETED] Add AuthenticatedUser to REST API endpoints
   - [IN PROGRESS] Add JWT auth to viewer WebSocket
   - [PENDING] Add agent authentication (API key or support code required)
   - [PENDING] Rebuild and deploy server

I need to continue adding authentication to the WebSocket handlers. Let me read the current state of relay/mod.rs to see what I've already done and continue from there.

The summary mentions I added StatusCode import but haven't implemented the actual auth checks yet.

Let me check the current state of relay/mod.rs and continue implementing:
1. JWT validation for viewer WebSocket (require token in query params)
2. Authentication for agent WebSocket (require valid support_code OR API key)


Continuing to secure the WebSocket handlers. Let me check the current state of relay/mod.rs and implement JWT validation for viewers and authentication for agents.


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

<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,
    http::StatusCode,
};
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::db::{self, Database};
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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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();
    let db = state.db.clone();

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    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);

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

    // 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;

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_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;

    // Database: end session and mark machine offline
    if let Some(ref db) = db {
        // End the session record
        let _ = db::sessions::end_session(db.pool(), session_id, "ended").await;

        // Mark machine as offline
        let _ = db::machines::mark_machine_offline(db.pool(), &agent_id).await;

        // Log session ended event
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::SESSION_ENDED,
            None, None, None, None,
        ).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;

            // Database: mark code as completed
            if let Some(ref db) = db {
                let _ = db::support_codes::mark_code_completed(db.pool(), 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,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

    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();
    let viewer_name_cleanup = viewer_name.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;

    // Database: log viewer left event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_LEFT,
            Some(&viewer_id_cleanup),
            Some(&viewer_name_cleanup),
            None, None,
        ).await;
    }

    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}


     1→//! Authentication module
     2→//!
     3→//! Handles JWT validation for dashboard users and API key
     4→//! validation for agents.
     5→
     6→pub mod jwt;
     7→pub mod password;
     8→
     9→pub use jwt::{Claims, JwtConfig};
    10→pub use password::{hash_password, verify_password, generate_random_password};
    11→
    12→use axum::{
    13→    extract::FromRequestParts,
    14→    http::{request::Parts, StatusCode},
    15→};
    16→use std::sync::Arc;
    17→
    18→/// Authenticated user from JWT
    19→#[derive(Debug, Clone)]
    20→pub struct AuthenticatedUser {
    21→    pub user_id: String,
    22→    pub username: String,
    23→    pub role: String,
    24→    pub permissions: Vec<String>,
    25→}
    26→
    27→impl AuthenticatedUser {
    28→    /// Check if user has a specific permission
    29→    pub fn has_permission(&self, permission: &str) -> bool {
    30→        if self.role == "admin" {
    31→            return true;
    32→        }
    33→        self.permissions.contains(&permission.to_string())
    34→    }
    35→
    36→    /// Check if user is admin
    37→    pub fn is_admin(&self) -> bool {
    38→        self.role == "admin"
    39→    }
    40→}
    41→
    42→impl From<Claims> for AuthenticatedUser {
    43→    fn from(claims: Claims) -> Self {
    44→        Self {
    45→            user_id: claims.sub,
    46→            username: claims.username,
    47→            role: claims.role,
    48→            permissions: claims.permissions,
    49→        }
    50→    }
    51→}
    52→
    53→/// Authenticated agent from API key
    54→#[derive(Debug, Clone)]
    55→pub struct AuthenticatedAgent {
    56→    pub agent_id: String,
    57→    pub org_id: String,
    58→}
    59→
    60→/// JWT configuration stored in app state
    61→#[derive(Clone)]
    62→pub struct AuthState {
    63→    pub jwt_config: Arc<JwtConfig>,
    64→}
    65→
    66→impl AuthState {
    67→    pub fn new(jwt_secret: String, expiry_hours: i64) -> Self {
    68→        Self {
    69→            jwt_config: Arc::new(JwtConfig::new(jwt_secret, expiry_hours)),
    70→        }
    71→    }
    72→}
    73→
    74→/// Extract authenticated user from request
    75→#[axum::async_trait]
    76→impl<S> FromRequestParts<S> for AuthenticatedUser
    77→where
    78→    S: Send + Sync,
    79→{
    80→    type Rejection = (StatusCode, &'static str);
    81→
    82→    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
    83→        // Get Authorization header
    84→        let auth_header = parts
    85→            .headers
    86→            .get("Authorization")
    87→            .and_then(|v| v.to_str().ok())
    88→            .ok_or((StatusCode::UNAUTHORIZED, "Missing Authorization header"))?;
    89→
    90→        // Extract Bearer token
    91→        let token = auth_header
    92→            .strip_prefix("Bearer ")
    93→            .ok_or((StatusCode::UNAUTHORIZED, "Invalid Authorization format"))?;
    94→
    95→        // Get JWT config from extensions (set by middleware)
    96→        let jwt_config = parts
    97→            .extensions
    98→            .get::<Arc<JwtConfig>>()
    99→            .ok_or((StatusCode::INTERNAL_SERVER_ERROR, "Auth not configured"))?;
   100→
   101→        // Validate token
   102→        let claims = jwt_config
   103→            .validate_token(token)
   104→            .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid or expired token"))?;
   105→
   106→        Ok(AuthenticatedUser::from(claims))
   107→    }
   108→}
   109→
   110→/// Optional authenticated user (doesn't reject if not authenticated)
   111→#[derive(Debug, Clone)]
   112→pub struct OptionalUser(pub Option<AuthenticatedUser>);
   113→
   114→#[axum::async_trait]
   115→impl<S> FromRequestParts<S> for OptionalUser
   116→where
   117→    S: Send + Sync,
   118→{
   119→    type Rejection = (StatusCode, &'static str);
   120→
   121→    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
   122→        match AuthenticatedUser::from_request_parts(parts, state).await {
   123→            Ok(user) => Ok(OptionalUser(Some(user))),
   124→            Err(_) => Ok(OptionalUser(None)),
   125→        }
   126→    }
   127→}
   128→
   129→/// Require admin role
   130→#[derive(Debug, Clone)]
   131→pub struct AdminUser(pub AuthenticatedUser);
   132→
   133→#[axum::async_trait]
   134→impl<S> FromRequestParts<S> for AdminUser
   135→where
   136→    S: Send + Sync,
   137→{
   138→    type Rejection = (StatusCode, &'static str);
   139→
   140→    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
   141→        let user = AuthenticatedUser::from_request_parts(parts, state).await?;
   142→        if user.is_admin() {
   143→            Ok(AdminUser(user))
   144→        } else {
   145→            Err((StatusCode::FORBIDDEN, "Admin access required"))
   146→        }
   147→    }
   148→}
   149→
   150→/// Validate an agent API key (placeholder for MVP)
   151→pub fn validate_agent_key(_api_key: &str) -> Option<AuthenticatedAgent> {
   152→    // TODO: Implement actual API key validation against database
   153→    // For now, accept any key for agent connections
   154→    Some(AuthenticatedAgent {
   155→        agent_id: "mvp-agent".to_string(),
   156→        org_id: "mvp-org".to_string(),
   157→    })
   158→}
   159→

<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.

pub mod jwt;
pub mod password;

pub use jwt::{Claims, JwtConfig};
pub use password::{hash_password, verify_password, generate_random_password};

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

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

impl AuthenticatedUser {
    /// Check if user has a specific permission
    pub fn has_permission(&self, permission: &str) -> bool {
        if self.role == "admin" {
            return true;
        }
        self.permissions.contains(&permission.to_string())
    }

    /// Check if user is admin
    pub fn is_admin(&self) -> bool {
        self.role == "admin"
    }
}

impl From<Claims> for AuthenticatedUser {
    fn from(claims: Claims) -> Self {
        Self {
            user_id: claims.sub,
            username: claims.username,
            role: claims.role,
            permissions: claims.permissions,
        }
    }
}

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

/// JWT configuration stored in app state
#[derive(Clone)]
pub struct AuthState {
    pub jwt_config: Arc<JwtConfig>,
}

impl AuthState {
    pub fn new(jwt_secret: String, expiry_hours: i64) -> Self {
        Self {
            jwt_config: Arc::new(JwtConfig::new(jwt_secret, expiry_hours)),
        }
    }
}

/// Extract authenticated user from request
#[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> {
        // Get Authorization header
        let auth_header = parts
            .headers
            .get("Authorization")
            .and_then(|v| v.to_str().ok())
            .ok_or((StatusCode::UNAUTHORIZED, "Missing Authorization header"))?;

        // Extract Bearer token
        let token = auth_header
            .strip_prefix("Bearer ")
            .ok_or((StatusCode::UNAUTHORIZED, "Invalid Authorization format"))?;

        // Get JWT config from extensions (set by middleware)
        let jwt_config = parts
            .extensions
            .get::<Arc<JwtConfig>>()
            .ok_or((StatusCode::INTERNAL_SERVER_ERROR, "Auth not configured"))?;

        // Validate token
        let claims = jwt_config
            .validate_token(token)
            .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid or expired token"))?;

        Ok(AuthenticatedUser::from(claims))
    }
}

/// Optional authenticated user (doesn't reject if not authenticated)
#[derive(Debug, Clone)]
pub struct OptionalUser(pub Option<AuthenticatedUser>);

#[axum::async_trait]
impl<S> FromRequestParts<S> for OptionalUser
where
    S: Send + Sync,
{
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        match AuthenticatedUser::from_request_parts(parts, state).await {
            Ok(user) => Ok(OptionalUser(Some(user))),
            Err(_) => Ok(OptionalUser(None)),
        }
    }
}

/// Require admin role
#[derive(Debug, Clone)]
pub struct AdminUser(pub AuthenticatedUser);

#[axum::async_trait]
impl<S> FromRequestParts<S> for AdminUser
where
    S: Send + Sync,
{
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let user = AuthenticatedUser::from_request_parts(parts, state).await?;
        if user.is_admin() {
            Ok(AdminUser(user))
        } else {
            Err((StatusCode::FORBIDDEN, "Admin access required"))
        }
    }
}

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


Looking at the code, I need to add authentication to the WebSocket handlers:

1. **Viewer WebSocket** - needs JWT token validation (passed as query param since WebSocket can't use headers initially)
2. **Agent WebSocket** - needs either:
   - Valid support code (for ad-hoc support)
   - Valid API key (for persistent managed agents)

For viewers, I'll add a `token` query param to ViewerParams and validate it before upgrading.

For agents, I need to decide:
- Support code sessions are still allowed (they get a code from dashboard)
- Persistent agents (no support code) should require some form of authentication

Looking at the current logic:
- `support_code.is_none()` makes agent "persistent"
- Currently ANY agent can connect without a code and become persistent

The fix: persistent agents MUST have an API key or be connecting with a valid support code.

Let me update the relay/mod.rs:

1. Add token to ViewerParams
2. Validate JWT for viewers before upgrade
3. Require support_code OR api_key for agents

I need to import the auth module and JwtConfig. Let me also check how to access the JwtConfig from AppState.


Now I understand the security gap. Let me implement proper authentication for both WebSocket handlers:

1. **Viewers**: Require JWT token in query params (validated before upgrade)
2. **Agents**: Require either a valid support code OR API key


Edit

#[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,
}

#[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>,
    /// API key for persistent (managed) agents
    #[serde(default)]
    api_key: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
    /// JWT token for authentication (required)
    #[serde(default)]
    token: Option<String>,
}


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\relay\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    21→use crate::session::SessionManager;
    22→use crate::db::{self, Database};
    23→use crate::AppState;
    24→
    25→#[derive(Debug, Deserialize)]
    26→pub struct AgentParams {
    27→    agent_id: String,
    28→    #[serde(default)]
    29→    agent_name: Option<String>,
    30→    #[serde(default)]
    31→    support_code: Option<String>,
    32→    #[serde(default)]
    33→    hostname: Option<String>,
    34→    /// API key for persistent (managed) agents
    35→    #[serde(default)]
    36→    api_key: Option<String>,
    37→}
    38→
    39→#[derive(Debug, Deserialize)]
    40→pub struct ViewerParams {
    41→    session_id: String,
    42→    #[serde(default = "default_viewer_name")]
    43→    viewer_name: String,
    44→    /// JWT token for authentication (required)
    45→    #[serde(default)]
    46→    token: Option<String>,
    47→}
    48→
    49→fn default_viewer_name() -> String {
    50→    "Technician".to_string()
    51→}

#[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>,
    /// API key for persistent (managed) agents
    #[serde(default)]
    api_key: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
    /// JWT token for authentication (required)
    #[serde(default)]
    token: Option<String>,
}

//! 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,
    http::StatusCode,
};
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::db::{self, Database};
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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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();
    let db = state.db.clone();

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    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);

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

    // 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;

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_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;

    // Database: end session and mark machine offline
    if let Some(ref db) = db {
        // End the session record
        let _ = db::sessions::end_session(db.pool(), session_id, "ended").await;

        // Mark machine as offline
        let _ = db::machines::mark_machine_offline(db.pool(), &agent_id).await;

        // Log session ended event
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::SESSION_ENDED,
            None, None, None, None,
        ).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;

            // Database: mark code as completed
            if let Some(ref db) = db {
                let _ = db::support_codes::mark_code_completed(db.pool(), 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,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

    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();
    let viewer_name_cleanup = viewer_name.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;

    // Database: log viewer left event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_LEFT,
            Some(&viewer_id_cleanup),
            Some(&viewer_name_cleanup),
            None, None,
        ).await;
    }

    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}

+    /// API key for persistent (managed) agents

+    /// JWT token for authentication (required)


Now I need to update the viewer_ws_handler to validate the JWT token before accepting the connection, and agent_ws_handler to require either support_code OR api_key.


/// 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();
    let db = state.db.clone();

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

/// WebSocket handler for viewer connections
pub async fn viewer_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<ViewerParams>,
) -> Result<impl IntoResponse, StatusCode> {
    // Require JWT token for viewers
    let token = params.token.ok_or_else(|| {
        warn!("Viewer connection rejected: missing token");
        StatusCode::UNAUTHORIZED
    })?;

    // Validate the token
    let claims = state.jwt_config.validate_token(&token).map_err(|e| {
        warn!("Viewer connection rejected: invalid token: {}", e);
        StatusCode::UNAUTHORIZED
    })?;

    info!("Viewer {} authenticated via JWT", claims.username);

    let session_id = params.session_id;
    let viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();
    let db = state.db.clone();

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