The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   132→        }
   133→        Some(Commands::Launch { url }) => {
   134→            run_launch(&url)
   135→        }
   136→        None => {
   137→            // Legacy mode: if a support code was provided, run as agent
   138→            if let Some(code) = cli.support_code {
   139→                run_agent_mode(Some(code))
   140→            } else {
   141→                // No args: check if protocol handler is installed
   142→                // If not, run install mode (user likely downloaded from web)
   143→                if !install::is_protocol_handler_registered() {
   144→                    info!("Protocol handler not registered, running installer");
   145→                    run_install(false)
   146→                } else {
   147→                    // Protocol handler exists, run as agent
   148→                    run_agent_mode(None)
   149→                }
   150→            }
   151→        }
   152→    }
   153→}
   154→
   155→/// Run in agent mode (receive remote connections)

        None => {
            // Legacy mode: if a support code was provided, run as agent
            if let Some(code) = cli.support_code {
                run_agent_mode(Some(code))
            } else {
                // No args: check if protocol handler is installed
                // If not, run install mode (user likely downloaded from web)
                if !install::is_protocol_handler_registered() {
                    info!("Protocol handler not registered, running installer");
                    run_install(false)
                } else {
                    // Protocol handler exists, run as agent
                    run_agent_mode(None)
                }
            }
        }

//! GuruConnect - Remote Desktop Agent and Viewer
//!
//! Single binary for both agent (receiving connections) and viewer (initiating connections).
//!
//! Usage:
//!   guruconnect agent              - Run as background agent
//!   guruconnect view <session_id>  - View a remote session
//!   guruconnect install            - Install and register protocol handler
//!   guruconnect launch <url>       - Handle guruconnect:// URL
//!   guruconnect [support_code]     - Legacy: run agent with support code

// Hide console window by default on Windows (release builds)
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

mod capture;
mod chat;
mod config;
mod encoder;
mod input;
mod install;
mod sas_client;
mod session;
mod startup;
mod transport;
mod tray;
mod viewer;

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

use anyhow::Result;
use clap::{Parser, Subcommand};
use tracing::{info, error, warn, Level};
use tracing_subscriber::FmtSubscriber;

#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION, MB_ICONERROR};
#[cfg(windows)]
use windows::core::PCWSTR;
#[cfg(windows)]
use windows::Win32::System::Console::{AllocConsole, GetConsoleWindow};
#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW};

/// GuruConnect Remote Desktop
#[derive(Parser)]
#[command(name = "guruconnect")]
#[command(version, about = "Remote desktop agent and viewer")]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,

    /// Support code for legacy mode (runs agent with code)
    #[arg(value_name = "SUPPORT_CODE")]
    support_code: Option<String>,

    /// Enable verbose logging
    #[arg(short, long, global = true)]
    verbose: bool,
}

#[derive(Subcommand)]
enum Commands {
    /// Run as background agent (receive remote connections)
    Agent {
        /// Support code for one-time session
        #[arg(short, long)]
        code: Option<String>,
    },

    /// View a remote session (connect to an agent)
    View {
        /// Session ID to connect to
        session_id: String,

        /// Server URL
        #[arg(short, long, default_value = "wss://connect.azcomputerguru.com/ws/viewer")]
        server: String,

        /// API key for authentication
        #[arg(short, long, default_value = "")]
        api_key: String,
    },

    /// Install GuruConnect and register protocol handler
    Install {
        /// Skip UAC elevation, install for current user only
        #[arg(long)]
        user_only: bool,

        /// Called internally when running elevated
        #[arg(long, hide = true)]
        elevated: bool,
    },

    /// Uninstall GuruConnect
    Uninstall,

    /// Handle a guruconnect:// protocol URL
    Launch {
        /// The guruconnect:// URL to handle
        url: String,
    },
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    // Initialize logging
    let level = if cli.verbose { Level::DEBUG } else { Level::INFO };
    FmtSubscriber::builder()
        .with_max_level(level)
        .with_target(true)
        .with_thread_ids(true)
        .init();

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

    match cli.command {
        Some(Commands::Agent { code }) => {
            run_agent_mode(code)
        }
        Some(Commands::View { session_id, server, api_key }) => {
            run_viewer_mode(&server, &session_id, &api_key)
        }
        Some(Commands::Install { user_only, elevated }) => {
            run_install(user_only || elevated)
        }
        Some(Commands::Uninstall) => {
            run_uninstall()
        }
        Some(Commands::Launch { url }) => {
            run_launch(&url)
        }
        None => {
            // Legacy mode: if a support code was provided, run as agent
            if let Some(code) = cli.support_code {
                run_agent_mode(Some(code))
            } else {
                // No args: check if we should auto-detect mode
                // For now, default to agent mode
                run_agent_mode(None)
            }
        }
    }
}

/// Run in agent mode (receive remote connections)
fn run_agent_mode(support_code: Option<String>) -> Result<()> {
    info!("Running in agent mode");

    // Check elevation status
    if install::is_elevated() {
        info!("Running with elevated (administrator) privileges");
    } else {
        info!("Running with standard user privileges");
    }

    // Also check for support code in filename (legacy compatibility)
    let code = support_code.or_else(extract_code_from_filename);
    if let Some(ref c) = code {
        info!("Support code: {}", c);
    }

    // Load configuration
    let mut config = config::Config::load()?;
    config.support_code = code;
    info!("Server: {}", config.server_url);

    // Run the agent
    let rt = tokio::runtime::Runtime::new()?;
    rt.block_on(run_agent(config))
}

/// Run in viewer mode (connect to remote session)
fn run_viewer_mode(server: &str, session_id: &str, api_key: &str) -> Result<()> {
    info!("Running in viewer mode");
    info!("Connecting to session: {}", session_id);

    let rt = tokio::runtime::Runtime::new()?;
    rt.block_on(viewer::run(server, session_id, api_key))
}

/// Handle guruconnect:// URL launch
fn run_launch(url: &str) -> Result<()> {
    info!("Handling protocol URL: {}", url);

    match install::parse_protocol_url(url) {
        Ok((server, session_id, token)) => {
            let api_key = token.unwrap_or_default();
            run_viewer_mode(&server, &session_id, &api_key)
        }
        Err(e) => {
            error!("Failed to parse URL: {}", e);
            show_error_box("GuruConnect", &format!("Invalid URL: {}", e));
            Err(e)
        }
    }
}

/// Install GuruConnect
fn run_install(force_user_install: bool) -> Result<()> {
    info!("Installing GuruConnect...");

    match install::install(force_user_install) {
        Ok(()) => {
            show_message_box("GuruConnect", "Installation complete!\n\nYou can now use guruconnect:// links.");
            Ok(())
        }
        Err(e) => {
            error!("Installation failed: {}", e);
            show_error_box("GuruConnect", &format!("Installation failed: {}", e));
            Err(e)
        }
    }
}

/// Uninstall GuruConnect
fn run_uninstall() -> Result<()> {
    info!("Uninstalling GuruConnect...");

    // Remove from startup
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }

    // TODO: Remove registry keys for protocol handler
    // TODO: Remove install directory

    show_message_box("GuruConnect", "Uninstall complete.");
    Ok(())
}

/// Extract a 6-digit support code from the executable's filename
fn extract_code_from_filename() -> Option<String> {
    let exe_path = std::env::current_exe().ok()?;
    let filename = exe_path.file_stem()?.to_str()?;

    // Look for a 6-digit number in the filename
    for part in filename.split(|c| c == '-' || c == '_' || c == '.') {
        let trimmed = part.trim();
        if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) {
            return Some(trimmed.to_string());
        }
    }

    // Check if the last 6 characters are digits
    if filename.len() >= 6 {
        let last_six = &filename[filename.len() - 6..];
        if last_six.chars().all(|c| c.is_ascii_digit()) {
            return Some(last_six.to_string());
        }
    }

    None
}

/// Show a message box (Windows only)
#[cfg(windows)]
fn show_message_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONINFORMATION,
        );
    }
}

#[cfg(not(windows))]
fn show_message_box(_title: &str, message: &str) {
    println!("{}", message);
}

/// Show an error message box (Windows only)
#[cfg(windows)]
fn show_error_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONERROR,
        );
    }
}

#[cfg(not(windows))]
fn show_error_box(_title: &str, message: &str) {
    eprintln!("ERROR: {}", message);
}

/// Show debug console window (Windows only)
#[cfg(windows)]
#[allow(dead_code)]
fn show_debug_console() {
    unsafe {
        let hwnd = GetConsoleWindow();
        if hwnd.0 == std::ptr::null_mut() {
            let _ = AllocConsole();
        } else {
            let _ = ShowWindow(hwnd, SW_SHOW);
        }
    }
}

#[cfg(not(windows))]
#[allow(dead_code)]
fn show_debug_console() {}

/// Clean up before exiting
fn cleanup_on_exit() {
    info!("Cleaning up before exit");
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }
}

/// Run the agent main loop
async fn run_agent(config: config::Config) -> Result<()> {
    let elevated = install::is_elevated();
    let mut session = session::SessionManager::new(config.clone(), elevated);
    let is_support_session = config.support_code.is_some();
    let hostname = config.hostname();

    // Add to startup
    if let Err(e) = startup::add_to_startup() {
        warn!("Failed to add to startup: {}", e);
    }

    // Create tray icon
    let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref(), is_support_session) {
        Ok(t) => {
            info!("Tray icon created");
            Some(t)
        }
        Err(e) => {
            warn!("Failed to create tray icon: {}", e);
            None
        }
    };

    // Create chat controller
    let chat_ctrl = chat::ChatController::new();

    // Connect to server and run main loop
    loop {
        info!("Connecting to server...");

        if is_support_session {
            if let Some(ref t) = tray {
                if t.exit_requested() {
                    info!("Exit requested by user");
                    cleanup_on_exit();
                    return Ok(());
                }
            }
        }

        match session.connect().await {
            Ok(_) => {
                info!("Connected to server");

                if let Some(ref t) = tray {
                    t.update_status("Status: Connected");
                }

                if let Err(e) = session.run_with_tray(tray.as_ref(), chat_ctrl.as_ref()).await {
                    let error_msg = e.to_string();

                    if error_msg.contains("USER_EXIT") {
                        info!("Session ended by user");
                        cleanup_on_exit();
                        return Ok(());
                    }

                    if error_msg.contains("SESSION_CANCELLED") {
                        info!("Session was cancelled by technician");
                        cleanup_on_exit();
                        show_message_box("Support Session Ended", "The support session was cancelled.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_DISCONNECT") {
                        info!("Session disconnected by administrator - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box("Remote Session Ended", "The session was ended by the administrator.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_UNINSTALL") {
                        info!("Uninstall command received from server - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box("GuruConnect Removed", "This computer has been removed from remote management.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_RESTART") {
                        info!("Restart command received - will reconnect");
                        // Don't exit, just let the loop continue to reconnect
                    } else {
                        error!("Session error: {}", e);
                    }
                }
            }
            Err(e) => {
                let error_msg = e.to_string();

                if error_msg.contains("cancelled") {
                    info!("Support code was cancelled");
                    cleanup_on_exit();
                    show_message_box("Support Session Cancelled", "This support session has been cancelled.");
                    return Ok(());
                }

                error!("Connection failed: {}", e);
            }
        }

        if is_support_session {
            info!("Support session ended, not reconnecting");
            cleanup_on_exit();
            return Ok(());
        }

        info!("Reconnecting in 5 seconds...");
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
    }
}

-                // No args: check if we should auto-detect mode

-                // For now, default to agent mode

+                // No args: check if protocol handler is installed

+                // If not, run install mode (user likely downloaded from web)

+                if !install::is_protocol_handler_registered() {

+                    info!("Protocol handler not registered, running installer");

+                    // Protocol handler exists, run as agent

+                    run_agent_mode(None)


Now I need to add the `is_protocol_handler_registered()` function to install.rs.


Now add the function to check if protocol handler is registered:


Edit

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {

/// Check if the guruconnect:// protocol handler is registered
#[cfg(windows)]
pub fn is_protocol_handler_registered() -> bool {
    use windows::Win32::System::Registry::{
        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,
    };

    unsafe {
        // Check system-wide registration (HKCR\guruconnect)
        let mut key = HKEY::default();
        let key_path = to_wide("guruconnect");
        if RegOpenKeyExW(
            HKEY_CLASSES_ROOT,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }

        // Check user-level registration (HKCU\Software\Classes\guruconnect)
        let key_path = to_wide(r"Software\Classes\guruconnect");
        if RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }
    }

    false
}

#[cfg(not(windows))]
pub fn is_protocol_handler_registered() -> bool {
    // On non-Windows, assume not registered (or check ~/.local/share/applications)
    false
}

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {


The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\install.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   300→
   301→    Ok(())
   302→}
   303→
   304→/// Check if the guruconnect:// protocol handler is registered
   305→#[cfg(windows)]
   306→pub fn is_protocol_handler_registered() -> bool {
   307→    use windows::Win32::System::Registry::{
   308→        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,
   309→    };
   310→
   311→    unsafe {
   312→        // Check system-wide registration (HKCR\guruconnect)
   313→        let mut key = HKEY::default();
   314→        let key_path = to_wide("guruconnect");
   315→        if RegOpenKeyExW(
   316→            HKEY_CLASSES_ROOT,
   317→            PCWSTR(key_path.as_ptr()),
   318→            0,
   319→            KEY_READ,
   320→            &mut key,
   321→        ).is_ok() {
   322→            let _ = RegCloseKey(key);
   323→            return true;
   324→        }
   325→
   326→        // Check user-level registration (HKCU\Software\Classes\guruconnect)
   327→        let key_path = to_wide(r"Software\Classes\guruconnect");
   328→        if RegOpenKeyExW(
   329→            HKEY_CURRENT_USER,
   330→            PCWSTR(key_path.as_ptr()),
   331→            0,
   332→            KEY_READ,
   333→            &mut key,
   334→        ).is_ok() {
   335→            let _ = RegCloseKey(key);
   336→            return true;
   337→        }
   338→    }
   339→
   340→    false
   341→}
   342→
   343→#[cfg(not(windows))]
   344→pub fn is_protocol_handler_registered() -> bool {
   345→    // On non-Windows, assume not registered (or check ~/.local/share/applications)
   346→    false
   347→}
   348→
   349→/// Parse a guruconnect:// URL and extract session parameters
   350→pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
   351→    // Expected formats:
   352→    // guruconnect://view/SESSION_ID
   353→    // guruconnect://view/SESSION_ID?token=API_KEY
   354→    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY

/// Check if the guruconnect:// protocol handler is registered
#[cfg(windows)]
pub fn is_protocol_handler_registered() -> bool {
    use windows::Win32::System::Registry::{
        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,
    };

    unsafe {
        // Check system-wide registration (HKCR\guruconnect)
        let mut key = HKEY::default();
        let key_path = to_wide("guruconnect");
        if RegOpenKeyExW(
            HKEY_CLASSES_ROOT,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }

        // Check user-level registration (HKCU\Software\Classes\guruconnect)
        let key_path = to_wide(r"Software\Classes\guruconnect");
        if RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }
    }

    false
}

#[cfg(not(windows))]
pub fn is_protocol_handler_registered() -> bool {
    // On non-Windows, assume not registered (or check ~/.local/share/applications)
    false
}

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {

//! Installation and protocol handler registration
//!
//! Handles:
//! - Self-installation to Program Files (with UAC) or LocalAppData (fallback)
//! - Protocol handler registration (guruconnect://)
//! - UAC elevation with graceful fallback

use anyhow::{anyhow, Result};
use tracing::{info, warn, error};

#[cfg(windows)]
use windows::{
    core::PCWSTR,
    Win32::Foundation::HANDLE,
    Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY},
    Win32::System::Threading::{GetCurrentProcess, OpenProcessToken},
    Win32::System::Registry::{
        RegCreateKeyExW, RegSetValueExW, RegCloseKey, HKEY, HKEY_CLASSES_ROOT,
        HKEY_CURRENT_USER, KEY_WRITE, REG_SZ, REG_OPTION_NON_VOLATILE,
    },
    Win32::UI::Shell::ShellExecuteW,
    Win32::UI::WindowsAndMessaging::SW_SHOWNORMAL,
};

#[cfg(windows)]
use std::ffi::OsStr;
#[cfg(windows)]
use std::os::windows::ffi::OsStrExt;

/// Install locations
pub const SYSTEM_INSTALL_PATH: &str = r"C:\Program Files\GuruConnect";
pub const USER_INSTALL_PATH: &str = r"GuruConnect"; // Relative to %LOCALAPPDATA%

/// Check if running with elevated privileges
#[cfg(windows)]
pub fn is_elevated() -> bool {
    unsafe {
        let mut token_handle = HANDLE::default();
        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
            return false;
        }

        let mut elevation = TOKEN_ELEVATION::default();
        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;

        let result = GetTokenInformation(
            token_handle,
            TokenElevation,
            Some(&mut elevation as *mut _ as *mut _),
            size,
            &mut size,
        );

        let _ = windows::Win32::Foundation::CloseHandle(token_handle);

        result.is_ok() && elevation.TokenIsElevated != 0
    }
}

#[cfg(not(windows))]
pub fn is_elevated() -> bool {
    unsafe { libc::geteuid() == 0 }
}

/// Get the install path based on elevation status
pub fn get_install_path(elevated: bool) -> std::path::PathBuf {
    if elevated {
        std::path::PathBuf::from(SYSTEM_INSTALL_PATH)
    } else {
        let local_app_data = std::env::var("LOCALAPPDATA")
            .unwrap_or_else(|_| {
                let home = std::env::var("USERPROFILE").unwrap_or_else(|_| ".".to_string());
                format!(r"{}\AppData\Local", home)
            });
        std::path::PathBuf::from(local_app_data).join(USER_INSTALL_PATH)
    }
}

/// Get the executable path
pub fn get_exe_path(install_path: &std::path::Path) -> std::path::PathBuf {
    install_path.join("guruconnect.exe")
}

/// Attempt to elevate and re-run with install command
#[cfg(windows)]
pub fn try_elevate_and_install() -> Result<bool> {
    let exe_path = std::env::current_exe()?;
    let exe_path_wide: Vec<u16> = OsStr::new(exe_path.as_os_str())
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let verb: Vec<u16> = OsStr::new("runas")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let params: Vec<u16> = OsStr::new("install --elevated")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let result = ShellExecuteW(
            None,
            PCWSTR(verb.as_ptr()),
            PCWSTR(exe_path_wide.as_ptr()),
            PCWSTR(params.as_ptr()),
            PCWSTR::null(),
            SW_SHOWNORMAL,
        );

        // ShellExecuteW returns > 32 on success
        if result.0 as usize > 32 {
            info!("UAC elevation requested");
            Ok(true)
        } else {
            warn!("UAC elevation denied or failed");
            Ok(false)
        }
    }
}

#[cfg(not(windows))]
pub fn try_elevate_and_install() -> Result<bool> {
    Ok(false)
}

/// Register the guruconnect:// protocol handler
#[cfg(windows)]
pub fn register_protocol_handler(elevated: bool) -> Result<()> {
    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);
    let exe_path_str = exe_path.to_string_lossy();

    // Command to execute: "C:\...\guruconnect.exe" "launch" "%1"
    let command = format!("\"{}\" launch \"%1\"", exe_path_str);

    // Choose registry root based on elevation
    let root_key = if elevated {
        HKEY_CLASSES_ROOT
    } else {
        // User-level registration under Software\Classes
        HKEY_CURRENT_USER
    };

    let base_path = if elevated {
        "guruconnect"
    } else {
        r"Software\Classes\guruconnect"
    };

    unsafe {
        // Create guruconnect key
        let mut protocol_key = HKEY::default();
        let key_path = to_wide(base_path);
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut protocol_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create protocol key: {:?}", result));
        }

        // Set default value (protocol description)
        let description = to_wide("GuruConnect Protocol");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&description)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set protocol description: {:?}", result));
        }

        // Set URL Protocol (empty string indicates this is a protocol handler)
        let url_protocol = to_wide("URL Protocol");
        let empty = to_wide("");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR(url_protocol.as_ptr()),
            0,
            REG_SZ,
            Some(&description_to_bytes(&empty)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set URL Protocol: {:?}", result));
        }

        let _ = RegCloseKey(protocol_key);

        // Create shell\open\command key
        let command_path = if elevated {
            r"guruconnect\shell\open\command"
        } else {
            r"Software\Classes\guruconnect\shell\open\command"
        };
        let command_key_path = to_wide(command_path);
        let mut command_key = HKEY::default();
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(command_key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut command_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create command key: {:?}", result));
        }

        // Set the command
        let command_wide = to_wide(&command);
        let result = RegSetValueExW(
            command_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&command_wide)),
        );
        if result.is_err() {
            let _ = RegCloseKey(command_key);
            return Err(anyhow!("Failed to set command: {:?}", result));
        }

        let _ = RegCloseKey(command_key);
    }

    info!("Protocol handler registered: guruconnect://");
    Ok(())
}

#[cfg(not(windows))]
pub fn register_protocol_handler(_elevated: bool) -> Result<()> {
    warn!("Protocol handler registration not supported on this platform");
    Ok(())
}

/// Install the application
pub fn install(force_user_install: bool) -> Result<()> {
    let elevated = is_elevated();

    // If not elevated and not forcing user install, try to elevate
    if !elevated && !force_user_install {
        info!("Attempting UAC elevation for system-wide install...");
        match try_elevate_and_install() {
            Ok(true) => {
                // Elevation was requested, exit this instance
                // The elevated instance will continue the install
                info!("Elevated process started, exiting current instance");
                std::process::exit(0);
            }
            Ok(false) => {
                info!("UAC denied, falling back to user install");
            }
            Err(e) => {
                warn!("Elevation failed: {}, falling back to user install", e);
            }
        }
    }

    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);

    info!("Installing to: {}", install_path.display());

    // Create install directory
    std::fs::create_dir_all(&install_path)?;

    // Copy ourselves to install location
    let current_exe = std::env::current_exe()?;
    if current_exe != exe_path {
        std::fs::copy(&current_exe, &exe_path)?;
        info!("Copied executable to: {}", exe_path.display());
    }

    // Register protocol handler
    register_protocol_handler(elevated)?;

    info!("Installation complete!");
    if elevated {
        info!("Installed system-wide to: {}", install_path.display());
    } else {
        info!("Installed for current user to: {}", install_path.display());
    }

    Ok(())
}

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY

    let url = url::Url::parse(url)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    let path = url.path().trim_start_matches('/');
    let parts: Vec<&str> = path.split('/').collect();

    if parts.is_empty() {
        return Err(anyhow!("Missing action in URL"));
    }

    let action = parts[0];
    let session_id = parts.get(1).map(|s| s.to_string())
        .ok_or_else(|| anyhow!("Missing session ID"))?;

    // Extract query parameters
    let mut server = None;
    let mut token = None;

    for (key, value) in url.query_pairs() {
        match key.as_ref() {
            "server" => server = Some(value.to_string()),
            "token" | "api_key" => token = Some(value.to_string()),
            _ => {}
        }
    }

    // Default server if not specified
    let server = server.unwrap_or_else(|| "wss://connect.azcomputerguru.com/ws/viewer".to_string());

    match action {
        "view" | "connect" => Ok((server, session_id, token)),
        _ => Err(anyhow!("Unknown action: {}", action)),
    }
}

// Helper functions for Windows registry operations
#[cfg(windows)]
fn to_wide(s: &str) -> Vec<u16> {
    OsStr::new(s)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect()
}

#[cfg(windows)]
fn description_to_bytes(wide: &[u16]) -> Vec<u8> {
    wide.iter()
        .flat_map(|w| w.to_le_bytes())
        .collect()
}

+/// Check if the guruconnect:// protocol handler is registered

+pub fn is_protocol_handler_registered() -> bool {

+    use windows::Win32::System::Registry::{

+        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,

+        // Check system-wide registration (HKCR\guruconnect)

+        let key_path = to_wide("guruconnect");

+        // Check user-level registration (HKCU\Software\Classes\guruconnect)

+        let key_path = to_wide(r"Software\Classes\guruconnect");

+    // On non-Windows, assume not registered (or check ~/.local/share/applications)


Now I need to rebuild the agent and upload it to the server. Let me rebuild.


Now rebuild and upload:


Bash

cd /c/Users/MikeSwanson/Claude/guru-connect && cargo build -p guruconnect --release 2>&1 | tail -20

Rebuild Windows agent


    |             ^^^^^^^^^^^^^^^^^^^^
    |
help: use `let _ = ...` to ignore the resulting value
    |
114 |             let _ = DeleteObject(bitmap);
    |             +++++++

warning: unused `BOOL` that must be used
   --> agent\src\capture\gdi.rs:115:13
    |
115 |             DeleteDC(mem_dc);
    |             ^^^^^^^^^^^^^^^^
    |
help: use `let _ = ...` to ignore the resulting value
    |
115 |             let _ = DeleteDC(mem_dc);
    |             +++++++

warning: `guruconnect` (bin "guruconnect") generated 59 warnings (run `cargo fix --bin "guruconnect" -p guruconnect` to apply 20 suggestions)
    Finished `release` profile [optimized] target(s) in 35.72s


Build succeeded. Now upload to the server.


Upload updated binary


Done. Let me also commit the changes to git.


cd /c/Users/MikeSwanson/Claude/guru-connect && git add -A && git commit -m "Auto-install protocol handler when exe run without args" && git push origin main

Commit and push changes


[main 9f36686] Auto-install protocol handler when exe run without args
 2 files changed, 54 insertions(+), 3 deletions(-)
remote: . Processing 1 references
remote: Processed 1 references in total
To https://git.azcomputerguru.com/azcomputerguru/guru-connect.git
   1b810a5..9f36686  main -> main


The user says there's an error after launching the protocol handler. They didn't include the error message. Let me ask or wait for them to provide more details.


Done. Now when users double-click the downloaded `guruconnect.exe`:
- If protocol not registered → runs install mode (registers handler, shows success message)
- If already registered → runs agent mode

What's the error you're seeing after launching the protocol handler?


The error says "Invalid URL: Missing session ID". This means the URL parsing is failing. Let me check the URL format being generated in the dashboard and compare it to what the parser expects.

In the dashboard, I'm generating:
```javascript
const protocolUrl = `guruconnect://view/${connectSessionId}?server=${serverUrl}`;
```

And in install.rs, the parser expects:
```
guruconnect://view/SESSION_ID
guruconnect://view/SESSION_ID?token=API_KEY
guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY
```

The parsing code:
```rust
let path = url.path().trim_start_matches('/');
let parts: Vec<&str> = path.split('/').collect();

if parts.is_empty() {
    return Err(anyhow!("Missing action in URL"));
}

let action = parts[0];
let session_id = parts.get(1).map(|s| s.to_string())
    .ok_or_else(|| anyhow!("Missing session ID"))?;
```

The issue is that when the URL is `guruconnect://view/SESSION_ID?server=...`, the `url.path()` returns `/view/SESSION_ID` but there might be an issue with how it's being parsed.

Wait, let me think about this more carefully. The URL structure for custom protocols can be tricky. Let me check what `url::Url::parse` does with `guruconnect://view/SESSION_ID`.

Actually, for custom protocol URLs like `guruconnect://view/SESSION_ID`, the URL parser might interpret it differently. Let me check:
- scheme: guruconnect
- host: view (this might be treated as the host!)
- path: /SESSION_ID

That's the problem! In URL format, `guruconnect://view/SESSION_ID` has:

- scheme: guruconnect
- host: view
- path: /SESSION_ID

The custom protocol parsing is causing the action to be misinterpreted as the host. This means when splitting the path, we're losing the "view" action. I'll need to modify the parsing strategy to correctly extract both the action and session ID. Potential solutions include adjusting the URL format or implementing a more robust parsing method that accounts for this quirk in custom protocol URL handling.

The most straightforward approach is to leverage the host field as the action identifier, which provides a clean and minimal solution to the current parsing challenge.


The issue is URL parsing - `guruconnect://view/SESSION_ID` treats "view" as the host, not part of the path. Let me fix the parser:


Edit

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY

    let url = url::Url::parse(url)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    let path = url.path().trim_start_matches('/');
    let parts: Vec<&str> = path.split('/').collect();

    if parts.is_empty() {
        return Err(anyhow!("Missing action in URL"));
    }

    let action = parts[0];
    let session_id = parts.get(1).map(|s| s.to_string())
        .ok_or_else(|| anyhow!("Missing session ID"))?;

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY
    //
    // Note: In URL parsing, "view" becomes the host, SESSION_ID is the path

    let url = url::Url::parse(url)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    // 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('/');
    let session_id = if path.is_empty() {
        return Err(anyhow!("Missing session ID"));
    } else {
        path.split('/').next().unwrap_or("").to_string()
    };

    if session_id.is_empty() {
        return Err(anyhow!("Missing session ID"));
    }


The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\install.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   345→    // On non-Windows, assume not registered (or check ~/.local/share/applications)
   346→    false
   347→}
   348→
   349→/// Parse a guruconnect:// URL and extract session parameters
   350→pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
   351→    // Expected formats:
   352→    // guruconnect://view/SESSION_ID
   353→    // guruconnect://view/SESSION_ID?token=API_KEY
   354→    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY
   355→    //
   356→    // Note: In URL parsing, "view" becomes the host, SESSION_ID is the path
   357→
   358→    let url = url::Url::parse(url)
   359→        .map_err(|e| anyhow!("Invalid URL: {}", e))?;
   360→
   361→    if url.scheme() != "guruconnect" {
   362→        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
   363→    }
   364→
   365→    // The "action" (view/connect) is parsed as the host
   366→    let action = url.host_str()
   367→        .ok_or_else(|| anyhow!("Missing action in URL"))?;
   368→
   369→    // The session ID is the first path segment
   370→    let path = url.path().trim_start_matches('/');
   371→    let session_id = if path.is_empty() {
   372→        return Err(anyhow!("Missing session ID"));
   373→    } else {
   374→        path.split('/').next().unwrap_or("").to_string()
   375→    };
   376→
   377→    if session_id.is_empty() {
   378→        return Err(anyhow!("Missing session ID"));
   379→    }
   380→
   381→    // Extract query parameters
   382→    let mut server = None;
   383→    let mut token = None;

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY
    //
    // Note: In URL parsing, "view" becomes the host, SESSION_ID is the path

    let url = url::Url::parse(url)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    // 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('/');
    let session_id = if path.is_empty() {
        return Err(anyhow!("Missing session ID"));
    } else {
        path.split('/').next().unwrap_or("").to_string()
    };

    if session_id.is_empty() {
        return Err(anyhow!("Missing session ID"));
    }

//! Installation and protocol handler registration
//!
//! Handles:
//! - Self-installation to Program Files (with UAC) or LocalAppData (fallback)
//! - Protocol handler registration (guruconnect://)
//! - UAC elevation with graceful fallback

use anyhow::{anyhow, Result};
use tracing::{info, warn, error};

#[cfg(windows)]
use windows::{
    core::PCWSTR,
    Win32::Foundation::HANDLE,
    Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY},
    Win32::System::Threading::{GetCurrentProcess, OpenProcessToken},
    Win32::System::Registry::{
        RegCreateKeyExW, RegSetValueExW, RegCloseKey, HKEY, HKEY_CLASSES_ROOT,
        HKEY_CURRENT_USER, KEY_WRITE, REG_SZ, REG_OPTION_NON_VOLATILE,
    },
    Win32::UI::Shell::ShellExecuteW,
    Win32::UI::WindowsAndMessaging::SW_SHOWNORMAL,
};

#[cfg(windows)]
use std::ffi::OsStr;
#[cfg(windows)]
use std::os::windows::ffi::OsStrExt;

/// Install locations
pub const SYSTEM_INSTALL_PATH: &str = r"C:\Program Files\GuruConnect";
pub const USER_INSTALL_PATH: &str = r"GuruConnect"; // Relative to %LOCALAPPDATA%

/// Check if running with elevated privileges
#[cfg(windows)]
pub fn is_elevated() -> bool {
    unsafe {
        let mut token_handle = HANDLE::default();
        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
            return false;
        }

        let mut elevation = TOKEN_ELEVATION::default();
        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;

        let result = GetTokenInformation(
            token_handle,
            TokenElevation,
            Some(&mut elevation as *mut _ as *mut _),
            size,
            &mut size,
        );

        let _ = windows::Win32::Foundation::CloseHandle(token_handle);

        result.is_ok() && elevation.TokenIsElevated != 0
    }
}

#[cfg(not(windows))]
pub fn is_elevated() -> bool {
    unsafe { libc::geteuid() == 0 }
}

/// Get the install path based on elevation status
pub fn get_install_path(elevated: bool) -> std::path::PathBuf {
    if elevated {
        std::path::PathBuf::from(SYSTEM_INSTALL_PATH)
    } else {
        let local_app_data = std::env::var("LOCALAPPDATA")
            .unwrap_or_else(|_| {
                let home = std::env::var("USERPROFILE").unwrap_or_else(|_| ".".to_string());
                format!(r"{}\AppData\Local", home)
            });
        std::path::PathBuf::from(local_app_data).join(USER_INSTALL_PATH)
    }
}

/// Get the executable path
pub fn get_exe_path(install_path: &std::path::Path) -> std::path::PathBuf {
    install_path.join("guruconnect.exe")
}

/// Attempt to elevate and re-run with install command
#[cfg(windows)]
pub fn try_elevate_and_install() -> Result<bool> {
    let exe_path = std::env::current_exe()?;
    let exe_path_wide: Vec<u16> = OsStr::new(exe_path.as_os_str())
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let verb: Vec<u16> = OsStr::new("runas")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let params: Vec<u16> = OsStr::new("install --elevated")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let result = ShellExecuteW(
            None,
            PCWSTR(verb.as_ptr()),
            PCWSTR(exe_path_wide.as_ptr()),
            PCWSTR(params.as_ptr()),
            PCWSTR::null(),
            SW_SHOWNORMAL,
        );

        // ShellExecuteW returns > 32 on success
        if result.0 as usize > 32 {
            info!("UAC elevation requested");
            Ok(true)
        } else {
            warn!("UAC elevation denied or failed");
            Ok(false)
        }
    }
}

#[cfg(not(windows))]
pub fn try_elevate_and_install() -> Result<bool> {
    Ok(false)
}

/// Register the guruconnect:// protocol handler
#[cfg(windows)]
pub fn register_protocol_handler(elevated: bool) -> Result<()> {
    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);
    let exe_path_str = exe_path.to_string_lossy();

    // Command to execute: "C:\...\guruconnect.exe" "launch" "%1"
    let command = format!("\"{}\" launch \"%1\"", exe_path_str);

    // Choose registry root based on elevation
    let root_key = if elevated {
        HKEY_CLASSES_ROOT
    } else {
        // User-level registration under Software\Classes
        HKEY_CURRENT_USER
    };

    let base_path = if elevated {
        "guruconnect"
    } else {
        r"Software\Classes\guruconnect"
    };

    unsafe {
        // Create guruconnect key
        let mut protocol_key = HKEY::default();
        let key_path = to_wide(base_path);
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut protocol_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create protocol key: {:?}", result));
        }

        // Set default value (protocol description)
        let description = to_wide("GuruConnect Protocol");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&description)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set protocol description: {:?}", result));
        }

        // Set URL Protocol (empty string indicates this is a protocol handler)
        let url_protocol = to_wide("URL Protocol");
        let empty = to_wide("");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR(url_protocol.as_ptr()),
            0,
            REG_SZ,
            Some(&description_to_bytes(&empty)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set URL Protocol: {:?}", result));
        }

        let _ = RegCloseKey(protocol_key);

        // Create shell\open\command key
        let command_path = if elevated {
            r"guruconnect\shell\open\command"
        } else {
            r"Software\Classes\guruconnect\shell\open\command"
        };
        let command_key_path = to_wide(command_path);
        let mut command_key = HKEY::default();
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(command_key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut command_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create command key: {:?}", result));
        }

        // Set the command
        let command_wide = to_wide(&command);
        let result = RegSetValueExW(
            command_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&command_wide)),
        );
        if result.is_err() {
            let _ = RegCloseKey(command_key);
            return Err(anyhow!("Failed to set command: {:?}", result));
        }

        let _ = RegCloseKey(command_key);
    }

    info!("Protocol handler registered: guruconnect://");
    Ok(())
}

#[cfg(not(windows))]
pub fn register_protocol_handler(_elevated: bool) -> Result<()> {
    warn!("Protocol handler registration not supported on this platform");
    Ok(())
}

/// Install the application
pub fn install(force_user_install: bool) -> Result<()> {
    let elevated = is_elevated();

    // If not elevated and not forcing user install, try to elevate
    if !elevated && !force_user_install {
        info!("Attempting UAC elevation for system-wide install...");
        match try_elevate_and_install() {
            Ok(true) => {
                // Elevation was requested, exit this instance
                // The elevated instance will continue the install
                info!("Elevated process started, exiting current instance");
                std::process::exit(0);
            }
            Ok(false) => {
                info!("UAC denied, falling back to user install");
            }
            Err(e) => {
                warn!("Elevation failed: {}, falling back to user install", e);
            }
        }
    }

    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);

    info!("Installing to: {}", install_path.display());

    // Create install directory
    std::fs::create_dir_all(&install_path)?;

    // Copy ourselves to install location
    let current_exe = std::env::current_exe()?;
    if current_exe != exe_path {
        std::fs::copy(&current_exe, &exe_path)?;
        info!("Copied executable to: {}", exe_path.display());
    }

    // Register protocol handler
    register_protocol_handler(elevated)?;

    info!("Installation complete!");
    if elevated {
        info!("Installed system-wide to: {}", install_path.display());
    } else {
        info!("Installed for current user to: {}", install_path.display());
    }

    Ok(())
}

/// Check if the guruconnect:// protocol handler is registered
#[cfg(windows)]
pub fn is_protocol_handler_registered() -> bool {
    use windows::Win32::System::Registry::{
        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,
    };

    unsafe {
        // Check system-wide registration (HKCR\guruconnect)
        let mut key = HKEY::default();
        let key_path = to_wide("guruconnect");
        if RegOpenKeyExW(
            HKEY_CLASSES_ROOT,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }

        // Check user-level registration (HKCU\Software\Classes\guruconnect)
        let key_path = to_wide(r"Software\Classes\guruconnect");
        if RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }
    }

    false
}

#[cfg(not(windows))]
pub fn is_protocol_handler_registered() -> bool {
    // On non-Windows, assume not registered (or check ~/.local/share/applications)
    false
}

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY

    let url = url::Url::parse(url)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    let path = url.path().trim_start_matches('/');
    let parts: Vec<&str> = path.split('/').collect();

    if parts.is_empty() {
        return Err(anyhow!("Missing action in URL"));
    }

    let action = parts[0];
    let session_id = parts.get(1).map(|s| s.to_string())
        .ok_or_else(|| anyhow!("Missing session ID"))?;

    // Extract query parameters
    let mut server = None;
    let mut token = None;

    for (key, value) in url.query_pairs() {
        match key.as_ref() {
            "server" => server = Some(value.to_string()),
            "token" | "api_key" => token = Some(value.to_string()),
            _ => {}
        }
    }

    // Default server if not specified
    let server = server.unwrap_or_else(|| "wss://connect.azcomputerguru.com/ws/viewer".to_string());

    match action {
        "view" | "connect" => Ok((server, session_id, token)),
        _ => Err(anyhow!("Unknown action: {}", action)),
    }
}

// Helper functions for Windows registry operations
#[cfg(windows)]
fn to_wide(s: &str) -> Vec<u16> {
    OsStr::new(s)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect()
}

#[cfg(windows)]
fn description_to_bytes(wide: &[u16]) -> Vec<u8> {
    wide.iter()
        .flat_map(|w| w.to_le_bytes())
        .collect()
}

+    // Note: In URL parsing, "view" becomes the host, SESSION_ID is the path

+    // The "action" (view/connect) is parsed as the host

+        .ok_or_else(|| anyhow!("Missing action in URL"))?;

+    // The session ID is the first path segment

-    let parts: Vec<&str> = path.split('/').collect();

+    let session_id = if path.is_empty() {

+        return Err(anyhow!("Missing session ID"));

+        path.split('/').next().unwrap_or("").to_string()

-        return Err(anyhow!("Missing action in URL"));

-    let session_id = parts.get(1).map(|s| s.to_string())

-        .ok_or_else(|| anyhow!("Missing session ID"))?;


Bash

cd /c/Users/MikeSwanson/Claude/guru-connect && cargo build -p guruconnect --release 2>&1 | grep -E "(Compiling guruconnect|Finished|error)"

Rebuild agent


   Compiling guruconnect v0.1.0 (C:\Users\MikeSwanson\Claude\guru-connect\agent)
warning: unused imports: `error` and `warn`
9 | use tracing::{info, warn, error};
warning: unused import: `error`
9 | use tracing::{info, warn, error};
warning: unused import: `error`
6 | use tracing::{info, warn, error};
    Finished `release` profile [optimized] target(s) in 5.44s


scp /c/Users/MikeSwanson/Claude/guru-connect/target/x86_64-pc-windows-msvc/release/guruconnect.exe guru@172.16.3.30:~/guru-connect/server/static/downloads/ && echo "Uploaded successfully"

Upload fixed binary


Uploaded successfully


cd /c/Users/MikeSwanson/Claude/guru-connect && git add -A && git commit -m "Fix protocol URL parsing - action is host, session is path" && git push origin main

Commit URL parsing fix


[main 56a9496] Fix protocol URL parsing - action is host, session is path
 1 file changed, 14 insertions(+), 7 deletions(-)
remote: . Processing 1 references
remote: Processed 1 references in total
To https://git.azcomputerguru.com/azcomputerguru/guru-connect.git
   9f36686..56a9496  main -> main


The user will need to re-download the guruconnect.exe and run it again to get the fix. But wait - they already have it installed, so the protocol handler is registered. They just need to replace the exe or re-run the installer.

Actually, the protocol handler points to the installed location (C:\Program Files\GuruConnect\ or %LOCALAPPDATA%\GuruConnect\), not the downloaded file. So they need to either:
1. Re-download and run the installer again
2. Manually copy the new exe to the install location

Let me tell the user to re-download and run install again.


Fixed. The URL parser now correctly handles `guruconnect://view/SESSION_ID` where "view" is the host and "SESSION_ID" is the path.

Re-download `guruconnect.exe` and run it again - it will update the installed copy. Then try the Native Viewer button again.


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.


The user is concerned that machines are showing up in the dashboard even though they implemented login requirements. This is a valid security concern - if machines can register/connect without authentication, that's a problem.

Let me think about what could be happening:

1. The dashboard shows machines from the `/api/sessions` endpoint
2. Machines connect via WebSocket to register themselves
3. The question is: are these endpoints protected by authentication?

Looking at the code I've seen:
- The dashboard requires login (JWT auth)
- But the API endpoints for listing sessions/machines - are they protected?
- The WebSocket connections for agents - are they authenticated?

Let me check:
1. The `/api/sessions` endpoint - is it behind auth middleware?
2. The agent WebSocket connection `/ws/agent` - does it require authentication?

From what I recall of the implementation:
- The auth middleware was added for user management endpoints
- But the session listing and agent connections may not have been protected

Let me explore the server code to understand the security model.


Good catch - this is a security concern. Let me investigate how agents are registering and if the endpoints are properly protected:


   180→        .route("/api/auth/login", post(api::auth::login))
   181→
   182→        // Auth endpoints (auth required)
   183→        .route("/api/auth/me", get(api::auth::get_me))
   184→        .route("/api/auth/change-password", post(api::auth::change_password))
   185→
   186→        // User management (admin only)
   187→        .route("/api/users", get(api::users::list_users))
   188→        .route("/api/users", post(api::users::create_user))
   189→        .route("/api/users/:id", get(api::users::get_user))
   190→        .route("/api/users/:id", put(api::users::update_user))
   191→        .route("/api/users/:id", delete(api::users::delete_user))
   192→        .route("/api/users/:id/permissions", put(api::users::set_permissions))
   193→        .route("/api/users/:id/clients", put(api::users::set_client_access))
   194→
   195→        // Portal API - Support codes
   196→        .route("/api/codes", post(create_code))
   197→        .route("/api/codes", get(list_codes))
   198→        .route("/api/codes/:code/validate", get(validate_code))
   199→        .route("/api/codes/:code/cancel", post(cancel_code))
   200→
   201→        // WebSocket endpoints
   202→        .route("/ws/agent", get(relay::agent_ws_handler))
   203→        .route("/ws/viewer", get(relay::viewer_ws_handler))
   204→
   205→        // REST API - Sessions
   206→        .route("/api/sessions", get(list_sessions))
   207→        .route("/api/sessions/:id", get(get_session))
   208→        .route("/api/sessions/:id", delete(disconnect_session))
   209→
   210→        // REST API - Machines
   211→        .route("/api/machines", get(list_machines))
   212→        .route("/api/machines/:agent_id", get(get_machine))
   213→        .route("/api/machines/:agent_id", delete(delete_machine))
   214→        .route("/api/machines/:agent_id/history", get(get_machine_history))
   215→
   216→        // HTML page routes (clean URLs)
   217→        .route("/login", get(serve_login))
   218→        .route("/dashboard", get(serve_dashboard))
   219→        .route("/users", get(serve_users))
   220→
   221→        // State and middleware
   222→        .with_state(state.clone())
   223→        .layer(middleware::from_fn_with_state(state, auth_layer))
   224→
   225→        // Serve static files for portal (fallback)
   226→        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))
   227→
   228→        // Middleware
   229→        .layer(TraceLayer::new_for_http())
   230→        .layer(
   231→            CorsLayer::new()
   232→                .allow_origin(Any)
   233→                .allow_methods(Any)
   234→                .allow_headers(Any),
   235→        );
   236→
   237→    // Start server
   238→    let addr: SocketAddr = listen_addr.parse()?;
   239→    let listener = tokio::net::TcpListener::bind(addr).await?;
   240→
   241→    info!("Server listening on {}", addr);
   242→
   243→    axum::serve(listener, app).await?;
   244→
   245→    Ok(())
   246→}
   247→
   248→async fn health() -> &'static str {
   249→    "OK"
   250→}
   251→
   252→// Support code API handlers
   253→
   254→async fn create_code(
   255→    State(state): State<AppState>,
   256→    Json(request): Json<CreateCodeRequest>,
   257→) -> Json<SupportCode> {
   258→    let code = state.support_codes.create_code(request).await;
   259→    info!("Created support code: {}", code.code);
   260→    Json(code)
   261→}
   262→
   263→async fn list_codes(
   264→    State(state): State<AppState>,
   265→) -> Json<Vec<SupportCode>> {
   266→    Json(state.support_codes.list_active_codes().await)
   267→}
   268→
   269→#[derive(Deserialize)]
   270→struct ValidateParams {
   271→    code: String,
   272→}
   273→
   274→async fn validate_code(
   275→    State(state): State<AppState>,
   276→    Path(code): Path<String>,
   277→) -> Json<CodeValidation> {
   278→    Json(state.support_codes.validate_code(&code).await)
   279→}

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

        .route("/api/auth/login", post(api::auth::login))

        // Auth endpoints (auth required)
        .route("/api/auth/me", get(api::auth::get_me))
        .route("/api/auth/change-password", post(api::auth::change_password))

        // User management (admin only)
        .route("/api/users", get(api::users::list_users))
        .route("/api/users", post(api::users::create_user))
        .route("/api/users/:id", get(api::users::get_user))
        .route("/api/users/:id", put(api::users::update_user))
        .route("/api/users/:id", delete(api::users::delete_user))
        .route("/api/users/:id/permissions", put(api::users::set_permissions))
        .route("/api/users/:id/clients", put(api::users::set_client_access))

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

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

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

        // REST API - Machines
        .route("/api/machines", get(list_machines))
        .route("/api/machines/:agent_id", get(get_machine))
        .route("/api/machines/:agent_id", delete(delete_machine))
        .route("/api/machines/:agent_id/history", get(get_machine_history))

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

        // State and middleware
        .with_state(state.clone())
        .layer(middleware::from_fn_with_state(state, auth_layer))

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

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

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

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

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

    Ok(())
}

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

// Support code API handlers

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

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

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

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


C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs:202:        .route("/ws/agent", get(relay::agent_ws_handler))
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs:206:        .route("/api/sessions", get(list_sessions))
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs:207:        .route("/api/sessions/:id", get(get_session))
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs:208:        .route("/api/sessions/:id", delete(disconnect_session))


I can see the issue. Looking at the routes:1. The auth middleware is added via `.layer(middleware::from_fn_with_state(state, auth_layer))` but this applies to ALL routes in the router.

2. However, looking at the auth_layer function, I need to see what it actually does. It probably just injects the JWT config into the request extensions but doesn't actually require authentication for all routes.

3. The key endpoints that should be protected:
   - `/api/sessions` - lists all connected machines (NOT protected, anyone can see!)
   - `/api/codes` - create/list support codes (NOT protected!)
   - `/ws/agent` - WebSocket for agent connections (NOT protected!)

This is the security hole. The endpoints themselves don't have authentication guards - they're just exposed publicly. The auth middleware only injects the JWT config, it doesn't enforce authentication.

Let me check the auth_layer function to confirm.


I see the issue. Let me check the auth_layer middleware:


C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs:49:async fn auth_layer(
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-50-    State(state): State<AppState>,
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-51-    mut request: Request,
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-52-    next: Next,
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-53-) -> impl IntoResponse {
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-54-    request.extensions_mut().insert(state.jwt_config.clone());
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-55-    next.run(request).await
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-56-}
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-57-
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-58-#[tokio::main]
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-59-async fn main() -> Result<()> {
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-60-    // Initialize logging
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-61-    let _subscriber = FmtSubscriber::builder()
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-62-        .with_max_level(Level::INFO)
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-63-        .with_target(true)
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-64-        .init();
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-65-
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-66-    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-67-
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-68-    // Load configuration
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-69-    let config = config::Config::load()?;
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-70-
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-71-    // Use port 3002 for GuruConnect
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-72-    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-73-    info!("Loaded configuration, listening on {}", listen_addr);
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-74-
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-75-    // JWT configuration
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-76-    let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| {
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-77-        tracing::warn!("JWT_SECRET not set, using default (INSECURE for production!)");
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-78-        "guruconnect-dev-secret-change-me-in-production".to_string()
C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs-79-    });


Found it! The `auth_layer` middleware (lines 49-56) does NOT enforce authentication. It only injects the JWT config into request extensions so that extractors like `AuthenticatedUser` can use it. But it doesn't actually block unauthorized requests.

This means:
1. **`/api/sessions`** - Anyone can list all connected machines (no auth required!)
2. **`/api/codes`** - Anyone can create support codes (no auth required!)
3. **`/ws/agent`** - Any agent can connect without authentication
4. **`/api/machines`** - Anyone can list machines

The authentication is only enforced on endpoints that use the `AuthenticatedUser` or `AdminUser` extractors as function parameters. Looking at the endpoint handlers:
- `api::users::*` - Uses `AdminUser` extractor (protected)
- `api::auth::get_me` - Uses `AuthenticatedUser` extractor (protected)
- But `list_sessions`, `create_code`, `list_codes`, etc. - Don't use any auth extractor (NOT protected!)

This is a significant security issue. The machines appearing in the dashboard could be:
1. Old agents that were connected before the auth update
2. New agents connecting via the unprotected `/ws/agent` endpoint

I need to:
1. Add `AuthenticatedUser` requirement to the API endpoints
2. Add agent authentication to the WebSocket handler

Let me check the relay module to see how agent connections work.


Found the security hole. The `auth_layer` only injects JWT config - it doesn't enforce auth. The actual protection comes from using `AuthenticatedUser` extractor in handlers.

**Unprotected endpoints:**
- `/api/sessions` - Anyone can list machines
- `/api/codes` - Anyone can create support codes
- `/ws/agent` - Any agent can connect
- `/api/machines` - Anyone can list/delete machines

Let me check how agents connect:


     1→//! WebSocket relay handlers
     2→//!
     3→//! Handles WebSocket connections from agents and viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→use axum::{
     7→    extract::{
     8→        ws::{Message, WebSocket, WebSocketUpgrade},
     9→        Query, State,
    10→    },
    11→    response::IntoResponse,
    12→};
    13→use futures_util::{SinkExt, StreamExt};
    14→use prost::Message as ProstMessage;
    15→use serde::Deserialize;
    16→use tracing::{error, info, warn};
    17→use uuid::Uuid;
    18→
    19→use crate::proto;
    20→use crate::session::SessionManager;
    21→use crate::db::{self, Database};
    22→use crate::AppState;
    23→
    24→#[derive(Debug, Deserialize)]
    25→pub struct AgentParams {
    26→    agent_id: String,
    27→    #[serde(default)]
    28→    agent_name: Option<String>,
    29→    #[serde(default)]
    30→    support_code: Option<String>,
    31→    #[serde(default)]
    32→    hostname: Option<String>,
    33→}
    34→
    35→#[derive(Debug, Deserialize)]
    36→pub struct ViewerParams {
    37→    session_id: String,
    38→    #[serde(default = "default_viewer_name")]
    39→    viewer_name: String,
    40→}
    41→
    42→fn default_viewer_name() -> String {
    43→    "Technician".to_string()
    44→}
    45→
    46→/// WebSocket handler for agent connections
    47→pub async fn agent_ws_handler(
    48→    ws: WebSocketUpgrade,
    49→    State(state): State<AppState>,
    50→    Query(params): Query<AgentParams>,
    51→) -> impl IntoResponse {
    52→    let agent_id = params.agent_id;
    53→    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    54→    let support_code = params.support_code;
    55→    let sessions = state.sessions.clone();
    56→    let support_codes = state.support_codes.clone();
    57→    let db = state.db.clone();
    58→
    59→    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, agent_id, agent_name, support_code))
    60→}
    61→
    62→/// WebSocket handler for viewer connections
    63→pub async fn viewer_ws_handler(
    64→    ws: WebSocketUpgrade,
    65→    State(state): State<AppState>,
    66→    Query(params): Query<ViewerParams>,
    67→) -> impl IntoResponse {
    68→    let session_id = params.session_id;
    69→    let viewer_name = params.viewer_name;
    70→    let sessions = state.sessions.clone();
    71→    let db = state.db.clone();
    72→
    73→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, db, session_id, viewer_name))
    74→}
    75→
    76→/// Handle an agent WebSocket connection
    77→async fn handle_agent_connection(
    78→    socket: WebSocket,
    79→    sessions: SessionManager,
    80→    support_codes: crate::support_codes::SupportCodeManager,
    81→    db: Option<Database>,
    82→    agent_id: String,
    83→    agent_name: String,
    84→    support_code: Option<String>,
    85→) {
    86→    info!("Agent connected: {} ({})", agent_name, agent_id);
    87→
    88→    let (mut ws_sender, mut ws_receiver) = socket.split();
    89→
    90→    // If a support code was provided, check if it's valid
    91→    if let Some(ref code) = support_code {
    92→        // Check if the code is cancelled or invalid
    93→        if support_codes.is_cancelled(code).await {
    94→            warn!("Agent tried to connect with cancelled code: {}", code);
    95→            // Send disconnect message to agent
    96→            let disconnect_msg = proto::Message {
    97→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
    98→                    reason: "Support session was cancelled by technician".to_string(),
    99→                })),
   100→            };

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

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

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

use crate::proto;
use crate::session::SessionManager;
use crate::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(),
                })),
            };


Now I understand the full picture:

**Security issues:**

1. **Agent WebSocket (`/ws/agent`)** - No authentication at all! Any agent can connect just by providing an `agent_id` parameter. There's no API key validation.

2. **Viewer WebSocket (`/ws/viewer`)** - No authentication! Anyone can connect to view a session if they know the session_id.

3. **REST APIs** - `/api/sessions`, `/api/codes`, `/api/machines` are all unprotected.

**How machines appeared:**
- Old agents from before the auth update were still connected
- OR agents could simply connect via WebSocket without any authentication
- The WebSocket handler accepts any `agent_id` parameter without validation

**What needs to be done:**
1. Add `AuthenticatedUser` to all REST API endpoints that list/manage sessions/machines
2. Add agent API key authentication to the WebSocket handshake
3. Add JWT auth for viewer WebSocket connections

For now, the immediate security concern is:
- Agents can connect without any authentication
- Anyone can view the session list/machine list without logging in

The user wants to understand how devices showed up. Let me explain and then fix it.


**Security Analysis:**

The problem is clear - **agents can connect without any authentication**:

```rust
// /ws/agent - No auth check, accepts any agent_id!
pub async fn agent_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<AgentParams>,  // Just needs agent_id param
) -> impl IntoResponse {
```

**How machines appeared:**
1. Any agent can connect to `/ws/agent?agent_id=xxx&hostname=yyy`
2. No API key, no token, no validation
3. They get listed in the dashboard immediately

**Unprotected endpoints:**
| Endpoint | Risk |
|----------|------|
| `/ws/agent` | Any agent can register |
| `/ws/viewer` | Anyone can view sessions |
| `/api/sessions` | Anyone can list machines |
| `/api/codes` | Anyone can create support codes |
| `/api/machines` | Anyone can list/delete machines |

**Recommended fixes:**
1. **Agent auth** - Require API key/enrollment token for agents
2. **Viewer auth** - Require JWT for viewer WebSocket
3. **API auth** - Add `AuthenticatedUser` to all REST endpoints

Want me to implement these fixes now? I'd suggest starting with requiring JWT auth on the REST APIs and viewer WebSocket, then adding an agent enrollment/API key system.


yes, fix it


The user wants me to fix the security issues. Let me plan the fixes:

1. **REST API endpoints** - Add `AuthenticatedUser` extractor to:
   - `/api/sessions` (list_sessions, get_session, disconnect_session)
   - `/api/codes` (create_code, list_codes, cancel_code) - but validate_code needs to stay public for agents
   - `/api/machines` (list_machines, get_machine, delete_machine, get_machine_history)

2. **Viewer WebSocket** - Add JWT validation to the WebSocket upgrade

3. **Agent WebSocket** - This is trickier. Agents need some form of authentication:
   - Option A: API key per agent (need enrollment process)
   - Option B: Support code validates the connection (already partially implemented)
   - Option C: Org-level API key that all agents share

For now, let me:
1. Secure all REST APIs with AuthenticatedUser
2. Secure viewer WebSocket with JWT
3. For agents: require either a valid support code OR an org API key

Let me start implementing.


TodoWrite

Add AuthenticatedUser to REST API endpoints (sessions, codes, machines)

Add JWT auth to viewer WebSocket

Add agent authentication (API key or support code required)

Rebuild and deploy server


Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

Add AuthenticatedUser to REST API endpoints (sessions, codes, machines)

Add JWT auth to viewer WebSocket

Add agent authentication (API key or support code required)

Rebuild and deploy server


Now let me start fixing the REST API endpoints. I need to add `AuthenticatedUser` to each handler. Let me look at the main.rs handlers for sessions, codes, and machines.


   250→}
   251→
   252→// Support code API handlers
   253→
   254→async fn create_code(
   255→    State(state): State<AppState>,
   256→    Json(request): Json<CreateCodeRequest>,
   257→) -> Json<SupportCode> {
   258→    let code = state.support_codes.create_code(request).await;
   259→    info!("Created support code: {}", code.code);
   260→    Json(code)
   261→}
   262→
   263→async fn list_codes(
   264→    State(state): State<AppState>,
   265→) -> Json<Vec<SupportCode>> {
   266→    Json(state.support_codes.list_active_codes().await)
   267→}
   268→
   269→#[derive(Deserialize)]
   270→struct ValidateParams {
   271→    code: String,
   272→}
   273→
   274→async fn validate_code(
   275→    State(state): State<AppState>,
   276→    Path(code): Path<String>,
   277→) -> Json<CodeValidation> {
   278→    Json(state.support_codes.validate_code(&code).await)
   279→}
   280→
   281→async fn cancel_code(
   282→    State(state): State<AppState>,
   283→    Path(code): Path<String>,
   284→) -> impl IntoResponse {
   285→    if state.support_codes.cancel_code(&code).await {
   286→        (StatusCode::OK, "Code cancelled")
   287→    } else {
   288→        (StatusCode::BAD_REQUEST, "Cannot cancel code")
   289→    }
   290→}
   291→
   292→// Session API handlers (updated to use AppState)
   293→
   294→async fn list_sessions(
   295→    State(state): State<AppState>,
   296→) -> Json<Vec<api::SessionInfo>> {
   297→    let sessions = state.sessions.list_sessions().await;
   298→    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
   299→}
   300→
   301→async fn get_session(
   302→    State(state): State<AppState>,
   303→    Path(id): Path<String>,
   304→) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
   305→    let session_id = uuid::Uuid::parse_str(&id)
   306→        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;
   307→
   308→    let session = state.sessions.get_session(session_id).await
   309→        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;
   310→
   311→    Ok(Json(api::SessionInfo::from(session)))
   312→}
   313→
   314→async fn disconnect_session(
   315→    State(state): State<AppState>,
   316→    Path(id): Path<String>,
   317→) -> impl IntoResponse {
   318→    let session_id = match uuid::Uuid::parse_str(&id) {
   319→        Ok(id) => id,
   320→        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
   321→    };
   322→
   323→    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
   324→        info!("Session {} disconnected by admin", session_id);
   325→        (StatusCode::OK, "Session disconnected")
   326→    } else {
   327→        (StatusCode::NOT_FOUND, "Session not found")
   328→    }
   329→}
   330→
   331→// Machine API handlers
   332→
   333→async fn list_machines(
   334→    State(state): State<AppState>,
   335→) -> Result<Json<Vec<api::MachineInfo>>, (StatusCode, &'static str)> {
   336→    let db = state.db.as_ref()
   337→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   338→
   339→    let machines = db::machines::get_all_machines(db.pool()).await
   340→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   341→
   342→    Ok(Json(machines.into_iter().map(api::MachineInfo::from).collect()))
   343→}
   344→
   345→async fn get_machine(
   346→    State(state): State<AppState>,
   347→    Path(agent_id): Path<String>,
   348→) -> Result<Json<api::MachineInfo>, (StatusCode, &'static str)> {
   349→    let db = state.db.as_ref()
   350→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   351→
   352→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   353→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
   354→        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;
   355→
   356→    Ok(Json(api::MachineInfo::from(machine)))
   357→}
   358→
   359→async fn get_machine_history(
   360→    State(state): State<AppState>,
   361→    Path(agent_id): Path<String>,
   362→) -> Result<Json<api::MachineHistory>, (StatusCode, &'static str)> {
   363→    let db = state.db.as_ref()
   364→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   365→
   366→    // Get machine
   367→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   368→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
   369→        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;
   370→
   371→    // Get sessions for this machine
   372→    let sessions = db::sessions::get_sessions_for_machine(db.pool(), machine.id).await
   373→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   374→
   375→    // Get events for this machine
   376→    let events = db::events::get_events_for_machine(db.pool(), machine.id).await
   377→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   378→
   379→    let history = api::MachineHistory {
   380→        machine: api::MachineInfo::from(machine),
   381→        sessions: sessions.into_iter().map(api::SessionRecord::from).collect(),
   382→        events: events.into_iter().map(api::EventRecord::from).collect(),
   383→        exported_at: chrono::Utc::now().to_rfc3339(),
   384→    };
   385→
   386→    Ok(Json(history))
   387→}
   388→
   389→async fn delete_machine(
   390→    State(state): State<AppState>,
   391→    Path(agent_id): Path<String>,
   392→    Query(params): Query<api::DeleteMachineParams>,
   393→) -> Result<Json<api::DeleteMachineResponse>, (StatusCode, &'static str)> {
   394→    let db = state.db.as_ref()
   395→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   396→
   397→    // Get machine first
   398→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   399→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?

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

}

// Support code API handlers

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

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

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

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

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

// Session API handlers (updated to use AppState)

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

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

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

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

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

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

// Machine API handlers

async fn list_machines(
    State(state): State<AppState>,
) -> Result<Json<Vec<api::MachineInfo>>, (StatusCode, &'static str)> {
    let db = state.db.as_ref()
        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;

    let machines = db::machines::get_all_machines(db.pool()).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;

    Ok(Json(machines.into_iter().map(api::MachineInfo::from).collect()))
}

async fn get_machine(
    State(state): State<AppState>,
    Path(agent_id): Path<String>,
) -> Result<Json<api::MachineInfo>, (StatusCode, &'static str)> {
    let db = state.db.as_ref()
        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;

    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;

    Ok(Json(api::MachineInfo::from(machine)))
}

async fn get_machine_history(
    State(state): State<AppState>,
    Path(agent_id): Path<String>,
) -> Result<Json<api::MachineHistory>, (StatusCode, &'static str)> {
    let db = state.db.as_ref()
        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;

    // Get machine
    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;

    // Get sessions for this machine
    let sessions = db::sessions::get_sessions_for_machine(db.pool(), machine.id).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;

    // Get events for this machine
    let events = db::events::get_events_for_machine(db.pool(), machine.id).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;

    let history = api::MachineHistory {
        machine: api::MachineInfo::from(machine),
        sessions: sessions.into_iter().map(api::SessionRecord::from).collect(),
        events: events.into_iter().map(api::EventRecord::from).collect(),
        exported_at: chrono::Utc::now().to_rfc3339(),
    };

    Ok(Json(history))
}

async fn delete_machine(
    State(state): State<AppState>,
    Path(agent_id): Path<String>,
    Query(params): Query<api::DeleteMachineParams>,
) -> Result<Json<api::DeleteMachineResponse>, (StatusCode, &'static str)> {
    let db = state.db.as_ref()
        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;

    // Get machine first
    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?


Now I need to add `AuthenticatedUser` to all these handlers. I'll need to:

1. Import `AuthenticatedUser` from `crate::auth`
2. Add `_user: AuthenticatedUser` parameter to each handler (the underscore prefix since we're not using the value, just requiring auth)

Let me update main.rs with these changes. First, let me check the imports at the top.


     1→//! GuruConnect Server - WebSocket Relay Server
     2→//!
     3→//! Handles connections from both agents and dashboard viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→mod config;
     7→mod relay;
     8→mod session;
     9→mod auth;
    10→mod api;
    11→mod db;
    12→mod support_codes;
    13→
    14→pub mod proto {
    15→    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
    16→}
    17→
    18→use anyhow::Result;
    19→use axum::{
    20→    Router,
    21→    routing::{get, post, put, delete},
    22→    extract::{Path, State, Json, Query, Request},
    23→    response::{Html, IntoResponse},
    24→    http::StatusCode,
    25→    middleware::{self, Next},
    26→};
    27→use std::net::SocketAddr;
    28→use std::sync::Arc;
    29→use tower_http::cors::{Any, CorsLayer};
    30→use tower_http::trace::TraceLayer;
    31→use tower_http::services::ServeDir;
    32→use tracing::{info, Level};
    33→use tracing_subscriber::FmtSubscriber;
    34→use serde::Deserialize;
    35→
    36→use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
    37→use auth::{JwtConfig, hash_password, generate_random_password};
    38→
    39→/// Application state
    40→#[derive(Clone)]
    41→pub struct AppState {
    42→    sessions: session::SessionManager,
    43→    support_codes: SupportCodeManager,
    44→    db: Option<db::Database>,
    45→    pub jwt_config: Arc<JwtConfig>,
    46→}
    47→
    48→/// Middleware to inject JWT config into request extensions
    49→async fn auth_layer(
    50→    State(state): State<AppState>,

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

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

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

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

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

use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
use auth::{JwtConfig, hash_password, generate_random_password};

/// Application state
#[derive(Clone)]
pub struct AppState {
    sessions: session::SessionManager,
    support_codes: SupportCodeManager,
    db: Option<db::Database>,
    pub jwt_config: Arc<JwtConfig>,
}

/// Middleware to inject JWT config into request extensions
async fn auth_layer(
    State(state): State<AppState>,


Now I need to:
1. Add `AuthenticatedUser` to the imports2. Add it to all the handler functions that need protectionLet me add the import first, then update each handler.


Edit

use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
use auth::{JwtConfig, hash_password, generate_random_password};

use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
use auth::{JwtConfig, hash_password, generate_random_password, AuthenticatedUser};