//! Session management for the agent //! //! Handles the lifecycle of a remote session including: //! - Connection to server //! - Idle mode (heartbeat only, minimal resources) //! - Active/streaming mode (capture and send frames) //! - Input event handling #[cfg(windows)] use windows::Win32::System::Console::{AllocConsole, GetConsoleWindow}; #[cfg(windows)] use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW}; use crate::capture::{self, Capturer, Display}; use crate::chat::{ChatController, ChatMessage as ChatMsg}; use crate::config::Config; use crate::encoder::{self, Encoder}; use crate::input::InputController; /// Show the debug console window (Windows only) #[cfg(windows)] fn show_debug_console() { unsafe { let hwnd = GetConsoleWindow(); if hwnd.0 == std::ptr::null_mut() { let _ = AllocConsole(); tracing::info!("Debug console window opened"); } else { let _ = ShowWindow(hwnd, SW_SHOW); tracing::info!("Debug console window shown"); } } } #[cfg(not(windows))] fn show_debug_console() { // No-op on non-Windows platforms } use crate::proto::{Message, message, ChatMessage, AgentStatus, Heartbeat, HeartbeatAck}; use crate::transport::WebSocketTransport; use crate::tray::{TrayController, TrayAction}; use anyhow::Result; use std::time::{Duration, Instant}; // Heartbeat interval (30 seconds) const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30); // Status report interval (60 seconds) const STATUS_INTERVAL: Duration = Duration::from_secs(60); // Update check interval (1 hour) const UPDATE_CHECK_INTERVAL: Duration = Duration::from_secs(3600); /// Session manager handles the remote control session pub struct SessionManager { config: Config, transport: Option, state: SessionState, // Lazy-initialized streaming resources capturer: Option>, encoder: Option>, input: Option, // Streaming state current_viewer_id: Option, // System info for status reports hostname: String, is_elevated: bool, start_time: Instant, } #[derive(Debug, Clone, PartialEq)] enum SessionState { Disconnected, Connecting, Idle, // Connected but not streaming - minimal resource usage Streaming, // Actively capturing and sending frames } impl SessionManager { /// Create a new session manager pub fn new(config: Config, is_elevated: bool) -> Self { let hostname = config.hostname(); Self { config, transport: None, state: SessionState::Disconnected, capturer: None, encoder: None, input: None, current_viewer_id: None, hostname, is_elevated, start_time: Instant::now(), } } /// Connect to the server pub async fn connect(&mut self) -> Result<()> { self.state = SessionState::Connecting; let transport = WebSocketTransport::connect( &self.config.server_url, &self.config.agent_id, &self.config.api_key, Some(&self.hostname), self.config.support_code.as_deref(), ).await?; self.transport = Some(transport); self.state = SessionState::Idle; // Start in idle mode tracing::info!("Connected to server, entering idle mode"); Ok(()) } /// Initialize streaming resources (capturer, encoder, input) fn init_streaming(&mut self) -> Result<()> { if self.capturer.is_some() { return Ok(()); // Already initialized } tracing::info!("Initializing streaming resources..."); tracing::info!("Capture config: use_dxgi={}, gdi_fallback={}, fps={}", self.config.capture.use_dxgi, self.config.capture.gdi_fallback, self.config.capture.fps); // Get primary display with panic protection tracing::debug!("Enumerating displays..."); let primary_display = match std::panic::catch_unwind(|| capture::primary_display()) { Ok(result) => result?, Err(e) => { tracing::error!("Panic during display enumeration: {:?}", e); return Err(anyhow::anyhow!("Display enumeration panicked")); } }; tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height); // Create capturer with panic protection // Force GDI mode if DXGI fails or panics tracing::debug!("Creating capturer (DXGI={})...", self.config.capture.use_dxgi); let capturer = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { capture::create_capturer( primary_display.clone(), self.config.capture.use_dxgi, self.config.capture.gdi_fallback, ) })) { Ok(result) => result?, Err(e) => { tracing::error!("Panic during capturer creation: {:?}", e); // Try GDI-only as last resort tracing::warn!("Attempting GDI-only capture after DXGI panic..."); capture::create_capturer(primary_display.clone(), false, false)? } }; self.capturer = Some(capturer); tracing::info!("Capturer created successfully"); // Create encoder with panic protection tracing::debug!("Creating encoder (codec={}, quality={})...", self.config.encoding.codec, self.config.encoding.quality); let encoder = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { encoder::create_encoder( &self.config.encoding.codec, self.config.encoding.quality, ) })) { Ok(result) => result?, Err(e) => { tracing::error!("Panic during encoder creation: {:?}", e); return Err(anyhow::anyhow!("Encoder creation panicked")); } }; self.encoder = Some(encoder); tracing::info!("Encoder created successfully"); // Create input controller with panic protection tracing::debug!("Creating input controller..."); let input = match std::panic::catch_unwind(InputController::new) { Ok(result) => result?, Err(e) => { tracing::error!("Panic during input controller creation: {:?}", e); return Err(anyhow::anyhow!("Input controller creation panicked")); } }; self.input = Some(input); tracing::info!("Streaming resources initialized successfully"); Ok(()) } /// Release streaming resources to save CPU/memory when idle fn release_streaming(&mut self) { if self.capturer.is_some() { tracing::info!("Releasing streaming resources"); self.capturer = None; self.encoder = None; self.input = None; self.current_viewer_id = None; } } /// Get display count for status reports fn get_display_count(&self) -> i32 { capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1) } /// Send agent status to server async fn send_status(&mut self) -> Result<()> { let status = AgentStatus { hostname: self.hostname.clone(), os_version: std::env::consts::OS.to_string(), is_elevated: self.is_elevated, uptime_secs: self.start_time.elapsed().as_secs() as i64, display_count: self.get_display_count(), is_streaming: self.state == SessionState::Streaming, agent_version: crate::build_info::short_version(), organization: self.config.company.clone().unwrap_or_default(), site: self.config.site.clone().unwrap_or_default(), tags: self.config.tags.clone(), }; let msg = Message { payload: Some(message::Payload::AgentStatus(status)), }; if let Some(transport) = self.transport.as_mut() { transport.send(msg).await?; } Ok(()) } /// Send heartbeat to server async fn send_heartbeat(&mut self) -> Result<()> { let heartbeat = Heartbeat { timestamp: chrono::Utc::now().timestamp_millis(), }; let msg = Message { payload: Some(message::Payload::Heartbeat(heartbeat)), }; if let Some(transport) = self.transport.as_mut() { transport.send(msg).await?; } Ok(()) } /// Run the session main loop with tray and chat event processing pub async fn run_with_tray(&mut self, tray: Option<&TrayController>, chat: Option<&ChatController>) -> Result<()> { if self.transport.is_none() { anyhow::bail!("Not connected"); } // Send initial status self.send_status().await?; // Timing for heartbeat and status let mut last_heartbeat = Instant::now(); let mut last_status = Instant::now(); let mut last_frame_time = Instant::now(); let mut last_update_check = Instant::now(); let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64); // Main loop loop { // Process tray events if let Some(t) = tray { if let Some(action) = t.process_events() { match action { TrayAction::EndSession => { tracing::info!("User requested session end via tray"); return Err(anyhow::anyhow!("USER_EXIT: Session ended by user")); } TrayAction::ShowDetails => { tracing::info!("User requested details (not yet implemented)"); } TrayAction::ShowDebugWindow => { show_debug_console(); } } } if t.exit_requested() { tracing::info!("Exit requested via tray"); return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user")); } } // Process incoming messages let messages: Vec = { let transport = self.transport.as_mut().unwrap(); let mut msgs = Vec::new(); while let Some(msg) = transport.try_recv()? { msgs.push(msg); } msgs }; for msg in messages { // Handle chat messages specially if let Some(message::Payload::ChatMessage(chat_msg)) = &msg.payload { if let Some(c) = chat { c.add_message(ChatMsg { id: chat_msg.id.clone(), sender: chat_msg.sender.clone(), content: chat_msg.content.clone(), timestamp: chat_msg.timestamp, }); } continue; } // Handle control messages that affect state if let Some(ref payload) = msg.payload { match payload { message::Payload::StartStream(start) => { tracing::info!("StartStream received from viewer: {}", start.viewer_id); if let Err(e) = self.init_streaming() { tracing::error!("Failed to init streaming: {}", e); } else { self.state = SessionState::Streaming; self.current_viewer_id = Some(start.viewer_id.clone()); tracing::info!("Now streaming to viewer {}", start.viewer_id); } continue; } message::Payload::StopStream(stop) => { tracing::info!("StopStream received for viewer: {}", stop.viewer_id); // Only stop if it matches current viewer if self.current_viewer_id.as_ref() == Some(&stop.viewer_id) { self.release_streaming(); self.state = SessionState::Idle; tracing::info!("Stopped streaming, returning to idle mode"); } continue; } message::Payload::Heartbeat(hb) => { // Respond to server heartbeat with ack let ack = HeartbeatAck { client_timestamp: hb.timestamp, server_timestamp: chrono::Utc::now().timestamp_millis(), }; let ack_msg = Message { payload: Some(message::Payload::HeartbeatAck(ack)), }; if let Some(transport) = self.transport.as_mut() { let _ = transport.send(ack_msg).await; } continue; } _ => {} } } // Handle other messages (input events, disconnect, etc.) self.handle_message(msg).await?; } // Check for outgoing chat messages if let Some(c) = chat { if let Some(outgoing) = c.poll_outgoing() { let chat_proto = ChatMessage { id: outgoing.id, sender: "client".to_string(), content: outgoing.content, timestamp: outgoing.timestamp, }; let msg = Message { payload: Some(message::Payload::ChatMessage(chat_proto)), }; let transport = self.transport.as_mut().unwrap(); transport.send(msg).await?; } } // State-specific behavior match self.state { SessionState::Idle => { // In idle mode, just send heartbeats and status periodically if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL { last_heartbeat = Instant::now(); if let Err(e) = self.send_heartbeat().await { tracing::warn!("Failed to send heartbeat: {}", e); } } if last_status.elapsed() >= STATUS_INTERVAL { last_status = Instant::now(); if let Err(e) = self.send_status().await { tracing::warn!("Failed to send status: {}", e); } } // Periodic update check (only for persistent agents, not support sessions) if self.config.support_code.is_none() && last_update_check.elapsed() >= UPDATE_CHECK_INTERVAL { last_update_check = Instant::now(); let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://"); match crate::update::check_for_update(&server_url).await { Ok(Some(version_info)) => { tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version); if let Err(e) = crate::update::perform_update(&version_info).await { tracing::error!("Auto-update failed: {}", e); } } Ok(None) => { tracing::debug!("No update available"); } Err(e) => { tracing::debug!("Update check failed: {}", e); } } } // Longer sleep in idle mode to reduce CPU usage tokio::time::sleep(Duration::from_millis(100)).await; } SessionState::Streaming => { // In streaming mode, capture and send frames if last_frame_time.elapsed() >= frame_interval { last_frame_time = Instant::now(); if let (Some(capturer), Some(encoder)) = (self.capturer.as_mut(), self.encoder.as_mut()) { if let Ok(Some(frame)) = capturer.capture() { if let Ok(encoded) = encoder.encode(&frame) { if encoded.size > 0 { let msg = Message { payload: Some(message::Payload::VideoFrame(encoded.frame)), }; let transport = self.transport.as_mut().unwrap(); if let Err(e) = transport.send(msg).await { tracing::warn!("Failed to send frame: {}", e); } } } } } } // Short sleep in streaming mode tokio::time::sleep(Duration::from_millis(1)).await; } _ => { // Disconnected or connecting - shouldn't be in main loop tokio::time::sleep(Duration::from_millis(100)).await; } } // Check if still connected if let Some(transport) = self.transport.as_ref() { if !transport.is_connected() { tracing::warn!("Connection lost"); break; } } else { tracing::warn!("Transport is None"); break; } } self.release_streaming(); self.state = SessionState::Disconnected; Ok(()) } /// Handle incoming message from server async fn handle_message(&mut self, msg: Message) -> Result<()> { match msg.payload { Some(message::Payload::MouseEvent(mouse)) => { if let Some(input) = self.input.as_mut() { use crate::proto::MouseEventType; use crate::input::MouseButton; match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) { MouseEventType::MouseMove => { input.mouse_move(mouse.x, mouse.y)?; } MouseEventType::MouseDown => { input.mouse_move(mouse.x, mouse.y)?; if let Some(ref buttons) = mouse.buttons { if buttons.left { input.mouse_click(MouseButton::Left, true)?; } if buttons.right { input.mouse_click(MouseButton::Right, true)?; } if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; } } } MouseEventType::MouseUp => { if let Some(ref buttons) = mouse.buttons { if buttons.left { input.mouse_click(MouseButton::Left, false)?; } if buttons.right { input.mouse_click(MouseButton::Right, false)?; } if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; } } } MouseEventType::MouseWheel => { input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?; } } } } Some(message::Payload::KeyEvent(key)) => { if let Some(input) = self.input.as_mut() { input.key_event(key.vk_code as u16, key.down)?; } } Some(message::Payload::SpecialKey(special)) => { if let Some(input) = self.input.as_mut() { use crate::proto::SpecialKey; match SpecialKey::try_from(special.key).ok() { Some(SpecialKey::CtrlAltDel) => { input.send_ctrl_alt_del()?; } _ => {} } } } Some(message::Payload::AdminCommand(cmd)) => { use crate::proto::AdminCommandType; tracing::info!("Admin command received: {:?} - {}", cmd.command, cmd.reason); match AdminCommandType::try_from(cmd.command).ok() { Some(AdminCommandType::AdminUninstall) => { tracing::warn!("Uninstall command received from server"); // Return special error to trigger uninstall in main loop return Err(anyhow::anyhow!("ADMIN_UNINSTALL: {}", cmd.reason)); } Some(AdminCommandType::AdminRestart) => { tracing::info!("Restart command received from server"); // For now, just disconnect - the auto-restart logic will handle it return Err(anyhow::anyhow!("ADMIN_RESTART: {}", cmd.reason)); } Some(AdminCommandType::AdminUpdate) => { tracing::info!("Update command received from server: {}", cmd.reason); // Trigger update check and perform update if available // The server URL is derived from the config let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://"); match crate::update::check_for_update(&server_url).await { Ok(Some(version_info)) => { tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version); if let Err(e) = crate::update::perform_update(&version_info).await { tracing::error!("Update failed: {}", e); } // If we get here, the update failed (perform_update exits on success) } Ok(None) => { tracing::info!("Already running latest version"); } Err(e) => { tracing::error!("Failed to check for updates: {}", e); } } } None => { tracing::warn!("Unknown admin command: {}", cmd.command); } } } Some(message::Payload::Disconnect(disc)) => { tracing::info!("Disconnect requested: {}", disc.reason); if disc.reason.contains("cancelled") { return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason)); } if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") { return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason)); } return Err(anyhow::anyhow!("Disconnect: {}", disc.reason)); } _ => { // Ignore unknown messages } } Ok(()) } }