diff --git a/crates/win-api-wrappers/src/process.rs b/crates/win-api-wrappers/src/process.rs index 6bad228d1f..a2f3598740 100644 --- a/crates/win-api-wrappers/src/process.rs +++ b/crates/win-api-wrappers/src/process.rs @@ -12,7 +12,7 @@ use anyhow::{Context, Result, bail}; use tracing::{error, warn}; use windows::Win32::Foundation::{ E_INVALIDARG, ERROR_INCORRECT_SIZE, ERROR_NO_MORE_FILES, FILETIME, FreeLibrary, HANDLE, HMODULE, HWND, LPARAM, - MAX_PATH, WAIT_EVENT, WAIT_FAILED, WPARAM, + MAX_PATH, WAIT_EVENT, WAIT_FAILED, WAIT_OBJECT_0, WAIT_TIMEOUT, WPARAM, }; use windows::Win32::Security::{TOKEN_ACCESS_MASK, TOKEN_ADJUST_PRIVILEGES, TOKEN_QUERY}; use windows::Win32::System::Com::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE}; @@ -31,8 +31,8 @@ use windows::Win32::System::Threading::{ CREATE_UNICODE_ENVIRONMENT, CreateProcessAsUserW, CreateRemoteThread, EXTENDED_STARTUPINFO_PRESENT, GetCurrentProcess, GetCurrentProcessId, GetExitCodeProcess, GetProcessTimes, INFINITE, LPPROC_THREAD_ATTRIBUTE_LIST, LPTHREAD_START_ROUTINE, OpenProcess, OpenProcessToken, PEB, PROCESS_ACCESS_RIGHTS, - PROCESS_BASIC_INFORMATION, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION, PROCESS_NAME_WIN32, PROCESS_TERMINATE, - QueryFullProcessImageNameW, STARTUPINFOEXW, STARTUPINFOW, STARTUPINFOW_FLAGS, TerminateProcess, + PROCESS_BASIC_INFORMATION, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION, PROCESS_NAME_WIN32, PROCESS_SYNCHRONIZE, + PROCESS_TERMINATE, QueryFullProcessImageNameW, STARTUPINFOEXW, STARTUPINFOW, STARTUPINFOW_FLAGS, TerminateProcess, WaitForSingleObject, }; use windows::Win32::UI::Shell::{SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOW, ShellExecuteExW}; @@ -961,14 +961,27 @@ fn is_process_running_impl(process_name: &str, session_id: Option) -> Resul } pub fn terminate_process_by_name(process_name: &str) -> Result { - terminate_process_by_name_impl(process_name, None) + terminate_process_by_name_impl(process_name, None, None) } pub fn terminate_process_by_name_in_session(process_name: &str, session_id: u32) -> Result { - terminate_process_by_name_impl(process_name, Some(session_id)) + terminate_process_by_name_impl(process_name, Some(session_id), None) } -fn terminate_process_by_name_impl(process_name: &str, session_id: Option) -> Result { +/// Terminates a process in the session and waits for its exit before returning success. +pub fn terminate_process_by_name_in_session_and_wait( + process_name: &str, + session_id: u32, + timeout_ms: u32, +) -> Result { + terminate_process_by_name_impl(process_name, Some(session_id), Some(timeout_ms)) +} + +fn terminate_process_by_name_impl( + process_name: &str, + session_id: Option, + timeout_ms: Option, +) -> Result { for process in ProcessEntry32Iterator::new()? { if let Some(session_id) = session_id { let actual_session = match process_id_to_session(process.process_id()) { @@ -984,22 +997,49 @@ fn terminate_process_by_name_impl(process_name: &str, session_id: Option) - } if str::eq_ignore_ascii_case(process.executable_name()?.as_str(), process_name) { - // SAFETY: `OpenProcess` is always safe to call and returns a valid handle on success. - let process = unsafe { OpenProcess(PROCESS_TERMINATE, false, process.process_id()) }; + let rights = if timeout_ms.is_some() { + PROCESS_TERMINATE | PROCESS_SYNCHRONIZE + } else { + PROCESS_TERMINATE + }; + // SAFETY: OpenProcess has no outstanding preconditions and returns an owned handle on success. + let process = unsafe { OpenProcess(rights, false, process.process_id()) }; match process { Ok(process) => { + // SAFETY: OpenProcess returned a valid handle owned by the caller. + let process = unsafe { Handle::new_owned(process) }?; // SAFETY: `OpenProcess` ensures that the handle is valid. - unsafe { - if let Err(error) = TerminateProcess(process, 1) { - warn!(process_name, session_id, %error, "TerminateProcess failed"); - return Ok(false); + if let Err(error) = unsafe { TerminateProcess(process.raw(), 1) } { + if timeout_ms.is_some() { + // SAFETY: The handle has SYNCHRONIZE access when termination is awaited. + let wait = unsafe { WaitForSingleObject(process.raw(), 0) }; + if termination_failure_is_complete(wait)? { + return Ok(true); + } + return Err(error.into()); + } + warn!(process_name, session_id, %error, "TerminateProcess failed"); + return Ok(false); + } + + if let Some(timeout_ms) = timeout_ms { + // SAFETY: The handle is valid and was opened with SYNCHRONIZE. + let wait = unsafe { WaitForSingleObject(process.raw(), timeout_ms) }; + if wait == WAIT_FAILED { + return Err(Error::last_error().into()); + } + if wait != WAIT_OBJECT_0 { + bail!("process termination did not complete within {timeout_ms} ms"); } } return Ok(true); } Err(error) => { + if timeout_ms.is_some() { + return Err(error.into()); + } warn!(process_name, session_id, %error, "OpenProcess failed"); continue; } @@ -1010,6 +1050,15 @@ fn terminate_process_by_name_impl(process_name: &str, session_id: Option) - Ok(false) } +fn termination_failure_is_complete(wait: WAIT_EVENT) -> Result { + match wait { + WAIT_OBJECT_0 => Ok(true), + WAIT_TIMEOUT => Ok(false), + WAIT_FAILED => Err(Error::last_error().into()), + status => bail!("unexpected process wait status {:#X}", status.0), + } +} + /// Get the Windows session ID for a given process ID. pub fn process_id_to_session(pid: u32) -> Result { let mut session_id = 0; @@ -1085,3 +1134,23 @@ pub fn post_message_for_pid(pid: u32, msg: u32, wparam: WPARAM, lparam: LPARAM) Ok(()) } + +#[cfg(test)] +mod tests { + use windows::Win32::Foundation::{WAIT_OBJECT_0, WAIT_TIMEOUT}; + + use super::termination_failure_is_complete; + + /// TerminateProcess may lose a race with normal process exit; a signaled process handle means termination is done. + // oracle: spec TerminateProcess §Remarks + #[test] + fn already_signaled_process_is_treated_as_terminated() { + assert!(termination_failure_is_complete(WAIT_OBJECT_0).expect("read process state")); + } + + // oracle: spec WaitForSingleObject §Return value + #[test] + fn running_process_is_not_treated_as_terminated() { + assert!(!termination_failure_is_complete(WAIT_TIMEOUT).expect("read process state")); + } +} diff --git a/crates/win-api-wrappers/src/wts.rs b/crates/win-api-wrappers/src/wts.rs index e509018456..acc2b969c1 100644 --- a/crates/win-api-wrappers/src/wts.rs +++ b/crates/win-api-wrappers/src/wts.rs @@ -196,6 +196,37 @@ pub fn get_session_domain_name(session_id: u32) -> anyhow::Result { query_session_information_string(session_id, RemoteDesktop::WTSDomainName) } +/// Checks for RDP protocol type 2, as defined by [`WTS_INFO_CLASS` "Constants"][protocol]. +/// +/// [protocol]: https://learn.microsoft.com/en-us/windows/win32/api/wtsapi32/ne-wtsapi32-wts_info_class#constants +pub fn session_uses_rdp(session_id: u32) -> anyhow::Result { + let mut buffer = windows::core::PWSTR::null(); + let mut len = 0; + + // SAFETY: Both output pointers are valid for the duration of the call. + unsafe { + WTSQuerySessionInformationW( + Some(WTS_CURRENT_SERVER_HANDLE), + session_id, + RemoteDesktop::WTSClientProtocolType, + &mut buffer, + &mut len, + )?; + } + + // SAFETY: A successful query returns memory allocated by WTS, released by WtsMemory. + let buffer = unsafe { WtsMemory::from_raw(buffer.as_ptr()) }; + if buffer.as_ptr().is_null() || len != 2 + /* ClientProtocolType */ + { + anyhow::bail!("WTSClientProtocolType returned an invalid buffer"); + } + + // SAFETY: WTSClientProtocolType returns one aligned USHORT, and its length was checked above. + let protocol = unsafe { *buffer.as_ptr() }; + Ok(protocol == 2) +} + pub fn send_message_to_session( session_id: u32, title: &str, diff --git a/devolutions-agent/src/session_manager.rs b/devolutions-agent/src/session_manager.rs index 28c461ca55..fa02c7b03a 100644 --- a/devolutions-agent/src/session_manager.rs +++ b/devolutions-agent/src/session_manager.rs @@ -2,17 +2,19 @@ use std::collections::BTreeMap; use std::fmt::Debug; +use std::time::{Duration, Instant}; use async_trait::async_trait; -use camino::Utf8PathBuf; +use camino::{Utf8Path, Utf8PathBuf}; use ceviche::controller::Session; use devolutions_gateway_task::{ShutdownSignal, Task}; use tokio::sync::{RwLock, mpsc}; use win_api_wrappers::process::{ - StartupInfo, create_process_in_session, is_process_running_in_session, terminate_process_by_name_in_session, + Process, StartupInfo, create_process_in_session, is_process_running_in_session, + terminate_process_by_name_in_session_and_wait, }; use win_api_wrappers::utils::{CommandLine, WideString}; -use win_api_wrappers::wts::session_has_logged_in_user; +use win_api_wrappers::wts::{WTSConnectState, get_sessions, session_has_logged_in_user, session_uses_rdp}; use windows::Win32::System::Threading::{ CREATE_NEW_CONSOLE, CREATE_UNICODE_ENVIRONMENT, NORMAL_PRIORITY_CLASS, STARTF_USESHOWWINDOW, }; @@ -73,6 +75,14 @@ struct GatewaySession { session: Session, kind: SessionKind, readiness: SessionReadiness, + /// Launch id of the last process this agent started in the session, if any. + launch_id: Option, + /// Number of consecutive quick failures of the session process. + consecutive_failures: u32, + /// Generation of a start request that has not started or adopted a process yet. + start_retry_id: Option, + /// A reconnect must finish terminating its old host before lookup can adopt another one. + termination_pending: bool, } // NOTE: `ceviche::controller::Session` do not implement `Debug` for session. @@ -82,6 +92,9 @@ impl Debug for GatewaySession { .field("session", &self.session.id) .field("kind", &self.kind) .field("readiness", &self.readiness) + .field("launch_id", &self.launch_id) + .field("consecutive_failures", &self.consecutive_failures) + .field("start_retry_id", &self.start_retry_id) .finish() } } @@ -92,6 +105,10 @@ impl GatewaySession { session, kind, readiness, + launch_id: None, + consecutive_failures: 0, + start_retry_id: None, + termination_pending: false, } } @@ -119,12 +136,58 @@ impl GatewaySession { fn set_waiting_for_logon(&mut self) { self.readiness = SessionReadiness::WaitingForLogon; + self.consecutive_failures = 0; + self.start_retry_id = None; } + + /// Whether the process started with `launch_id` is the one this session still wants running. + fn wants_launch(&self, launch_id: u64) -> bool { + self.kind().is_remote() && self.readiness == SessionReadiness::Ready && self.launch_id == Some(launch_id) + } +} + +/// Exit of a session process started by this agent. +#[derive(Debug, Clone, Copy)] +struct ProcessExit { + session_id: u32, + launch_id: u64, + exit_code: u32, + run_time: Duration, +} + +/// Events the session manager sends to itself. +#[derive(Debug)] +enum InternalEvent { + ProcessExited(ProcessExit), + RelaunchDue { session_id: u32, launch_id: u64 }, + StartRetryDue { session_id: u32, retry_id: u64 }, +} + +/// Delay before relaunching a process that failed for the first time. +const RELAUNCH_DELAY_FLOOR: Duration = Duration::from_secs(2); + +/// Upper bound of the relaunch delay. +const RELAUNCH_DELAY_CAP: Duration = Duration::from_secs(60); + +/// Run time after which a failure no longer counts as failing on start. +const STABLE_RUN_TIME: Duration = Duration::from_secs(60); + +/// Number of consecutive quick failures after which relaunching stops until the next logon or unlock. +const MAX_CONSECUTIVE_FAILURES: u32 = 5; + +/// Delay before relaunching after the `consecutive_failures`-th consecutive quick failure, counting from 1. +fn relaunch_delay(consecutive_failures: u32) -> Duration { + let doublings = consecutive_failures.saturating_sub(1); + 2u32.checked_pow(doublings) + .and_then(|factor| RELAUNCH_DELAY_FLOOR.checked_mul(factor)) + .map_or(RELAUNCH_DELAY_CAP, |delay| delay.min(RELAUNCH_DELAY_CAP)) } #[derive(Default, Debug)] struct SessionManagerCtx { sessions: BTreeMap, + last_launch_id: u64, + last_start_retry_id: u64, } impl SessionManagerCtx { @@ -148,6 +211,195 @@ impl SessionManagerCtx { .get(&session.to_string()) .is_some_and(|session| session.is_ready_to_start(event)) } + + /// Starts the session process in a registered session unless it is already running. + /// + /// The session is marked ready only once its process is known to run, so that a failure leaves it waiting for the + /// next readiness event. + /// A process started here gets a fresh launch id, passed to `start` and recorded on the session, so that reports + /// about an earlier process can be told apart. + /// A process found already running keeps the recorded launch id, because nothing tells whether this agent started it. + fn start_process_if_not_running( + &mut self, + session: &Session, + is_running: impl FnOnce(&Session) -> anyhow::Result, + start: impl FnOnce(&Session, u64) -> anyhow::Result<()>, + ) -> bool { + let launch_id = self.last_launch_id + 1; + + let Some(gw_session) = self.get_session_mut(session) else { + warn!(%session, "Session is not yet registered"); + return false; + }; + + match is_running(session) { + Ok(true) => { + gw_session.set_session_ready(); + return true; + } + Ok(false) => {} + Err(error) => { + // Starting a process without knowing whether one already runs could start a second one. + error!( + error = format!("{error:#}"), + %session, + "Failed to check whether the session process is running" + ); + return false; + } + } + + info!(%session, "Starting session process"); + + match start(session, launch_id) { + Ok(()) => { + info!(%session, launch_id, "Session process started"); + gw_session.set_session_ready(); + gw_session.launch_id = Some(launch_id); + self.last_launch_id = launch_id; + true + } + Err(error) => { + error!(%error, %session, "Failed to start session process"); + false + } + } + } + + /// Decides what to do about the exit of a session process this agent started. + /// + /// Returns the delay after which to relaunch the process, or `None` when it must not be relaunched. + fn handle_process_exit(&mut self, exit: &ProcessExit) -> Option { + let ProcessExit { + session_id, + launch_id, + exit_code, + run_time, + } = *exit; + + let Some(gw_session) = self + .get_session_mut(&Session::new(session_id)) + .filter(|gw_session| gw_session.wants_launch(launch_id)) + else { + debug!( + session = session_id, + launch_id, exit_code, "Ignored exit of an unwanted session process" + ); + return None; + }; + + if exit_code == 0 { + info!(session = session_id, launch_id, "Session process ended"); + return None; + } + + if STABLE_RUN_TIME <= run_time { + gw_session.consecutive_failures = 0; + } + warn!( + session = session_id, + launch_id, + exit_code = format!("{exit_code:#010X}"), + ?run_time, + "Session process failed" + ); + + self.retry_failed_launch(session_id, launch_id) + } + + fn retry_failed_launch(&mut self, session_id: u32, launch_id: u64) -> Option { + let gw_session = self + .get_session_mut(&Session::new(session_id)) + .filter(|gw_session| gw_session.wants_launch(launch_id))?; + gw_session.consecutive_failures += 1; + + if MAX_CONSECUTIVE_FAILURES <= gw_session.consecutive_failures { + error!( + session = session_id, + launch_id, + consecutive_failures = gw_session.consecutive_failures, + "Session process recovery keeps failing; waiting for the next logon or unlock" + ); + gw_session.readiness = SessionReadiness::WaitingForLogonOrUnlock; + gw_session.consecutive_failures = 0; + return None; + } + + let delay = relaunch_delay(gw_session.consecutive_failures); + warn!( + session = session_id, + launch_id, + ?delay, + "Scheduling session process recovery" + ); + + Some(delay) + } + + fn begin_start_retry(&mut self, session: &Session) -> Option { + let key = session.to_string(); + if !self.sessions.contains_key(&key) { + return None; + } + + self.last_start_retry_id += 1; + let retry_id = self.last_start_retry_id; + let gw_session = self.sessions.get_mut(&key)?; + gw_session.start_retry_id = Some(retry_id); + gw_session.consecutive_failures = 0; + Some(retry_id) + } + + fn complete_start_retry(&mut self, session_id: u32, retry_id: u64) { + if let Some(gw_session) = self + .get_session_mut(&Session::new(session_id)) + .filter(|gw_session| gw_session.start_retry_id == Some(retry_id)) + { + gw_session.start_retry_id = None; + } + } + + fn retry_failed_start(&mut self, session_id: u32, retry_id: u64) -> Option { + let gw_session = self + .get_session_mut(&Session::new(session_id)) + .filter(|gw_session| gw_session.start_retry_id == Some(retry_id))?; + gw_session.consecutive_failures += 1; + + if MAX_CONSECUTIVE_FAILURES <= gw_session.consecutive_failures { + error!( + session = session_id, + retry_id, + consecutive_failures = gw_session.consecutive_failures, + "Session process recovery keeps failing; waiting for the next logon or unlock" + ); + gw_session.readiness = SessionReadiness::WaitingForLogonOrUnlock; + gw_session.consecutive_failures = 0; + gw_session.start_retry_id = None; + return None; + } + + let delay = relaunch_delay(gw_session.consecutive_failures); + warn!( + session = session_id, + retry_id, + ?delay, + "Scheduling session process start retry" + ); + Some(delay) + } + + fn should_retry_start(&self, session_id: u32, retry_id: u64) -> bool { + self.sessions + .get(&Session::new(session_id).to_string()) + .is_some_and(|gw_session| gw_session.start_retry_id == Some(retry_id)) + } + + /// Whether a relaunch scheduled after the exit of the process started with `launch_id` is still wanted. + fn should_relaunch(&self, session_id: u32, launch_id: u64) -> bool { + self.sessions + .get(&Session::new(session_id).to_string()) + .is_some_and(|gw_session| gw_session.wants_launch(launch_id)) + } } pub struct SessionManager { @@ -180,7 +432,75 @@ impl Task for SessionManager { const NAME: &'static str = "SessionManager"; - async fn run(self, mut shutdown_signal: ShutdownSignal) -> anyhow::Result<()> { + async fn run(self, shutdown_signal: ShutdownSignal) -> anyhow::Result<()> { + self.run_with(shutdown_signal, &WindowsSessionSystem).await + } +} + +trait SessionSystem: Sync { + fn active_remote_sessions(&self) -> anyhow::Result>; + fn has_logged_in_user(&self, session: &Session) -> anyhow::Result; + fn is_running(&self, session: &Session) -> anyhow::Result; + fn start_watched( + &self, + session: &Session, + launch_id: u64, + internal_tx: &mpsc::UnboundedSender, + ) -> anyhow::Result<()>; + fn terminate(&self, session: &Session) -> anyhow::Result<()>; +} + +struct WindowsSessionSystem; + +impl SessionSystem for WindowsSessionSystem { + fn active_remote_sessions(&self) -> anyhow::Result> { + let mut active_sessions = Vec::new(); + for session in get_sessions()? { + if session.state != WTSConnectState::Active { + continue; + } + + match session_uses_rdp(session.session_id) { + Ok(true) => active_sessions.push(Session::new(session.session_id)), + Ok(false) => {} + Err(error) => { + warn!( + session_id = session.session_id, + %error, + "Failed to query the protocol of an existing session" + ); + } + } + } + Ok(active_sessions) + } + + fn has_logged_in_user(&self, session: &Session) -> anyhow::Result { + session_has_logged_in_user(session.id) + } + + fn is_running(&self, session: &Session) -> anyhow::Result { + is_session_running_in_session(session) + } + + fn start_watched( + &self, + session: &Session, + launch_id: u64, + internal_tx: &mpsc::UnboundedSender, + ) -> anyhow::Result<()> { + let process = start_session_process(session)?; + watch_session_process(process, session.id, launch_id, internal_tx.clone()); + Ok(()) + } + + fn terminate(&self, session: &Session) -> anyhow::Result<()> { + terminate_session_process(session) + } +} + +impl SessionManager { + async fn run_with(self, mut shutdown_signal: ShutdownSignal, system: &dyn SessionSystem) -> anyhow::Result<()> { let Self { mut service_event_rx, ctx, @@ -189,8 +509,67 @@ impl Task for SessionManager { info!("Starting session manager..."); + // The loop keeps a sender, so the receiver never reports a closed channel. + let (internal_tx, mut internal_rx) = mpsc::unbounded_channel(); + + // Keep queued SCM events so disconnects during enumeration still stop their hosts. + // UNTESTED: SCM delivers session changes while WTS enumeration runs. + // When changing service startup, verify this ordering in a live RDP session. + match system.active_remote_sessions() { + Ok(sessions) => { + let mut ctx = ctx.write().await; + for session in sessions { + ctx.register_session(&session, SessionKind::Remote, SessionReadiness::WaitingForLogonOrUnlock); + request_session_process_start(&mut ctx, system, &session, false, &internal_tx); + } + } + Err(error) => { + error!(%error, "Failed to enumerate existing remote sessions"); + } + } + loop { tokio::select! { + Some(event) = internal_rx.recv() => { + let mut ctx = ctx.write().await; + match event { + InternalEvent::ProcessExited(exit) => { + if let Some(delay) = ctx.handle_process_exit(&exit) { + schedule_relaunch(&internal_tx, exit.session_id, exit.launch_id, delay); + } + } + InternalEvent::RelaunchDue { session_id, launch_id } => { + if ctx.should_relaunch(session_id, launch_id) { + info!(session = session_id, "Relaunching session process"); + if !start_session_process_if_not_running(&mut ctx, system, &Session::new(session_id), &internal_tx) + && let Some(delay) = ctx.retry_failed_launch(session_id, launch_id) + { + schedule_relaunch(&internal_tx, session_id, launch_id, delay); + } + } + } + InternalEvent::StartRetryDue { session_id, retry_id } => { + if ctx.should_retry_start(session_id, retry_id) { + match system.active_remote_sessions() { + Ok(sessions) if sessions.iter().any(|session| session.id == session_id) => { + if start_session_process_if_not_running(&mut ctx, system, &Session::new(session_id), &internal_tx) { + ctx.complete_start_retry(session_id, retry_id); + } else if let Some(delay) = ctx.retry_failed_start(session_id, retry_id) { + schedule_start_retry(&internal_tx, session_id, retry_id, delay); + } + } + Ok(_) => ctx.complete_start_retry(session_id, retry_id), + Err(error) => { + error!(%error, session = session_id, "Failed to check whether the remote session is active for start retry"); + if let Some(delay) = ctx.retry_failed_start(session_id, retry_id) { + schedule_start_retry(&internal_tx, session_id, retry_id, delay); + } + } + } + } + } + } + } event = service_event_rx.recv() => { info!("Received service event"); @@ -215,15 +594,14 @@ impl Task for SessionManager { } AgentServiceEvent::SessionDisconnect(id) => { info!(%id, "Session disconnected"); - terminate_session_process(&id); + if let Err(error) = system.terminate(&id) { + error!(%id, %error, "Failed to terminate disconnected session process"); + } ctx.write().await.unregister_session(&id); } AgentServiceEvent::SessionRemoteConnect(id) => { info!(%id, "Remote session connected"); - // Terminate old session process if it is already running. - terminate_session_process(&id); - - let readiness = match session_has_logged_in_user(id.id) { + let readiness = match system.has_logged_in_user(&id) { Ok(true) => SessionReadiness::WaitingForUnlock, Ok(false) => SessionReadiness::WaitingForLogon, Err(error) => { @@ -239,6 +617,8 @@ impl Task for SessionManager { { let mut ctx = ctx.write().await; ctx.register_session(&id, SessionKind::Remote, readiness); + ctx.get_session_mut(&id).expect("session was just registered").termination_pending = true; + request_session_process_start(&mut ctx, system, &id, true, &internal_tx); } } AgentServiceEvent::SessionRemoteDisconnect(id) => { @@ -246,7 +626,9 @@ impl Task for SessionManager { // Terminate session process when remote session is disconnected // (NOTE: depending on the system settings, session could // still be running in the background after RDP disconnect). - terminate_session_process(&id); + if let Err(error) = system.terminate(&id) { + error!(%id, %error, "Failed to terminate disconnected session process"); + } ctx.write().await.unregister_session(&id); } AgentServiceEvent::SessionLogon(id) => { @@ -257,7 +639,7 @@ impl Task for SessionManager { // session app on logon, if not yet started. let mut ctx = ctx.write().await; if ctx.should_start_session(&id, SessionReadinessEvent::Logon) { - start_session_process_if_not_running(&mut ctx, &id)?; + request_session_process_start(&mut ctx, system, &id, false, &internal_tx); } } AgentServiceEvent::SessionUnlock(id) => { @@ -265,7 +647,7 @@ impl Task for SessionManager { let mut ctx = ctx.write().await; if ctx.should_start_session(&id, SessionReadinessEvent::Unlock) { - start_session_process_if_not_running(&mut ctx, &id)?; + request_session_process_start(&mut ctx, system, &id, false, &internal_tx); } } AgentServiceEvent::SessionLogoff(id) => { @@ -295,47 +677,168 @@ impl Task for SessionManager { } /// Starts Devolutions Session process in the target session. -fn start_session_process_if_not_running(ctx: &mut SessionManagerCtx, session: &Session) -> anyhow::Result<()> { - match ctx.get_session_mut(session) { - Some(gw_session) => { - if is_session_running_in_session(session)? { - gw_session.set_session_ready(); - return Ok(()); - } +/// +/// A process started here is watched, and its exit is reported through `internal_tx`. +/// A process found already running, such as one started by an earlier agent instance, is not watched, so its exit is +/// not relaunched until the next logon or unlock. +fn start_session_process_if_not_running( + ctx: &mut SessionManagerCtx, + system: &dyn SessionSystem, + session: &Session, + internal_tx: &mpsc::UnboundedSender, +) -> bool { + if !finish_session_termination(ctx, system, session) { + return false; + } + start_registered_session_process(ctx, system, session, internal_tx) +} - info!(%session, "Starting session process"); +fn start_registered_session_process( + ctx: &mut SessionManagerCtx, + system: &dyn SessionSystem, + session: &Session, + internal_tx: &mpsc::UnboundedSender, +) -> bool { + ctx.start_process_if_not_running( + session, + |session| system.is_running(session), + |session, launch_id| system.start_watched(session, launch_id, internal_tx), + ) +} - match start_session_process(session) { - Ok(()) => { - info!(%session, "Session process started"); - gw_session.set_session_ready(); - } - Err(error) => { - error!(%error, %session, "Failed to start session process"); +fn request_session_process_start( + ctx: &mut SessionManagerCtx, + system: &dyn SessionSystem, + session: &Session, + verify_active_remote_session: bool, + internal_tx: &mpsc::UnboundedSender, +) { + let Some(retry_id) = ctx.begin_start_retry(session) else { + return; + }; + + let mut should_start = finish_session_termination(ctx, system, session); + if should_start && verify_active_remote_session { + match system.active_remote_sessions() { + Ok(sessions) => { + should_start = sessions.iter().any(|active| active.id == session.id); + if !should_start { + ctx.complete_start_retry(session.id, retry_id); + return; } } + Err(error) => { + error!(%error, %session, "Failed to check whether the remote session is active"); + should_start = false; + } } - None => { - warn!(%session, "Session is not yet registered"); - } + } + + if should_start && start_registered_session_process(ctx, system, session, internal_tx) { + ctx.complete_start_retry(session.id, retry_id); + } else if let Some(delay) = ctx.retry_failed_start(session.id, retry_id) { + schedule_start_retry(internal_tx, session.id, retry_id, delay); + } +} + +fn finish_session_termination(ctx: &mut SessionManagerCtx, system: &dyn SessionSystem, session: &Session) -> bool { + let Some(gw_session) = ctx.get_session_mut(session) else { + return false; }; + if gw_session.termination_pending { + if let Err(error) = system.terminate(session) { + error!(%session, %error, "Failed to finish terminating the previous session process"); + return false; + } + gw_session.termination_pending = false; + } + true +} - Ok(()) +fn schedule_start_retry( + internal_tx: &mpsc::UnboundedSender, + session_id: u32, + retry_id: u64, + delay: Duration, +) { + let internal_tx = internal_tx.clone(); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + let _ = internal_tx.send(InternalEvent::StartRetryDue { session_id, retry_id }); + }); +} + +/// Reports the exit of `process` to the session manager from a dedicated thread. +/// +/// The thread blocks until the process exits, which may outlive the session manager. +fn watch_session_process( + process: Process, + session_id: u32, + launch_id: u64, + internal_tx: mpsc::UnboundedSender, +) { + let started_at = Instant::now(); + + // UNTESTED: the watcher is a detached `std::thread`, never a `tokio::task::spawn_blocking` task, and is never + // joined, because the runtime waits for blocking tasks at shutdown and a long-lived process would then hang the + // service stop. + // When changing this, check that stopping the service with a running session process returns promptly. + let spawn_result = std::thread::Builder::new() + .name(format!("session-{session_id}-process-watcher")) + .spawn(move || { + let exit_code = match process.wait(None).and_then(|_| process.exit_code()) { + Ok(exit_code) => exit_code, + Err(error) => { + error!(%error, session = session_id, "Failed to wait for session process exit"); + return; + } + }; + + let exit = ProcessExit { + session_id, + launch_id, + exit_code, + run_time: started_at.elapsed(), + }; + + // The session manager is gone once it shuts down, and then nothing is left to relaunch the process. + let _ = internal_tx.send(InternalEvent::ProcessExited(exit)); + }); + + if let Err(error) = spawn_result { + error!(%error, session = session_id, "Failed to watch session process"); + } +} + +/// Sends a relaunch-due event for the process started with `launch_id` after `delay`. +fn schedule_relaunch( + internal_tx: &mpsc::UnboundedSender, + session_id: u32, + launch_id: u64, + delay: Duration, +) { + let internal_tx = internal_tx.clone(); + + // UNTESTED: the timer is a detached task that nothing awaits, so that shutting down the session manager does not + // wait for a pending relaunch. + // When changing this, check that stopping the service during a relaunch delay returns promptly. + tokio::spawn(async move { + tokio::time::sleep(delay).await; + let _ = internal_tx.send(InternalEvent::RelaunchDue { session_id, launch_id }); + }); } /// Terminates Devolutions Session process in the target session. -fn terminate_session_process(session: &Session) { - match terminate_process_by_name_in_session(SESSION_BINARY, session.id) { - Ok(false) => { +fn terminate_session_process(session: &Session) -> anyhow::Result<()> { + match terminate_process_by_name_in_session_and_wait(SESSION_BINARY, session.id, 1_000)? { + false => { trace!(%session, "Session process is not running"); } - Ok(true) => { + true => { info!(%session, "Session process terminated"); } - Err(error) => { - error!(%error, %session, "Failed to terminate session process"); - } } + Ok(()) } fn is_session_running_in_session(session: &Session) -> anyhow::Result { @@ -343,13 +846,13 @@ fn is_session_running_in_session(session: &Session) -> anyhow::Result { Ok(is_running) } -fn start_session_process(session: &Session) -> anyhow::Result<()> { +fn start_session_process(session: &Session) -> anyhow::Result { if !session_has_logged_in_user(session.id)? { anyhow::bail!("Session {} does not have a logged in user", session); } let session_app_path = session_app_path(); - let command_line = CommandLine::new(vec!["--session".to_owned(), session.to_string()]); + let command_line = session_command_line(&session_app_path, session); info!(%session, "Starting `{session_app_path}`"); @@ -378,9 +881,9 @@ fn start_session_process(session: &Session) -> anyhow::Result<()> { ); match start_result { - Ok(_) => { - info!(%session, "{SESSION_BINARY} started"); - Ok(()) + Ok(process_information) => { + info!(%session, process_id = process_information.process_id, "{SESSION_BINARY} started"); + Ok(process_information.process) } Err(error) => { error!(%error, %session, "Failed to start {SESSION_BINARY}"); @@ -389,6 +892,18 @@ fn start_session_process(session: &Session) -> anyhow::Result<()> { } } +/// Builds the session process command line. +/// +/// The process receives the command line unchanged as its whole command line, so it must start with the module name +/// that the process reads as `argv[0]`. +fn session_command_line(session_app_path: &Utf8Path, session: &Session) -> CommandLine { + CommandLine::new(vec![ + session_app_path.as_str().to_owned(), + "--session".to_owned(), + session.to_string(), + ]) +} + fn session_app_path() -> Utf8PathBuf { let mut current_dir = Utf8PathBuf::from_path_buf(std::env::current_exe().expect("BUG: can't get current exe path")) .expect("BUG: OS should always return valid UTF-8 executable path"); @@ -401,12 +916,499 @@ fn session_app_path() -> Utf8PathBuf { #[cfg(test)] mod tests { - use super::{GatewaySession, Session, SessionKind, SessionReadiness, SessionReadinessEvent}; + use std::collections::{BTreeSet, VecDeque}; + use std::sync::Mutex; + use std::time::Duration; + + use camino::Utf8Path; + use devolutions_gateway_task::{ChildTask, ShutdownHandle}; + use tokio::sync::mpsc; + use win_api_wrappers::utils::CommandLine; + + use super::{ + AgentServiceEvent, GatewaySession, InternalEvent, ProcessExit, Session, SessionKind, SessionManager, + SessionManagerCtx, SessionReadiness, SessionReadinessEvent, SessionSystem, relaunch_delay, + session_command_line, + }; + + #[derive(Debug, PartialEq, Eq)] + enum HostAction { + Started(u32), + Stopped(u32), + FoundRunning(u32), + } + + struct TestSessionSystem { + active_sessions: Vec, + events: mpsc::Sender, + events_during_snapshot: Mutex>, + lookup_failures: Mutex>, + launch_failures: Mutex>, + termination_failures: Mutex>, + running: Mutex>, + exit_on_start: Mutex>, + actions: mpsc::UnboundedSender, + } + + impl SessionSystem for TestSessionSystem { + fn active_remote_sessions(&self) -> anyhow::Result> { + for event in std::mem::take(&mut *self.events_during_snapshot.lock().expect("snapshot events lock")) { + assert!(self.events.try_send(event).is_ok(), "session event queue has capacity"); + } + Ok(self.active_sessions.iter().copied().map(Session::new).collect()) + } + + fn has_logged_in_user(&self, _session: &Session) -> anyhow::Result { + Ok(true) + } + + fn is_running(&self, session: &Session) -> anyhow::Result { + if self + .lookup_failures + .lock() + .expect("lookup failure lock") + .pop_front() + .unwrap_or(false) + { + anyhow::bail!("process lookup failed"); + } + let running = self + .running + .lock() + .expect("running processes lock") + .contains(&session.id); + if running { + self.actions + .send(HostAction::FoundRunning(session.id)) + .expect("host observer remains open"); + } + Ok(running) + } + + fn start_watched( + &self, + session: &Session, + launch_id: u64, + internal_tx: &mpsc::UnboundedSender, + ) -> anyhow::Result<()> { + if self + .launch_failures + .lock() + .expect("launch failure lock") + .pop_front() + .unwrap_or(false) + { + anyhow::bail!("session process launch failed"); + } + assert!( + self.running.lock().expect("running processes lock").insert(session.id), + "duplicate host launch" + ); + self.actions + .send(HostAction::Started(session.id)) + .expect("host observer remains open"); + if let Some(exit_code) = self.exit_on_start.lock().expect("process exit lock").take() { + self.running.lock().expect("running processes lock").remove(&session.id); + internal_tx + .send(InternalEvent::ProcessExited(ProcessExit { + session_id: session.id, + launch_id, + exit_code, + run_time: Duration::ZERO, + })) + .expect("session manager remains open"); + } + Ok(()) + } + + fn terminate(&self, session: &Session) -> anyhow::Result<()> { + if self + .termination_failures + .lock() + .expect("termination failure lock") + .pop_front() + .unwrap_or(false) + { + anyhow::bail!("previous process is still terminating"); + } + if self.running.lock().expect("running processes lock").remove(&session.id) { + self.actions + .send(HostAction::Stopped(session.id)) + .expect("host observer remains open"); + } + Ok(()) + } + } + + fn start_test_manager( + active_sessions: Vec, + running: BTreeSet, + events_during_snapshot: Vec, + lookup_failures: Vec, + exit_on_start: Option, + termination_failures: Vec, + ) -> ( + ShutdownHandle, + ChildTask>, + mpsc::UnboundedReceiver, + ) { + start_test_manager_with_launch_failures( + active_sessions, + running, + events_during_snapshot, + lookup_failures, + Vec::new(), + exit_on_start, + termination_failures, + ) + } + + fn start_test_manager_with_launch_failures( + active_sessions: Vec, + running: BTreeSet, + events_during_snapshot: Vec, + lookup_failures: Vec, + launch_failures: Vec, + exit_on_start: Option, + termination_failures: Vec, + ) -> ( + ShutdownHandle, + ChildTask>, + mpsc::UnboundedReceiver, + ) { + let manager = SessionManager::default(); + let (shutdown, signal) = ShutdownHandle::new(); + let (actions, observations) = mpsc::unbounded_channel(); + let system = TestSessionSystem { + active_sessions, + events: manager.service_event_tx(), + events_during_snapshot: Mutex::new(events_during_snapshot), + lookup_failures: Mutex::new(lookup_failures.into()), + launch_failures: Mutex::new(launch_failures.into()), + termination_failures: Mutex::new(termination_failures.into()), + running: Mutex::new(running), + exit_on_start: Mutex::new(exit_on_start), + actions, + }; + let task = ChildTask::spawn(async move { manager.run_with(signal, &system).await }); + (shutdown, task, observations) + } + + async fn next_action(actions: &mut mpsc::UnboundedReceiver) -> HostAction { + tokio::time::timeout(Duration::from_secs(10), actions.recv()) + .await + .expect("session manager did not act") + .expect("session manager stopped") + } + + // oracle: reference deployed Agent 2BD6884A9E49A9D6FE80F6036907C48BECF180C0DBDD7AD93C3C7B80D571FA45 + // launches a missing host in an already-active RDP session without new logon or unlock events. + #[tokio::test] + async fn already_active_remote_session_launches_without_new_events() { + let (shutdown, task, mut actions) = start_test_manager(vec![42], BTreeSet::new(), vec![], vec![], None, vec![]); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + /// A transient startup process launch failure is retried without waiting for another readiness event. + // oracle: mutant an active startup session remains without a host after one launch failure. + #[tokio::test] + async fn startup_launch_failure_is_retried_without_a_new_event() { + let (shutdown, task, mut actions) = start_test_manager_with_launch_failures( + vec![42], + BTreeSet::new(), + vec![], + vec![false, false], + vec![true, false], + None, + vec![], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant replayed connection events leave an active session waiting for an already-past logon. + #[tokio::test] + async fn reconnect_during_snapshot_restores_an_already_active_host() { + let (shutdown, task, mut actions) = start_test_manager( + vec![42], + BTreeSet::new(), + vec![ + AgentServiceEvent::SessionRemoteDisconnect(Session::new(42)), + AgentServiceEvent::SessionRemoteConnect(Session::new(42)), + ], + vec![], + None, + vec![], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Stopped(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: reference deployed Agent 2BD6884A9E49A9D6FE80F6036907C48BECF180C0DBDD7AD93C3C7B80D571FA45 + // adopts an existing host on restart without replacing it. + #[tokio::test] + async fn already_running_startup_host_is_adopted() { + let (shutdown, task, mut actions) = + start_test_manager(vec![42], BTreeSet::from([42]), vec![], vec![], None, vec![]); + + assert_eq!(next_action(&mut actions).await, HostAction::FoundRunning(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + assert!( + actions.try_recv().is_err(), + "startup must neither stop nor duplicate the host" + ); + } + + // oracle: mutant discarding notifications during the initial snapshot leaves a disconnected host running. + #[tokio::test] + async fn disconnect_during_snapshot_stops_the_host() { + let (shutdown, task, mut actions) = start_test_manager( + vec![42], + BTreeSet::new(), + vec![AgentServiceEvent::SessionRemoteDisconnect(Session::new(42))], + vec![], + None, + vec![], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Stopped(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant a session absent from the snapshot loses its queued connection or logon event. + #[tokio::test] + async fn connection_during_snapshot_launches_after_logon() { + let (shutdown, task, mut actions) = start_test_manager( + vec![], + BTreeSet::new(), + vec![ + AgentServiceEvent::SessionRemoteConnect(Session::new(42)), + AgentServiceEvent::SessionLogon(Session::new(42)), + ], + vec![], + None, + vec![], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant a transient lookup error kills the manager instead of allowing a readiness retry. + #[tokio::test] + async fn process_lookup_failure_does_not_stop_later_readiness_events() { + let (shutdown, task, mut actions) = start_test_manager( + vec![], + BTreeSet::new(), + vec![ + AgentServiceEvent::SessionRemoteConnect(Session::new(42)), + AgentServiceEvent::SessionLogon(Session::new(42)), + AgentServiceEvent::SessionUnlock(Session::new(42)), + ], + vec![true], + None, + vec![], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant startup launches bypass the process-exit notification and relaunch path. + #[tokio::test] + async fn failed_startup_host_is_relaunched_through_the_watched_path() { + let (shutdown, task, mut actions) = + start_test_manager(vec![42], BTreeSet::new(), vec![], vec![], Some(1), vec![]); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + /// A transient startup process lookup failure is retried without waiting for another readiness event. + // oracle: mutant an active startup session remains without a host after one process lookup failure. + #[tokio::test] + async fn startup_lookup_failure_is_retried_without_a_new_event() { + let (shutdown, task, mut actions) = + start_test_manager(vec![42], BTreeSet::new(), vec![], vec![true, false], None, vec![]); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + /// A reconnect termination failure is retried without another logon or unlock event. + // oracle: mutant a transient reconnect termination failure prevents host recovery until another service event. + #[tokio::test] + async fn reconnect_termination_failure_is_retried_without_a_new_event() { + let (shutdown, task, mut actions) = start_test_manager( + vec![42], + BTreeSet::new(), + vec![AgentServiceEvent::SessionRemoteConnect(Session::new(42))], + vec![], + None, + vec![true, false], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Stopped(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant a transient lookup failure during relaunch leaves the session Ready without another retry. + #[tokio::test] + async fn relaunch_lookup_failure_is_retried_without_another_service_event() { + let (shutdown, task, mut actions) = + start_test_manager(vec![42], BTreeSet::new(), vec![], vec![false, true], Some(1), vec![]); + + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } + + // oracle: mutant reconnect adopts the old host before its asynchronous termination completes. + #[tokio::test] + async fn reconnect_does_not_adopt_a_host_still_terminating() { + let (shutdown, task, mut actions) = start_test_manager( + vec![42], + BTreeSet::from([42]), + vec![ + AgentServiceEvent::SessionRemoteConnect(Session::new(42)), + AgentServiceEvent::SessionUnlock(Session::new(42)), + ], + vec![], + None, + vec![true], + ); + + assert_eq!(next_action(&mut actions).await, HostAction::FoundRunning(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Stopped(42)); + assert_eq!(next_action(&mut actions).await, HostAction::Started(42)); + + shutdown.signal(); + task.join() + .await + .expect("session manager task") + .expect("session manager result"); + } fn make_session(kind: SessionKind, readiness: SessionReadiness) -> GatewaySession { GatewaySession::new(Session::new(1), kind, readiness) } + /// The session process reads its arguments after `argv[0]`, so a command line without the module name loses the + /// first one. + #[test] + fn session_command_line_starts_with_the_module_name() { + let session_app_path = Utf8Path::new(r"C:\Program Files\Devolutions\Agent\DevolutionsSession.exe"); + + let command_line = session_command_line(session_app_path, &Session::new(3)).to_command_line(); + + let argv = CommandLine::from_command_line(&command_line).0; + // oracle: spec CreateProcessAsUserW §Parameters + assert_eq!(argv[0], session_app_path.as_str()); + // oracle: mutant session arguments are dropped, reordered, or merged into the module name + assert_eq!(argv[1..], ["--session", "3"]); + } + + /// A failure to look up the session process must neither end the session manager nor mark the session ready, so + /// that the next readiness event retries. + // oracle: mutant a process lookup failure ends the session manager or marks the session ready + #[test] + fn process_lookup_failure_leaves_the_session_waiting() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + ctx.register_session(&session, SessionKind::Remote, SessionReadiness::WaitingForLogon); + + ctx.start_process_if_not_running( + &session, + |_| Err(anyhow::anyhow!("process lookup failed")), + |_, _| panic!("the process was started although the lookup failed"), + ); + + assert!(ctx.should_start_session(&session, SessionReadinessEvent::Logon)); + } + + /// Start retries use the existing exponential backoff and stop after the configured failure bound. + // oracle: mutant prelaunch recovery retries forever or skips its exponential backoff. + #[test] + fn prelaunch_retries_are_bounded_until_a_new_readiness_event() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + ctx.register_session(&session, SessionKind::Remote, SessionReadiness::WaitingForLogon); + let retry_id = ctx.begin_start_retry(&session).expect("session is registered"); + + assert_eq!(ctx.retry_failed_start(1, retry_id), Some(Duration::from_secs(2))); + assert_eq!(ctx.retry_failed_start(1, retry_id), Some(Duration::from_secs(4))); + assert_eq!(ctx.retry_failed_start(1, retry_id), Some(Duration::from_secs(8))); + assert_eq!(ctx.retry_failed_start(1, retry_id), Some(Duration::from_secs(16))); + assert_eq!(ctx.retry_failed_start(1, retry_id), None); + assert!(!ctx.should_retry_start(1, retry_id)); + assert!(ctx.should_start_session(&session, SessionReadinessEvent::Logon)); + } + #[test] fn new_user_session_waits_for_logon() { assert!(SessionReadiness::WaitingForLogon.is_satisfied_by(SessionReadinessEvent::Logon)); @@ -454,6 +1456,240 @@ mod tests { assert!(session.is_ready_to_start(SessionReadinessEvent::Unlock)); } + /// Registers a remote session and starts its process through the production path, returning its launch id. + fn start_ready_remote_session(ctx: &mut SessionManagerCtx, session: &Session) -> u64 { + ctx.register_session(session, SessionKind::Remote, SessionReadiness::WaitingForLogon); + start_again(ctx, session) + } + + /// Starts the process of an already registered session as a relaunch would, returning its launch id. + fn start_again(ctx: &mut SessionManagerCtx, session: &Session) -> u64 { + let mut launched = None; + ctx.start_process_if_not_running( + session, + |_| Ok(false), + |_, launch_id| { + launched = Some(launch_id); + Ok(()) + }, + ); + launched.expect("the process was not started") + } + + fn failed_exit(session: &Session, launch_id: u64, run_time: Duration) -> ProcessExit { + ProcessExit { + session_id: session.id, + launch_id, + exit_code: 0xC000_0005, + run_time, + } + } + + const QUICK_RUN: Duration = Duration::from_secs(1); + + /// A process that crashes while its session still wants it must come back, or the channel it hosts stays unserved. + // oracle: mutant a failed exit of a ready remote session schedules no relaunch, or one before the floor delay + #[test] + fn failed_exit_of_ready_remote_session_schedules_relaunch_after_the_floor_delay() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + + let delay = ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + + assert_eq!(delay, Some(Duration::from_secs(2))); + assert!(ctx.should_relaunch(session.id, launch_id)); + } + + /// Back-off keeps a process that fails on start from being relaunched in a tight loop. + // oracle: mutant the relaunch delay does not double per consecutive failure, or exceeds its cap + #[test] + fn relaunch_delay_doubles_per_consecutive_failure_and_is_capped() { + assert_eq!(relaunch_delay(1), Duration::from_secs(2)); + assert_eq!(relaunch_delay(2), Duration::from_secs(4)); + assert_eq!(relaunch_delay(3), Duration::from_secs(8)); + assert_eq!(relaunch_delay(6), Duration::from_secs(60)); + assert_eq!(relaunch_delay(u32::MAX), Duration::from_secs(60)); + } + + /// Consecutive quick failures of relaunched processes grow the delay. + // oracle: mutant a second quick failure keeps the first delay + #[test] + fn second_quick_failure_doubles_the_delay() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + let launch_id = start_again(&mut ctx, &session); + + let delay = ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + + assert_eq!(delay, Some(Duration::from_secs(4))); + } + + /// A process that ran for a while before failing is not failing on start, so it is relaunched promptly. + // oracle: mutant a failure after a long run keeps the grown delay + #[test] + fn long_run_before_failure_resets_the_delay() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + let launch_id = start_again(&mut ctx, &session); + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + let launch_id = start_again(&mut ctx, &session); + + let delay = ctx.handle_process_exit(&failed_exit(&session, launch_id, Duration::from_secs(60))); + + assert_eq!(delay, Some(Duration::from_secs(2))); + } + + /// A process that keeps failing on start is given up on, and the next logon or unlock tries again. + // oracle: mutant repeated quick failures relaunch forever, or block the next logon or unlock + #[test] + fn repeated_quick_failures_give_up_until_the_next_logon_or_unlock() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let mut launch_id = start_ready_remote_session(&mut ctx, &session); + for _ in 0..4 { + assert!( + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)) + .is_some() + ); + launch_id = start_again(&mut ctx, &session); + } + + let delay = ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + + assert_eq!(delay, None); + assert!(ctx.should_start_session(&session, SessionReadinessEvent::Logon)); + assert!(ctx.should_start_session(&session, SessionReadinessEvent::Unlock)); + } + + /// Exit code 0 means the process ended on purpose, such as when the client has no listener for the channel. + // oracle: mutant a clean exit is relaunched + #[test] + fn clean_exit_is_not_relaunched() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + let exit = ProcessExit { + exit_code: 0, + ..failed_exit(&session, launch_id, QUICK_RUN) + }; + + assert_eq!(ctx.handle_process_exit(&exit), None); + assert!(!ctx.should_start_session(&session, SessionReadinessEvent::Logon)); + } + + /// The exit of a process replaced by a later launch says nothing about the current one. + // oracle: mutant the exit of a replaced launch relaunches the process + #[test] + fn exit_of_a_replaced_launch_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let old_launch_id = start_ready_remote_session(&mut ctx, &session); + start_again(&mut ctx, &session); + + assert_eq!( + ctx.handle_process_exit(&failed_exit(&session, old_launch_id, QUICK_RUN)), + None + ); + } + + /// A reconnected session terminates the old process, whose exit must not start a second one. + // oracle: mutant an exit after a remote reconnect starts a second process + #[test] + fn exit_after_remote_reconnect_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.register_session(&session, SessionKind::Remote, SessionReadiness::WaitingForUnlock); + + assert_eq!( + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)), + None + ); + } + + /// A disconnected session terminates its process and wants none. + // oracle: mutant an exit after a disconnect relaunches the process + #[test] + fn exit_after_disconnect_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.unregister_session(&session); + + assert_eq!( + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)), + None + ); + } + + /// The system stops the process at logoff, and the next logon starts it. + // oracle: mutant an exit after logoff relaunches the process + #[test] + fn exit_after_logoff_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.get_session_mut(&session) + .expect("the session is registered") + .set_waiting_for_logon(); + + assert_eq!( + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)), + None + ); + } + + /// A console session never hosts the process, so a stray report for one must not start it. + // oracle: mutant an exit reported for the console session starts the process + #[test] + fn exit_in_console_session_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.register_session(&session, SessionKind::Console, SessionReadiness::Ready); + ctx.get_session_mut(&session) + .expect("the session is registered") + .launch_id = Some(launch_id); + + assert_eq!( + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)), + None + ); + } + + /// A session that logged off during the delay must not have the process started when the delay ends. + // oracle: mutant a relaunch falling due after logoff starts the process + #[test] + fn relaunch_due_after_logoff_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + ctx.get_session_mut(&session) + .expect("the session is registered") + .set_waiting_for_logon(); + + assert!(!ctx.should_relaunch(session.id, launch_id)); + } + + /// A relaunch scheduled for a launch that was since replaced must not start another process. + // oracle: mutant a relaunch falling due for a replaced launch starts another process + #[test] + fn relaunch_due_for_a_replaced_launch_is_ignored() { + let session = Session::new(1); + let mut ctx = SessionManagerCtx::default(); + let launch_id = start_ready_remote_session(&mut ctx, &session); + ctx.handle_process_exit(&failed_exit(&session, launch_id, QUICK_RUN)); + start_again(&mut ctx, &session); + + assert!(!ctx.should_relaunch(session.id, launch_id)); + } + #[test] fn ready_remote_session_is_not_started_twice() { let session = make_session(SessionKind::Remote, SessionReadiness::Ready); diff --git a/devolutions-session/Cargo.toml b/devolutions-session/Cargo.toml index 9f225aa690..b4f86bddea 100644 --- a/devolutions-session/Cargo.toml +++ b/devolutions-session/Cargo.toml @@ -61,7 +61,11 @@ optional = true features = [ "Win32_Foundation", "Win32_Globalization", + "Win32_Security", "Win32_Storage_FileSystem", + "Win32_System_IO", + "Win32_System_JobObjects", + "Win32_System_Pipes", "Win32_System_Shutdown", "Win32_UI_Accessibility", "Win32_UI_WindowsAndMessaging", diff --git a/devolutions-session/src/dvc/io.rs b/devolutions-session/src/dvc/io.rs index a4dcc6a109..1573109807 100644 --- a/devolutions-session/src/dvc/io.rs +++ b/devolutions-session/src/dvc/io.rs @@ -1,8 +1,9 @@ +use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::time::{Duration, Instant}; use anyhow::Context; -use now_proto_pdu::NowMessage; +use now_proto_pdu::{NowChannelHeartbeatMsg, NowMessage}; use tokio::sync::Notify; use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::error::TrySendError; @@ -11,10 +12,10 @@ use win_api_wrappers::event::Event; use win_api_wrappers::utils::Pipe; use win_api_wrappers::wts::WtsVirtualChannel; use windows::Win32::Foundation::{ - ERROR_GEN_FAILURE, ERROR_IO_PENDING, GetLastError, HANDLE, WAIT_EVENT, WAIT_OBJECT_0, WAIT_TIMEOUT, + ERROR_GEN_FAILURE, ERROR_IO_PENDING, GetLastError, HANDLE, WAIT_OBJECT_0, WAIT_TIMEOUT, }; use windows::Win32::Storage::FileSystem::{ReadFile, WriteFile}; -use windows::Win32::System::IO::{GetOverlappedResult, OVERLAPPED}; +use windows::Win32::System::IO::{CancelIoEx, GetOverlappedResult, OVERLAPPED}; use windows::Win32::System::RemoteDesktop::CHANNEL_PDU_HEADER; use windows::Win32::System::Threading::{INFINITE, WaitForMultipleObjects, WaitForSingleObject}; use windows::core::Owned; @@ -78,15 +79,18 @@ impl DvcInitializationError { } } -/// Counters and signalling shared between the DVC IO thread and the DVC task. +/// Counters and signalling shared between the DVC IO thread and the DVC task for one channel. /// /// Beyond diagnostics, this gates the negotiation deadline, which must not start before the -/// channel is open. +/// channel is open, and records whether negotiation completed, which decides whether a failed +/// channel is reopened. #[derive(Debug, Default)] pub struct DvcInitializationProgress { open_attempts: AtomicU32, channel_opened: AtomicBool, channel_opened_notify: Notify, + negotiated: AtomicBool, + heartbeats_enabled: AtomicBool, } impl DvcInitializationProgress { @@ -132,148 +136,356 @@ impl DvcInitializationProgress { pub fn channel_opened(&self) -> bool { self.channel_opened.load(Ordering::Relaxed) } + + pub fn mark_negotiated(&self) { + self.negotiated.store(true, Ordering::Relaxed); + } + + /// Whether the client completed capability negotiation on this channel. + pub fn negotiated(&self) -> bool { + self.negotiated.load(Ordering::Relaxed) + } + + /// Records whether the negotiated capabilities let the server send heartbeats. + pub fn set_heartbeats_enabled(&self, enabled: bool) { + self.heartbeats_enabled.store(enabled, Ordering::Relaxed); + } + + /// Whether the client accepted heartbeats on this channel. + /// + /// A client that negotiated a zero interval disabled them, so it is never sent one. + pub fn heartbeats_enabled(&self) -> bool { + self.heartbeats_enabled.load(Ordering::Relaxed) + } } struct DvcIoContext { + // Declared first so that its pending IO is cancelled before the WTS channel closes. + io: Box, _wts: WtsVirtualChannel, - channel_file: Owned, - pdu_chunk_buffer: Box<[u8]>, - overlapped: OVERLAPPED, - bytes_read: u32, - message_dissector: NowMessageDissector, - read_event: Event, } /// Run main DVC IO loop for `Devolutions::Now::Agent` channel. pub fn run_dvc_io( - mut write_rx: WinapiSignaledReceiver>, + write_rx: WinapiSignaledReceiver>, read_tx: Sender>, stop_event: Event, progress: &DvcInitializationProgress, + io_loop_ended_tx: tokio::sync::oneshot::Sender<()>, ) -> Result<(), anyhow::Error> { let Some(mut context) = initialize_dvc_with_retry(&stop_event, progress)? else { info!("DVC IO thread stopped during initialization"); return Ok(()); }; - loop { - let events = [context.read_event.raw(), write_rx.raw_wait_handle(), stop_event.raw()]; - - const WAIT_OBJECT_READ_DVC: WAIT_EVENT = WAIT_OBJECT_0; - const WAIT_OBJECT_WRITE_DVC: WAIT_EVENT = WAIT_EVENT(WAIT_OBJECT_0.0 + 1); - const WAIT_OBJECT_STOP: WAIT_EVENT = WAIT_EVENT(WAIT_OBJECT_0.0 + 2); - - // SAFETY: No preconditions. - let wait_status = unsafe { WaitForMultipleObjects(&events, false, INFINITE) }; - - match wait_status { - // Read event is signaled (incoming data from DVC channel). - WAIT_OBJECT_READ_DVC => { - trace!("DVC channel read event is signaled"); - - // SAFETY: No preconditions. - unsafe { - GetOverlappedResult( - *context.channel_file, - &context.overlapped, - &mut context.bytes_read, - false, - ) - }?; - - if context.bytes_read - < u32::try_from(size_of::()) - .expect("CHANNEL_PDU_HEADER size always fits into u32") + let result = context.io.serve(write_rx, &read_tx, &stop_event, progress); + + // Notify processing before cancellation cleanup, which can block until pending IO completes. + let _ = io_loop_ended_tx.send(()); + result +} + +/// Interval at which decoded messages are offered again to a full processing queue. +const INBOUND_RETRY_INTERVAL_MS: u32 = 10; + +/// Longest time without a write while decoded messages are held back, on a channel with heartbeats enabled. +/// +/// No read is pending while messages are held back, so only a write can observe a closed channel. +/// Without inbound progress, a probe or timeout is needed to observe a lost client. +const BACKLOG_PROBE_INTERVAL: Duration = Duration::from_secs(1); +/// Maximum time a disabled-heartbeat channel may remain without progress while messages are held back. +const BACKLOG_STALL_TIMEOUT: Duration = Duration::from_secs(30); + +/// Overlapped IO state of one open DVC channel. +/// +/// Reads and writes are both overlapped, so that a client slow to read never stops the loop from reading, and the +/// loop never blocks on the processing task. +/// Must stay boxed while IO is pending, because the OS writes to the `OVERLAPPED` structures and buffers. +struct DvcIo { + file: Owned, + + read_buffer: Box<[u8]>, + read_overlapped: OVERLAPPED, + read_event: Event, + read_pending: bool, + message_dissector: NowMessageDissector, + /// Decoded messages the processing task has no room for yet. + /// + /// No read is issued while this is not empty, which bounds it to the messages of one read. + inbound: VecDeque>, + backlog_stall_timeout: Duration, + inbound_blocked_since: Option, + + write_buffer: Vec, + write_overlapped: OVERLAPPED, + write_event: Event, + write_pending: bool, + last_write_started: Instant, +} + +#[derive(Clone, Copy)] +enum DvcIoWake { + Stop, + ReadCompleted, + WriteCompleted, + Outbound, +} + +impl DvcIo { + fn new(file: Owned) -> anyhow::Result> { + let read_event = Event::new_unnamed()?; + let write_event = Event::new_unnamed()?; + + // All DVC messages should be under CHANNEL_CHUNK_LENGTH size, but sometimes RDP stack + // sends a few messages together; 128Kb buffer should be enough to hold a few dozen messages. + Ok(Box::new(Self { + file, + read_buffer: vec![0u8; 128 * 1024].into_boxed_slice(), + read_overlapped: OVERLAPPED { + hEvent: read_event.raw(), + ..Default::default() + }, + read_event, + read_pending: false, + message_dissector: NowMessageDissector::default(), + inbound: VecDeque::new(), + backlog_stall_timeout: BACKLOG_STALL_TIMEOUT, + inbound_blocked_since: None, + write_buffer: Vec::new(), + write_overlapped: OVERLAPPED { + hEvent: write_event.raw(), + ..Default::default() + }, + write_event, + write_pending: false, + last_write_started: Instant::now(), + })) + } + + fn serve( + &mut self, + mut write_rx: WinapiSignaledReceiver>, + read_tx: &Sender>, + stop_event: &Event, + progress: &DvcInitializationProgress, + ) -> anyhow::Result<()> { + loop { + self.forward_inbound(read_tx)?; + + if self.inbound.is_empty() { + if !self.read_pending { + self.start_read()?; + } + } else { + if !progress.heartbeats_enabled() + && self + .inbound_blocked_since + .is_some_and(|since| since.elapsed() >= self.backlog_stall_timeout) { - // Channel is closed abruptly; abort loop. - return Ok(()); + anyhow::bail!( + "inbound DVC messages remained backpressured for {:?}", + self.backlog_stall_timeout + ); } - let chunk_data_size = usize::try_from(context.bytes_read) - .expect("read size can't be breater than CHANNEL_CHUNK_LENGTH, therefore it should fit into usize") - .checked_sub(size_of::()) - .expect("read size is less than header size; Correctness of this should be ensured by the OS"); + if progress.heartbeats_enabled() + && !self.write_pending + && self.last_write_started.elapsed() >= BACKLOG_PROBE_INTERVAL + { + // A heartbeat is the only message the server may send unprompted, and only when the client + // enabled heartbeats; NOW clients take one at any time, since they only watch for missing ones. + // UNTESTED: the probe must wait for any pending write, because `start_write` replaces the buffer + // the OS is still sending, and must stay spaced by `BACKLOG_PROBE_INTERVAL` so a held-back + // channel is not flooded; when changing this condition, check both still hold. + trace!("Probing the DVC channel while inbound messages are held back"); + self.start_write(&NowChannelHeartbeatMsg::default().into()) + .context("probing the DVC channel while inbound messages are held back")?; + } + } - const HEADER_SIZE: usize = size_of::(); + // The stop event comes first so that it wins over a busy channel. + let mut wakes = vec![(stop_event.raw(), DvcIoWake::Stop)]; - let messages = context - .message_dissector - .dissect(&context.pdu_chunk_buffer[HEADER_SIZE..HEADER_SIZE + chunk_data_size]) - .context("failed to dissect DVC messages"); + if self.read_pending { + wakes.push((self.read_event.raw(), DvcIoWake::ReadCompleted)); + } - let messages = match messages { - Ok(messages) => messages, - Err(err) => { - error!(?err, "Failed to dissect DVC messages"); - return Err(err); - } - }; + // One write at a time: further messages wait in `write_rx`, which holds back the processing task. + if self.write_pending { + wakes.push((self.write_event.raw(), DvcIoWake::WriteCompleted)); + } else { + wakes.push((write_rx.raw_wait_handle(), DvcIoWake::Outbound)); + } - // Send all messages over the channel. - for message in messages { - debug!(?message, "DVC message received"); - forward_to_processor(&read_tx, message)?; - } - // Prepare async read file operation one more time. - // SAFETY: No preconditions. - let result = unsafe { - ReadFile( - *context.channel_file, - Some(&mut context.pdu_chunk_buffer), - None, - Some(&mut context.overlapped), - ) - }; + let handles: Vec = wakes.iter().map(|(handle, _)| *handle).collect(); + let timeout = if self.inbound.is_empty() { + INFINITE + } else { + INBOUND_RETRY_INTERVAL_MS + }; - ensure_overlapped_io_result(result)?; + // SAFETY: All handles are valid for the duration of the call. + let wait_status = unsafe { WaitForMultipleObjects(&handles, false, timeout) }; + + if wait_status == WAIT_TIMEOUT { + continue; } - // Write event is signaled (outgoing data to DVC channel). - WAIT_OBJECT_WRITE_DVC => { - trace!("DVC channel write event is signaled"); - let message_to_write = write_rx.try_recv()?; - let message_bytes = now_proto_pdu::ironrdp_core::encode_vec(&message_to_write)?; + let wake = usize::try_from(wait_status.0.wrapping_sub(WAIT_OBJECT_0.0)) + .ok() + .and_then(|index| wakes.get(index)) + .map(|(_, wake)| *wake); - let mut dw_written: u32 = 0; + match wake { + Some(DvcIoWake::Stop) => { + info!("DVC IO thread is stopped"); + return Ok(()); + } + Some(DvcIoWake::ReadCompleted) => self.complete_read()?, + Some(DvcIoWake::WriteCompleted) => self.complete_write()?, + Some(DvcIoWake::Outbound) => { + trace!("DVC channel write event is signaled"); - // SAFETY: No preconditions. - unsafe { WriteFile(*context.channel_file, Some(&message_bytes), Some(&mut dw_written), None)? } + let message = write_rx.try_recv()?; + self.start_write(&message)?; + } + None => { + return Err(windows::core::Error::from_win32()) + .with_context(|| format!("waiting for DVC IO failed with status {:#X}", wait_status.0)); + } } - WAIT_OBJECT_STOP => { - info!("DVC IO thread is stopped"); - // Stop event is signaled; abort loop. - return Ok(()); + } + } + + fn start_read(&mut self) -> windows::core::Result<()> { + // SAFETY: The buffer and `OVERLAPPED` live in this boxed struct, whose `Drop` waits for the read to finish. + let result = unsafe { + ReadFile( + *self.file, + Some(&mut self.read_buffer), + None, + Some(&mut self.read_overlapped), + ) + }; + ensure_overlapped_io_result(result)?; + self.read_pending = true; + + Ok(()) + } + + fn complete_read(&mut self) -> anyhow::Result<()> { + trace!("DVC channel read event is signaled"); + + self.read_pending = false; + let mut bytes_read: u32 = 0; + + // SAFETY: The read was issued on this handle with this `OVERLAPPED`, and has completed. + unsafe { GetOverlappedResult(*self.file, &self.read_overlapped, &mut bytes_read, false) }?; + + const HEADER_SIZE: usize = size_of::(); + + let bytes_read = usize::try_from(bytes_read).expect("u32 fits into usize on Windows"); + let payload = self.read_buffer.get(HEADER_SIZE..bytes_read).with_context(|| { + format!("DVC read returned {bytes_read} bytes, fewer than the {HEADER_SIZE}-byte channel PDU header") + })?; + + let messages = match self.message_dissector.dissect(payload) { + Ok(messages) => messages, + Err(error) => { + error!(?error, "Failed to dissect DVC messages"); + return Err(error.context("failed to dissect DVC messages")); } - _ => { - // Spurious wakeup or wait failure + }; + + for message in messages { + debug!(?message, "DVC message received"); + self.inbound.push_back(message); + } + + Ok(()) + } + + /// Hands decoded messages to the processing task, keeping those it has no room for. + /// + /// The send never blocks, because the loop also services outbound writes: blocking here while the processing task + /// awaits an outbound send would deadlock both sides. + /// Kept messages stop further reads until the queue has room, so a slow processing task holds back the client + /// rather than losing a message. + /// NOW capability responses carry no request id, so a lost request would silently shift every later response the + /// client matches by send order. + fn forward_inbound(&mut self, read_tx: &Sender>) -> anyhow::Result<()> { + while let Some(message) = self.inbound.pop_front() { + match read_tx.try_send(message) { + Ok(()) => self.inbound_blocked_since = None, + Err(TrySendError::Full(message)) => { + self.inbound_blocked_since.get_or_insert_with(Instant::now); + trace!(pending = self.inbound.len() + 1, "DVC processing queue is full"); + self.inbound.push_front(message); + break; + } + Err(TrySendError::Closed(_)) => { + anyhow::bail!("the DVC message processing task has stopped"); + } } + } + + Ok(()) + } + + fn start_write(&mut self, message: &NowMessage<'static>) -> anyhow::Result<()> { + self.write_buffer = now_proto_pdu::ironrdp_core::encode_vec(message)?; + + // SAFETY: The buffer and `OVERLAPPED` live in this boxed struct, whose `Drop` waits for the write to finish. + let result = unsafe { + WriteFile( + *self.file, + Some(&self.write_buffer), + None, + Some(&mut self.write_overlapped), + ) }; + self.last_write_started = Instant::now(); + ensure_overlapped_io_result(result)?; + self.write_pending = true; + + Ok(()) } -} -/// Hands a decoded message to the processing task without blocking. -/// -/// The send is non-blocking because the IO thread also services outbound writes: -/// blocking here while the processing task awaits an outbound send would deadlock both sides. -/// A full queue therefore ends the channel rather than dropping the message. -/// NOW capability responses carry no request id, so a lost request would silently shift every later response the -/// client matches by send order. -/// -/// Returning an error makes [`run_dvc_io`] exit, which drops [`DvcIoContext`] and with it the WTS channel handle. -// UNTESTED: the client observes a DVC close when this returns an error. -// The close happens through `Drop` of `DvcIoContext` when `run_dvc_io` returns, and requires a live RDP session to -// observe. -// When changing the IO loop, keep `run_dvc_io` returning the error rather than continuing. -fn forward_to_processor(read_tx: &Sender>, message: NowMessage<'static>) -> anyhow::Result<()> { - match read_tx.try_send(message) { - Ok(()) => Ok(()), - Err(TrySendError::Full(_)) => { - error!("Closing the DVC channel because the processing task fell behind incoming messages"); - anyhow::bail!("DVC input queue overflow"); + fn complete_write(&mut self) -> anyhow::Result<()> { + self.write_pending = false; + let mut bytes_written: u32 = 0; + + // SAFETY: The write was issued on this handle with this `OVERLAPPED`, and has completed. + unsafe { GetOverlappedResult(*self.file, &self.write_overlapped, &mut bytes_written, false) }?; + + let bytes_written = usize::try_from(bytes_written).expect("u32 fits into usize on Windows"); + + // The rest of a short write cannot be sent in a second write: the client decodes each DVC message on its own. + if bytes_written != self.write_buffer.len() { + anyhow::bail!("DVC write sent {bytes_written} of {} bytes", self.write_buffer.len()); } - Err(error @ TrySendError::Closed(_)) => { - error!("DVC message was dropped due to closed channel"); - Err(error.into()) + + Ok(()) + } +} + +impl Drop for DvcIo { + fn drop(&mut self) { + let pending = [ + (self.read_pending, &self.read_overlapped), + (self.write_pending, &self.write_overlapped), + ]; + + for (_, overlapped) in pending.into_iter().filter(|(pending, _)| *pending) { + let mut bytes_transferred: u32 = 0; + + // The OS writes to the buffer and `OVERLAPPED` until the operation finishes, so it must finish before + // they are freed. Errors are ignored: the operation may complete before it is cancelled, and a failed + // operation is finished all the same. + // SAFETY: The handle and `OVERLAPPED` are those of an issued operation, and are still alive. + let _ = unsafe { CancelIoEx(*self.file, Some(overlapped)) }; + // SAFETY: Same as above. + let _ = unsafe { GetOverlappedResult(*self.file, overlapped, &mut bytes_transferred, true) }; } } } @@ -369,42 +581,16 @@ fn initialize_dvc() -> Result, DvcInitializationError> { error, })?; - // All DVC messages should be under CHANNEL_CHUNK_LENGTH size, but sometimes RDP stack - // sends a few messages together; 128Kb buffer should be enough to hold a few dozen messages. - let pdu_chunk_buffer = vec![0u8; 128 * 1024].into_boxed_slice(); - let read_event = Event::new_unnamed().map_err(|error| DvcInitializationError { + let io = DvcIo::new(channel_file).map_err(|error| DvcInitializationError { stage: DvcInitializationStage::StartRead, error, })?; - let overlapped = OVERLAPPED { - hEvent: read_event.raw(), - ..Default::default() - }; - Ok(Box::new(DvcIoContext { - _wts: wts, - channel_file, - pdu_chunk_buffer, - overlapped, - bytes_read: 0, - message_dissector: NowMessageDissector::default(), - read_event, - })) + Ok(Box::new(DvcIoContext { io, _wts: wts })) } fn start_initial_read(context: &mut DvcIoContext) -> Result<(), windows::core::Error> { - // Prepare async read operation. - // SAFETY: Both `channel_file` and event passed to `overlapped` are valid during this call, - // therefore it is safe to call. - let read_result: Result<(), windows::core::Error> = unsafe { - ReadFile( - *context.channel_file, - Some(&mut context.pdu_chunk_buffer), - None, - Some(&mut context.overlapped), - ) - }; - ensure_overlapped_io_result(read_result)?; + context.io.start_read()?; trace!("DVC channel opened"); @@ -469,38 +655,464 @@ impl IoRedirectionPipes { #[cfg(test)] mod tests { - use std::time::Duration; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::time::{Duration, Instant}; use futures::poll; - use now_proto_pdu::{NowChannelHeartbeatMsg, NowMessage}; + use now_proto_pdu::ironrdp_core::encode_vec; + use now_proto_pdu::{NowChannelHeartbeatMsg, NowExecDataMsg, NowExecDataStreamKind, NowMessage}; use tokio::sync::mpsc; + use win_api_wrappers::event::Event; + use windows::Win32::Foundation::{ERROR_MORE_DATA, GENERIC_READ, GENERIC_WRITE, HANDLE}; + use windows::Win32::Storage::FileSystem::{ + CreateFileW, FILE_FLAG_FIRST_PIPE_INSTANCE, FILE_FLAG_OVERLAPPED, FILE_FLAGS_AND_ATTRIBUTES, FILE_SHARE_NONE, + OPEN_EXISTING, PIPE_ACCESS_DUPLEX, ReadFile, WriteFile, + }; + use windows::Win32::System::Pipes::{ + CreateNamedPipeW, PIPE_READMODE_MESSAGE, PIPE_REJECT_REMOTE_CLIENTS, PIPE_TYPE_MESSAGE, PIPE_WAIT, + PeekNamedPipe, + }; + use windows::core::{Owned, PCWSTR}; + + use super::{DVC_RETRY_DELAYS, DvcInitializationProgress, DvcIo, dvc_retry_delay}; + use crate::dvc::channel::{WinapiSignaledSender, winapi_signaled_mpsc_channel}; + + /// A connected message-mode pipe standing in for the DVC: the loop serves the overlapped server end, and the test + /// plays the client on the synchronous end. + fn dvc_pipe() -> (usize, Owned) { + static NEXT_PIPE: AtomicU32 = AtomicU32::new(0); + + let name = format!( + r"\\.\pipe\devolutions-session-dvc-io-test-{}-{}", + std::process::id(), + NEXT_PIPE.fetch_add(1, Ordering::Relaxed) + ); + let name: Vec = name.encode_utf16().chain(Some(0)).collect(); + + // SAFETY: The name is a valid null-terminated string. + let server = unsafe { + CreateNamedPipeW( + PCWSTR(name.as_ptr()), + PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED | FILE_FLAG_FIRST_PIPE_INSTANCE, + PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT | PIPE_REJECT_REMOTE_CLIENTS, + 1, + PIPE_BUFFER_SIZE, + PIPE_BUFFER_SIZE, + 0, + None, + ) + }; + assert!(!server.is_invalid(), "{}", windows::core::Error::from_win32()); + + // SAFETY: The name is a valid null-terminated string. + let client = unsafe { + CreateFileW( + PCWSTR(name.as_ptr()), + (GENERIC_READ | GENERIC_WRITE).0, + FILE_SHARE_NONE, + None, + OPEN_EXISTING, + FILE_FLAGS_AND_ATTRIBUTES(0), + None, + ) + } + .expect("connect to the test pipe"); + + // The raw value crosses to the IO thread, which takes ownership of it. + // SAFETY: Both handles were just opened and are owned by nothing else. + (server.0 as usize, unsafe { Owned::new(client) }) + } + + const PIPE_BUFFER_SIZE: u32 = 4 * 1024; + + struct RunningDvcIo { + write_tx: WinapiSignaledSender>, + read_rx: mpsc::Receiver>, + stop_event: Event, + thread: std::thread::JoinHandle>, + } + + fn run_dvc_io_on(server: usize, read_queue_capacity: usize) -> RunningDvcIo { + run_dvc_io_with_heartbeats(server, read_queue_capacity, true) + } + + fn run_dvc_io_with_heartbeats(server: usize, read_queue_capacity: usize, heartbeats_enabled: bool) -> RunningDvcIo { + run_dvc_io_with_timeout( + server, + read_queue_capacity, + heartbeats_enabled, + super::BACKLOG_STALL_TIMEOUT, + ) + } + + fn run_dvc_io_with_timeout( + server: usize, + read_queue_capacity: usize, + heartbeats_enabled: bool, + backlog_stall_timeout: Duration, + ) -> RunningDvcIo { + let (write_tx, write_rx) = winapi_signaled_mpsc_channel().expect("create the write queue"); + let (read_tx, read_rx) = mpsc::channel(read_queue_capacity); + let stop_event = Event::new_unnamed().expect("create the stop event"); + let thread_stop_event = stop_event.clone(); + + let thread = std::thread::spawn(move || { + let progress = DvcInitializationProgress::default(); + progress.set_heartbeats_enabled(heartbeats_enabled); + // SAFETY: The handle was opened by `dvc_pipe`, which gave up ownership of it. + let file = unsafe { Owned::new(HANDLE(server as *mut _)) }; + let mut io = DvcIo::new(file)?; + io.backlog_stall_timeout = backlog_stall_timeout; + io.serve(write_rx, &read_tx, &thread_stop_event, &progress) + }); + + RunningDvcIo { + write_tx, + read_rx, + stop_event, + thread, + } + } + + impl RunningDvcIo { + fn recv(&mut self) -> NowMessage<'static> { + let deadline = Instant::now() + TEST_TIMEOUT; + + loop { + if let Ok(message) = self.read_rx.try_recv() { + return message; + } + + assert!(Instant::now() < deadline, "no message reached the processing queue"); + std::thread::sleep(Duration::from_millis(5)); + } + } + + fn join(self) -> anyhow::Result<()> { + let deadline = Instant::now() + TEST_TIMEOUT; + + while !self.thread.is_finished() { + assert!(Instant::now() < deadline, "the DVC IO loop did not stop"); + std::thread::sleep(Duration::from_millis(5)); + } + + self.thread.join().expect("the DVC IO loop panicked") + } + + fn stop(self) -> anyhow::Result<()> { + self.stop_event.set().expect("set the stop event"); + self.join() + } + } + + /// Writes one channel read as the RDP stack delivers it: an 8-byte `CHANNEL_PDU_HEADER`, then the data. + fn write_pdu(client: &Owned, data: &[u8]) { + const CHANNEL_FLAG_FIRST_AND_LAST: u32 = 0x3; + + let mut pdu = u32::try_from(data.len()) + .expect("test data fits") + .to_le_bytes() + .to_vec(); + pdu.extend_from_slice(&CHANNEL_FLAG_FIRST_AND_LAST.to_le_bytes()); + pdu.extend_from_slice(data); + + write_raw(client, &pdu); + } + + fn write_raw(client: &Owned, bytes: &[u8]) { + let mut written = 0; + // SAFETY: The handle and buffer are valid for the duration of the synchronous call. + unsafe { WriteFile(**client, Some(bytes), Some(&mut written), None) }.expect("write to the test pipe"); + assert_eq!(written as usize, bytes.len()); + } + + fn read_exact(client: &Owned, len: usize) -> Vec { + let mut bytes = vec![0; len]; + let mut filled = 0; + + while filled < len { + let mut read = 0; + // SAFETY: The handle and buffer are valid for the duration of the synchronous call. + let result = unsafe { ReadFile(**client, Some(&mut bytes[filled..]), Some(&mut read), None) }; + + // A message larger than the buffer is read in parts, each reported as `ERROR_MORE_DATA`. + if let Err(error) = result { + assert_eq!(error.code(), ERROR_MORE_DATA.to_hresult(), "read from the test pipe"); + } + + filled += read as usize; + } + + bytes + } + + /// Waits until the client has bytes to read, so that a following [`read_exact`] cannot block forever. + fn wait_for_outbound(client: &Owned, failure: &str) { + let deadline = Instant::now() + TEST_TIMEOUT; + + loop { + let mut available = 0; + // SAFETY: The handle is valid, and only the byte count is requested. + unsafe { PeekNamedPipe(**client, None, 0, None, Some(&mut available), None) }.expect("peek the test pipe"); + + if available > 0 { + return; + } + + assert!(Instant::now() < deadline, "{failure}"); + std::thread::sleep(Duration::from_millis(5)); + } + } + + fn encoded(message: &NowMessage<'_>) -> Vec { + encode_vec(message).expect("encode a NOW message") + } + + fn exec_data(session_id: u32, data: Vec) -> NowMessage<'static> { + NowExecDataMsg::new(session_id, NowExecDataStreamKind::Stdin, false, data) + .expect("build an exec data message") + .into() + } + + // oracle: mutant a read shorter than the channel header ends the loop as a close, or loses its length + #[test] + fn read_shorter_than_the_channel_header_is_an_error_with_its_length() { + let (server, client) = dvc_pipe(); + let running = run_dvc_io_on(server, 1); + + write_raw(&client, &[0; 4]); + + let error = running.join().expect_err("a 4-byte read cannot hold a channel header"); + let message = format!("{error:#}"); + assert!(message.contains("4 bytes"), "unexpected error: {message}"); + } + + /// An empty channel PDU carries no message and must not end the loop, so a later message still arrives. + // oracle: mutant an empty channel PDU ends the loop or drops the next message + #[test] + fn empty_channel_pdu_is_skipped() { + let (server, client) = dvc_pipe(); + let mut running = run_dvc_io_on(server, 1); + let message = exec_data(1, b"after".to_vec()); + + write_pdu(&client, &[]); + write_pdu(&client, &encoded(&message)); + + assert_eq!(encoded(&running.recv()), encoded(&message)); + running.stop().expect("the loop stops cleanly"); + } + + /// A client that does not read our output must not stop us from reading its input; with a synchronous write the + /// loop would block until the client reads. + // oracle: mutant a pending write blocks reads from the client + #[test] + fn reads_continue_while_a_large_write_is_pending() { + let (server, client) = dvc_pipe(); + let mut running = run_dvc_io_on(server, 1); + + let outbound = exec_data(1, vec![0xA5; 1024 * 1024]); + let outbound_bytes = encoded(&outbound); + running.write_tx.try_send(outbound).expect("queue the outbound message"); + + // The write is in flight once its first bytes are visible to the client. + wait_for_outbound(&client, "the outbound write never started"); + + let inbound = exec_data(2, b"stdin".to_vec()); + write_pdu(&client, &encoded(&inbound)); + assert_eq!(encoded(&running.recv()), encoded(&inbound)); + + // The outbound message arrives whole, as a single pipe message. + assert_eq!(read_exact(&client, outbound_bytes.len()), outbound_bytes); + + running.stop().expect("the loop stops cleanly"); + } + + /// A full processing queue must hold messages back rather than drop them or end the channel: the client matches + /// responses by send order, so a gap would desynchronize it. + // oracle: mutant a full processing queue drops, reorders, or ends on held-back messages + #[test] + fn full_processing_queue_holds_messages_back_in_order() { + let (server, client) = dvc_pipe(); + let mut running = run_dvc_io_on(server, 1); + let messages: Vec<_> = (1..=3).map(|session_id| exec_data(session_id, vec![0; 16])).collect(); + + for message in &messages { + write_pdu(&client, &encoded(message)); + } + + // Give the loop time to read past the queue's capacity before anything is received. + std::thread::sleep(Duration::from_millis(200)); + assert!(!running.thread.is_finished(), "the loop ended on a full queue"); + + for message in &messages { + assert_eq!(encoded(&running.recv()), encoded(message)); + } + + running.stop().expect("the loop stops cleanly"); + } + + /// A client that goes away while the processing queue is full must still end the loop, so the channel is reopened. + /// + /// No read is pending while messages are held back, so only a write can observe the closed channel. + // oracle: mutant a closed channel goes unnoticed while messages are held back + #[test] + fn closed_channel_ends_the_loop_while_messages_are_held_back() { + let (server, client) = dvc_pipe(); + let running = run_dvc_io_on(server, 1); + + for session_id in 1..=3 { + write_pdu(&client, &encoded(&exec_data(session_id, vec![0; 16]))); + } + + std::thread::sleep(Duration::from_millis(200)); + drop(client); + + running.join().expect_err("a closed channel ends the loop"); + } + + /// Probing a held-back channel must not disturb it: the client gets heartbeats, and every message still arrives. + // oracle: mutant probing a held-back channel sends no heartbeat, or disturbs held-back messages + #[test] + fn held_back_channel_is_probed_with_heartbeats() { + let (server, client) = dvc_pipe(); + let mut running = run_dvc_io_on(server, 1); + let messages: Vec<_> = (1..=3).map(|session_id| exec_data(session_id, vec![0; 16])).collect(); - use super::{DVC_RETRY_DELAYS, DvcInitializationProgress, dvc_retry_delay, forward_to_processor}; + for message in &messages { + write_pdu(&client, &encoded(message)); + } + + let heartbeat = encoded(&NowChannelHeartbeatMsg::default().into()); + wait_for_outbound(&client, "the held-back channel was never probed"); + assert_eq!(read_exact(&client, heartbeat.len()), heartbeat); + + for message in &messages { + assert_eq!(encoded(&running.recv()), encoded(message)); + } + + running.stop().expect("the loop stops cleanly"); + } + + /// A client that disabled heartbeats must never be sent one, even to probe a held-back channel. + // oracle: mutant a client that disabled heartbeats is sent probe heartbeats + #[test] + fn held_back_channel_is_not_probed_when_heartbeats_are_disabled() { + let (server, client) = dvc_pipe(); + let mut running = run_dvc_io_with_heartbeats(server, 1, false); + let messages: Vec<_> = (1..=3).map(|session_id| exec_data(session_id, vec![0; 16])).collect(); + + for message in &messages { + write_pdu(&client, &encoded(message)); + } + + // Twice the probe interval, so a probe would have been written by now. + std::thread::sleep(super::BACKLOG_PROBE_INTERVAL * 2); + + let mut available = 0; + // SAFETY: The handle is valid, and only the byte count is requested. + unsafe { PeekNamedPipe(*client, None, 0, None, Some(&mut available), None) }.expect("peek the test pipe"); + assert_eq!(available, 0, "a client that disabled heartbeats was sent one"); + + for message in &messages { + assert_eq!(encoded(&running.recv()), encoded(message)); + } + + running.stop().expect("the loop stops cleanly"); + } - /// A full input queue must end the IO loop instead of dropping the message: the processing - /// task is still alive, so a silent drop would desynchronize the client's order-based - /// response matching with nothing to signal it. + /// The disabled-heartbeat backlog timeout remains at the chosen 30-second recovery limit. + // oracle: mutant the backlog timeout changes away from 30 seconds. #[test] - fn full_input_queue_is_an_error_not_a_drop() { - let (read_tx, mut read_rx) = mpsc::channel::>(1); - let heartbeat = || NowMessage::from(NowChannelHeartbeatMsg::default()); + fn disabled_heartbeat_backlog_timeout_is_thirty_seconds() { + assert_eq!(super::BACKLOG_STALL_TIMEOUT, Duration::from_secs(30)); + } - forward_to_processor(&read_tx, heartbeat()).expect("first message fits the queue"); + /// A blocked stdin sender and full processing queue cannot hide a closed DVC when heartbeats are disabled. + // oracle: mutant disabled-heartbeat backpressure can stall indefinitely after peer closure. + #[tokio::test(flavor = "multi_thread")] + async fn closed_channel_with_disabled_heartbeats_ends_saturated_backlog() { + let (server, client) = dvc_pipe(); + let running = run_dvc_io_with_timeout(server, 1, false, Duration::from_secs(1)); + let RunningDvcIo { + write_tx, + mut read_rx, + stop_event: _, + thread, + } = running; + let (stdin_tx, _stdin_rx) = mpsc::channel(1); + stdin_tx.try_send(()).expect("fill the stdin queue"); + let (blocked_tx, blocked_rx) = tokio::sync::oneshot::channel(); + let processing = tokio::spawn(async move { + let message = read_rx.recv().await.expect("first DVC message reaches processing"); + blocked_tx.send(message).expect("test is waiting for the first message"); + stdin_tx.send(()).await.expect("stdin queue remains open"); + }); + + write_pdu(&client, &encoded(&exec_data(1, vec![0; 16]))); + blocked_rx.await.expect("processing consumes the first message"); + write_pdu(&client, &encoded(&exec_data(2, vec![0; 16]))); + write_pdu(&client, &encoded(&exec_data(3, vec![0; 16]))); + + tokio::time::sleep(Duration::from_millis(100)).await; + drop(client); + + let deadline = Instant::now() + TEST_TIMEOUT; + while !thread.is_finished() { + assert!(Instant::now() < deadline, "the stalled saturated DVC did not end"); + std::thread::sleep(Duration::from_millis(5)); + } - let error = forward_to_processor(&read_tx, heartbeat()).expect_err("second message overflows the queue"); - assert!(error.to_string().contains("overflow"), "unexpected error: {error}"); + let error = thread + .join() + .expect("the DVC IO loop does not panic") + .expect_err("a stalled saturated channel ends with an error"); + assert!( + format!("{error:#}").contains("remained backpressured"), + "unexpected error: {error:#}" + ); + processing.abort(); + let _ = processing.await; + drop(write_tx); + } + + /// Stopping with a read and a write in flight must cancel both and return. + // oracle: mutant stopping with IO in flight hangs or leaves the IO pending + #[test] + fn stop_cancels_pending_io() { + let (server, _client) = dvc_pipe(); + let running = run_dvc_io_on(server, 1); - // The queued message is intact, and no further message was forwarded. - assert!(read_rx.try_recv().is_ok()); - assert!(read_rx.try_recv().is_err()); + running + .write_tx + .try_send(exec_data(1, vec![0; 1024 * 1024])) + .expect("queue the outbound message"); + std::thread::sleep(Duration::from_millis(100)); + + running.stop().expect("the loop stops cleanly"); } + // oracle: mutant a closed processing queue ends the loop without an error #[test] - fn closed_input_queue_is_an_error() { - let (read_tx, read_rx) = mpsc::channel::>(1); + fn closed_processing_queue_is_an_error() { + let (server, client) = dvc_pipe(); + let running = run_dvc_io_on(server, 1); + let RunningDvcIo { + write_tx, + read_rx, + stop_event, + thread, + } = running; drop(read_rx); - forward_to_processor(&read_tx, NowChannelHeartbeatMsg::default().into()).expect_err("closed queue is reported"); + write_pdu(&client, &encoded(&exec_data(1, Vec::new()))); + + let running = RunningDvcIo { + write_tx, + read_rx: mpsc::channel(1).1, + stop_event, + thread, + }; + running.join().expect_err("a closed processing queue ends the loop"); } #[test] diff --git a/devolutions-session/src/dvc/process.rs b/devolutions-session/src/dvc/process.rs index 95efb007e1..de3eeebf45 100644 --- a/devolutions-session/src/dvc/process.rs +++ b/devolutions-session/src/dvc/process.rs @@ -1,27 +1,38 @@ use std::collections::HashMap; +use std::os::windows::io::AsRawHandle as _; +use std::thread::JoinHandle; +use std::time::{Duration, Instant}; use now_proto_pdu::{NowExecDataStreamKind, NowStatusError}; -use tokio::sync::mpsc::Sender; -use tracing::{error, info, trace}; +use tokio::sync::mpsc::{Receiver, Sender}; +use tracing::{debug, error, info, trace, warn}; use win_api_wrappers::event::Event; use win_api_wrappers::handle::Handle; use win_api_wrappers::process::{Process, post_message_for_pid}; use win_api_wrappers::security::attributes::SecurityAttributesInit; +use win_api_wrappers::thread::Thread; use win_api_wrappers::utils::{Pipe, WideString}; use windows::Win32::Foundation::{ - CloseHandle, ERROR_BROKEN_PIPE, ERROR_HANDLE_EOF, GetLastError, LPARAM, WAIT_EVENT, WAIT_OBJECT_0, WPARAM, + ERROR_BROKEN_PIPE, ERROR_HANDLE_EOF, ERROR_NOT_FOUND, GetLastError, HANDLE, LPARAM, WAIT_EVENT, WAIT_OBJECT_0, + WPARAM, }; use windows::Win32::Storage::FileSystem::{ReadFile, WriteFile}; -use windows::Win32::System::IO::{GetOverlappedResult, OVERLAPPED}; +use windows::Win32::System::IO::{CancelSynchronousIo, GetOverlappedResult, OVERLAPPED}; +use windows::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOBOBJECT_BASIC_LIMIT_INFORMATION, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, SetInformationJobObject, +}; use windows::Win32::System::Threading::{ - CREATE_NEW_CONSOLE, CREATE_NEW_PROCESS_GROUP, CREATE_UNICODE_ENVIRONMENT, CreateProcessW, INFINITE, - NORMAL_PRIORITY_CLASS, PROCESS_INFORMATION, STARTF_USESHOWWINDOW, STARTF_USESTDHANDLES, STARTUPINFOW, - WaitForMultipleObjects, + CREATE_NEW_CONSOLE, CREATE_NEW_PROCESS_GROUP, CREATE_SUSPENDED, CREATE_UNICODE_ENVIRONMENT, CreateProcessW, + INFINITE, NORMAL_PRIORITY_CLASS, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION, STARTF_USESHOWWINDOW, + STARTF_USESTDHANDLES, STARTUPINFOW, WaitForMultipleObjects, }; use windows::Win32::UI::WindowsAndMessaging::{SW_HIDE, WM_QUIT}; use windows::core::PCWSTR; -use crate::dvc::channel::{WinapiSignaledReceiver, WinapiSignaledSender, winapi_signaled_mpsc_channel}; +use crate::dvc::channel::{ + WinapiSignaledReceiver, WinapiSignaledSender, bounded_mpsc_channel, winapi_signaled_mpsc_channel, +}; use crate::dvc::encoding::{DataEncoding, InputEncoder, OutputDecoder}; use crate::dvc::env::make_environment_block; use crate::dvc::fs::TmpFileGuard; @@ -62,15 +73,17 @@ impl From> for pub enum ProcessIoInputEvent { AbortExecution(u32), CancelExecution, - DataIn { - data: Vec, - /// If last is set, then stdin pipe should be closed. Any consecutive DataIn messages - /// will be ignored. - last: bool, - }, TerminateIo, } +/// Input for the process stdin. +#[derive(Debug)] +pub struct StdinChunk { + data: Vec, + /// Closes stdin after this chunk; later chunks are ignored. + last: bool, +} + /// Message, sent from Process IO thread to task to finalize process execution. #[derive(Debug)] pub enum ServerChannelEvent { @@ -195,9 +208,6 @@ impl WinApiProcessCtx { // wait for process to exit continue; } - ProcessIoInputEvent::DataIn { .. } => { - // DataIn messages ignored. - } ProcessIoInputEvent::TerminateIo => { info!(session_id, "Terminating IO loop"); @@ -226,10 +236,9 @@ impl WinApiProcessCtx { info!(session_id, ?self.encoding, "Process IO redirection loop has started"); - // Encoding transcoders for stdin/stdout/stderr. + // Encoding transcoders for stdout/stderr. let mut stdout_decoder = OutputDecoder::new(self.encoding); let mut stderr_decoder = OutputDecoder::new(self.encoding); - let mut stdin_encoder = InputEncoder::new(self.encoding); // Events for overlapped IO let stdout_read_event = Event::new_unnamed()?; @@ -325,61 +334,6 @@ impl WinApiProcessCtx { // wait for process to exit continue; } - ProcessIoInputEvent::DataIn { data, last } => { - trace!(session_id, "Received data to write to stdin pipe"); - - let pipe_handle = match self.stdin_write_pipe.as_ref() { - Some(pipe) => pipe.handle.raw(), - None => { - // Ignore data, as stdin pipe was already closed. - continue; - } - }; - - // Transcode UTF-8 input to the process encoding. - let encoded_data = stdin_encoder.encode(&data); - - if !encoded_data.is_empty() { - let mut bytes_written: u32 = 0; - - // Send data to stdin pipe in a blocking manner. - // SAFETY: pipe and input buffer are ensured to be valid via - // borrowing rules. - unsafe { - WriteFile( - pipe_handle, - Some(&*encoded_data), - Some(&mut bytes_written as *mut _), - None, - ) - }?; - } - - if last { - // Flush any remaining encoder state before closing. - let flushed = stdin_encoder.flush(); - if !flushed.is_empty() { - let mut bytes_written: u32 = 0; - // SAFETY: pipe and input buffer are ensured to be valid via - // borrowing rules. - unsafe { - WriteFile( - self.stdin_write_pipe - .as_ref() - .expect("BUG: stdin pipe closed unexpectedly") - .handle - .raw(), - Some(&flushed), - Some(&mut bytes_written as *mut _), - None, - ) - }?; - } - - // Close stdin pipe. - self.stdin_write_pipe = None; - } - } ProcessIoInputEvent::TerminateIo => { info!(session_id, "Terminating IO loop"); @@ -554,6 +508,106 @@ impl WinApiProcessCtx { } } +/// Writes stdin chunks to the process until the last one, then closes its stdin. +/// +/// Runs on a thread of its own because the write blocks until the process reads: in the IO loop, it would stop +/// stdout and stderr from being drained, and a process blocked on full output pipes never reads its input. +/// After a failed write, later input is discarded: the process no longer reads stdin, and its exit is reported by the +/// IO loop. +fn write_stdin(session_id: u32, pipe: Pipe, mut encoder: InputEncoder, mut stdin_rx: Receiver) { + while let Some(StdinChunk { data, last }) = stdin_rx.blocking_recv() { + let mut bytes = encoder.encode(&data).into_owned(); + + if last { + bytes.extend(encoder.flush()); + } + + if let Err(error) = write_all(&pipe, &bytes) { + warn!(session_id, %error, "Failed to write to the process stdin; discarding further input"); + return; + } + + if last { + break; + } + } + + trace!(session_id, "Closing the process stdin"); +} + +fn write_all(pipe: &Pipe, mut bytes: &[u8]) -> anyhow::Result<()> { + while !bytes.is_empty() { + let mut bytes_written: u32 = 0; + + // SAFETY: The pipe is a synchronous handle, and the buffer outlives the call. + unsafe { WriteFile(pipe.handle.raw(), Some(bytes), Some(&mut bytes_written), None) }?; + + let bytes_written = usize::try_from(bytes_written).expect("u32 fits into usize on Windows"); + + if bytes_written == 0 { + anyhow::bail!("stdin write made no progress"); + } + + bytes = bytes.get(bytes_written..).unwrap_or_default(); + } + + Ok(()) +} + +const STDIN_WRITER_STOP_TIMEOUT: Duration = Duration::from_secs(2); + +struct StdinWriter { + sender: Option>, + thread: Option>, +} + +impl StdinWriter { + fn stop(&mut self) -> anyhow::Result<()> { + drop(self.sender.take()); + + let Some(thread) = self.thread.take() else { + return Ok(()); + }; + + let deadline = Instant::now() + STDIN_WRITER_STOP_TIMEOUT; + while !thread.is_finished() && Instant::now() < deadline { + // SAFETY: The join handle owns this valid thread handle until it is dropped. + if let Err(error) = unsafe { CancelSynchronousIo(HANDLE(thread.as_raw_handle() as *mut _)) } + && error.code() != ERROR_NOT_FOUND.to_hresult() + { + warn!(%error, "Failed to cancel synchronous process stdin write"); + } + + std::thread::sleep(Duration::from_millis(10)); + } + + if thread.is_finished() { + if thread.join().is_err() { + anyhow::bail!("process stdin writer thread panicked"); + } + return Ok(()); + } + + // UNTESTED: A writer that does not stop remains detached with its pipe and write buffer alive, because + // freeing either while WriteFile is pending would be unsafe. When changing this timeout, check that cleanup + // remains bounded and the pending write's memory stays owned. + error!( + timeout = ?STDIN_WRITER_STOP_TIMEOUT, + "Process stdin writer did not stop after cancellation" + ); + drop(thread); + anyhow::bail!("process stdin writer did not stop within {STDIN_WRITER_STOP_TIMEOUT:?}") + } +} + +impl Drop for StdinWriter { + fn drop(&mut self) { + if let Err(error) = self.stop() { + error!(%error, "Failed to stop process stdin writer"); + } + } +} + /// Builder for process execution session. pub struct WinApiProcessBuilder { executable: String, @@ -618,6 +672,7 @@ impl WinApiProcessBuilder { mut self, session_id: u32, io_notification_tx: Option>, + job: Option<&ProcessJob>, detached: bool, ) -> Result, ExecError> { let command_line = format!("\"{}\" {}", self.executable, self.command_line) @@ -645,26 +700,27 @@ impl WinApiProcessBuilder { let encoding = self.encoding; let process_ctx = if io_redirection { - prepare_process_with_io_redirection(session_id, command_line, current_directory, self.env, encoding)? + prepare_process_with_io_redirection(session_id, command_line, current_directory, self.env, encoding, job)? } else { - prepare_process(session_id, command_line, current_directory, self.env, encoding)? + prepare_process(session_id, command_line, current_directory, self.env, encoding, job)? }; - // For detached mode, spawn a thread that waits for process exit and keeps temp files alive - if detached && !temp_files.is_empty() { - std::thread::spawn(move || { - let _temp_files = temp_files; + if detached { + if !temp_files.is_empty() { + std::thread::spawn(move || { + let _temp_files = temp_files; - // Wait for process to exit (indefinitely) - if let Err(error) = process_ctx.process.wait(None) { - error!(%error, session_id, "Failed to wait for detached process"); - return; - } + // Wait for process exit to keep its temporary files alive. + if let Err(error) = process_ctx.process.wait(None) { + error!(%error, session_id, "Failed to wait for detached process"); + return; + } - info!(session_id, "Detached process exited"); + info!(session_id, "Detached process exited"); - // Temp files will be cleaned up when this thread exits - }); + // Temporary files are cleaned up when this thread exits. + }); + } info!(session_id, "Detached process started successfully"); return Ok(None); @@ -673,6 +729,24 @@ impl WinApiProcessBuilder { // Create channel for `task` -> `Process IO thread` communication let (input_event_tx, input_event_rx) = winapi_signaled_mpsc_channel()?; + let mut process_ctx = process_ctx; + let stdin_writer = match process_ctx.stdin_write_pipe.take() { + Some(pipe) => { + let (stdin_tx, stdin_rx) = bounded_mpsc_channel()?; + let encoder = InputEncoder::new(encoding); + + let thread = std::thread::Builder::new() + .name(format!("process-{session_id}-stdin")) + .spawn(move || write_stdin(session_id, pipe, encoder, stdin_rx))?; + + Some(StdinWriter { + sender: Some(stdin_tx), + thread: Some(thread), + }) + } + None => None, + }; + let io_notification_tx = io_notification_tx.expect("BUG: io_notification_tx must be Some for non-detached mode"); @@ -695,19 +769,23 @@ impl WinApiProcessBuilder { Ok(Some(WinApiProcess { input_event_tx, + stdin_writer, join_handle, _temp_files: temp_files, })) } /// Starts process execution and spawns IO thread to redirect stdio to/from dvc. + /// + /// The process joins `job` before it runs, so it and every process it starts are killed when the job is dropped. pub fn run( self, session_id: u32, io_notification_tx: Sender, + job: &ProcessJob, ) -> Result { Ok(self - .run_impl(session_id, Some(io_notification_tx), false)? + .run_impl(session_id, Some(io_notification_tx), Some(job), false)? .expect("result should be non-optional when running in non-detached mode")) } @@ -716,17 +794,101 @@ impl WinApiProcessBuilder { /// Returns immediately after spawning. pub fn run_detached(self, session_id: u32) -> Result<(), ExecError> { // Result always empty and therefore ignored in detached mode. - self.run_impl(session_id, None, true)?; + self.run_impl(session_id, None, None, true)?; Ok(()) } } +/// Job object whose processes are killed when it is dropped, or when the host exits or crashes. +pub struct ProcessJob(Handle); + +impl ProcessJob { + pub fn new() -> anyhow::Result { + // SAFETY: Creating an unnamed job object with default security attributes has no preconditions. + // The handle is not inheritable, so no child process can keep the job open. + let job = unsafe { CreateJobObjectW(None, PCWSTR::null()) }?; + // SAFETY: `job` was just created and nothing else owns it. + let job = unsafe { Handle::new_owned(job) }?; + + let limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION { + BasicLimitInformation: JOBOBJECT_BASIC_LIMIT_INFORMATION { + LimitFlags: JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + ..Default::default() + }, + ..Default::default() + }; + + // SAFETY: `limits` is a JOBOBJECT_EXTENDED_LIMIT_INFORMATION of the size passed, and outlives the call. + unsafe { + SetInformationJobObject( + job.raw(), + JobObjectExtendedLimitInformation, + (&raw const limits).cast(), + u32::try_from(size_of_val(&limits)).expect("JOBOBJECT_EXTENDED_LIMIT_INFORMATION size fits into u32"), + ) + }?; + + Ok(Self(job)) + } +} + +fn creation_flags(job: Option<&ProcessJob>, has_environment_block: bool) -> PROCESS_CREATION_FLAGS { + let mut flags = NORMAL_PRIORITY_CLASS | CREATE_NEW_PROCESS_GROUP | CREATE_NEW_CONSOLE; + + if has_environment_block { + flags |= CREATE_UNICODE_ENVIRONMENT; + } + + // A process that ran before joining its job could start processes outside it. + if job.is_some() { + flags |= CREATE_SUSPENDED; + } + + flags +} + +/// Takes ownership of the handles returned by `CreateProcessW` for a process created with [`creation_flags`], and +/// lets the process run. +/// +/// With a job, the process joins it first, and is killed if it cannot. +fn start_created_process( + process_information: &PROCESS_INFORMATION, + job: Option<&ProcessJob>, +) -> Result { + // SAFETY: `CreateProcessW` returned these handles, and nothing else owns them. + let thread = Thread::from(unsafe { Handle::new_owned(process_information.hThread) }.map_err(anyhow::Error::from)?); + // SAFETY: Same as above. + let process = + Process::from(unsafe { Handle::new_owned(process_information.hProcess) }.map_err(anyhow::Error::from)?); + + if let Some(job) = job { + // SAFETY: Both handles are valid for the duration of the call. + let started = unsafe { AssignProcessToJobObject(job.0.raw(), process.handle.raw()) } + .map_err(anyhow::Error::from) + .and_then(|()| thread.resume()); + + if let Err(error) = started { + // UNTESTED: a process that cannot join its job must be killed, not left running where no channel can + // end it. No test makes AssignProcessToJobObject or ResumeThread fail; when changing this path, check + // that every failure after CreateProcessW still terminates the suspended process. + if let Err(error) = process.terminate(1) { + error!(%error, "Failed to kill a process that could not join its job"); + } + + return Err(error.context("failed to start the process in its job").into()); + } + } + + Ok(process) +} + fn prepare_process( session_id: u32, mut command_line: WideString, current_directory: WideString, env: HashMap, encoding: DataEncoding, + job: Option<&ProcessJob>, ) -> Result { let mut process_information = PROCESS_INFORMATION::default(); @@ -738,12 +900,6 @@ fn prepare_process( let environment_block = (!env.is_empty()).then(|| make_environment_block(env)).transpose()?; - let mut creation_flags = NORMAL_PRIORITY_CLASS | CREATE_NEW_PROCESS_GROUP | CREATE_NEW_CONSOLE; - - if environment_block.is_some() { - creation_flags |= CREATE_UNICODE_ENVIRONMENT; - } - // SAFETY: All parameters constructed above are valid and safe to use. unsafe { CreateProcessW( @@ -752,7 +908,7 @@ fn prepare_process( None, None, true, - creation_flags, + creation_flags(job, environment_block.is_some()), environment_block.as_ref().map(|block| block.as_ptr() as *const _), current_directory.as_pcwstr(), &startup_info as *const _, @@ -760,18 +916,7 @@ fn prepare_process( )?; } - // The thread handle returned by CreateProcessW is only needed if you want to manage or - // wait on the primary thread of the new process. Since we only need to manage the process - // itself and not its main thread, we close the thread handle immediately to avoid - // resource leaks. - - // SAFETY: FFI call with no outstanding precondition. - unsafe { CloseHandle(process_information.hThread) }?; - - let process_handle = Process::from( - // SAFETY: process_information is valid and contains valid process handle. - unsafe { Handle::new_owned(process_information.hProcess) }.map_err(anyhow::Error::from)?, - ); + let process_handle = start_created_process(&process_information, job)?; let pid = process_information.dwProcessId; @@ -792,6 +937,7 @@ fn prepare_process_with_io_redirection( current_directory: WideString, env: HashMap, encoding: DataEncoding, + job: Option<&ProcessJob>, ) -> Result { let mut process_information = PROCESS_INFORMATION::default(); @@ -822,11 +968,6 @@ fn prepare_process_with_io_redirection( let environment_block = (!env.is_empty()).then(|| make_environment_block(env)).transpose()?; - let mut creation_flags = NORMAL_PRIORITY_CLASS | CREATE_NEW_PROCESS_GROUP | CREATE_NEW_CONSOLE; - if environment_block.is_some() { - creation_flags |= CREATE_UNICODE_ENVIRONMENT; - } - // SAFETY: All parameters constructed above are valid and safe to use. unsafe { CreateProcessW( @@ -835,7 +976,7 @@ fn prepare_process_with_io_redirection( Some(security_attributes.as_ptr()), None, true, - creation_flags, + creation_flags(job, environment_block.is_some()), environment_block.as_ref().map(|block| block.as_ptr() as *const _), current_directory.as_pcwstr(), &startup_info as *const _, @@ -843,19 +984,13 @@ fn prepare_process_with_io_redirection( )?; } - // SAFETY: FFI call with no outstanding precondition. - unsafe { CloseHandle(process_information.hThread) }?; - // Handles were duplicated by CreateProcessW, so we can close them immediately. // Explicitly drop handles just for clarity. drop(stdout_write_pipe); drop(stderr_write_pipe); drop(stdin_read_pipe); - let process_handle = Process::from( - // SAFETY: process_information is valid and contains valid process handle. - unsafe { Handle::new_owned(process_information.hProcess) }.map_err(anyhow::Error::from)?, - ); + let process_handle = start_created_process(&process_information, job)?; let pid = process_information.dwProcessId; @@ -875,7 +1010,9 @@ fn prepare_process_with_io_redirection( /// Represents spawned process with IO redirection. pub struct WinApiProcess { input_event_tx: WinapiSignaledSender, - join_handle: std::thread::JoinHandle<()>, + /// `None` without IO redirection. + stdin_writer: Option, + join_handle: JoinHandle<()>, _temp_files: Vec, } @@ -898,10 +1035,27 @@ impl WinApiProcess { self.input_event_tx.send(ProcessIoInputEvent::CancelExecution).await } - pub async fn send_stdin(&self, data: Vec, last: bool) -> anyhow::Result<()> { - self.input_event_tx - .send(ProcessIoInputEvent::DataIn { data, last }) - .await + /// Returns a send of `data` to the process stdin that waits while the stdin queue is full. + /// + /// The send owns what it needs, so the caller can keep serving process output while it waits: the process may not + /// read its stdin until its output is drained. + /// Input for a process that no longer reads stdin is discarded. + pub fn send_stdin(&self, data: Vec, last: bool) -> impl Future + Send + 'static { + let stdin_tx = self + .stdin_writer + .as_ref() + .and_then(|writer| writer.sender.as_ref()) + .cloned(); + + async move { + let Some(stdin_tx) = stdin_tx else { + return; + }; + + if stdin_tx.send(StdinChunk { data, last }).await.is_err() { + debug!("Discarding input for a closed process stdin"); + } + } } pub async fn shutdown(&self) -> anyhow::Result<()> { @@ -913,3 +1067,100 @@ impl WinApiProcess { self.join_handle.is_finished() } } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + // oracle: regression #2156 detached Process without temporary files must not panic after launch. + #[test] + fn detached_process_without_temp_files_returns_success() { + WinApiProcessBuilder::new("cmd.exe") + .with_command_line("/d /c exit 0") + .run_detached(1) + .expect("detached launch returns successfully without an IO notification channel"); + } + + /// A process that fills its output before reading stdin must still get all of a large input. + /// + /// The input exceeds the stdin pipe buffer, so a stdin write that blocks the output drain deadlocks. + // oracle: spec CreatePipe §Remarks + #[tokio::test(flavor = "multi_thread")] + async fn large_stdin_after_output_completes() { + const SESSION_ID: u32 = 1; + const OUTPUT_LINES: u32 = 3000; + const INPUT_LINES: u32 = 5000; + + let (io_notification_tx, mut io_notification_rx) = tokio::sync::mpsc::channel(100); + let job = ProcessJob::new().expect("create job"); + let process = WinApiProcessBuilder::new("cmd.exe") + .with_command_line(&format!( + "/d /c (for /L %i in (1,1,{OUTPUT_LINES}) do @echo output-line-%i) & sort" + )) + .with_io_redirection(true) + .with_encoding(DataEncoding::Raw) + .run(SESSION_ID, io_notification_tx, &job) + .expect("start cmd.exe"); + + let input: String = (0..INPUT_LINES).map(|i| format!("input-line-{i:05}\r\n")).collect(); + let feed = tokio::spawn(process.send_stdin(input.into_bytes(), true)); + + let mut stdout = Vec::new(); + let exit_code = tokio::time::timeout(Duration::from_secs(60), async { + loop { + match io_notification_rx.recv().await.expect("notification channel open") { + ServerChannelEvent::SessionDataOut { + stream: NowExecDataStreamKind::Stdout, + data, + .. + } => stdout.extend(data), + ServerChannelEvent::SessionExited { exit_code, .. } => break exit_code, + ServerChannelEvent::SessionFailed { error, .. } => panic!("process failed: {error}"), + _ => {} + } + } + }) + .await + .expect("process exits instead of deadlocking on stdin"); + + feed.await.expect("stdin feed completes"); + + let stdout = String::from_utf8(stdout).expect("ASCII output"); + assert_eq!(exit_code, 0); + assert!(stdout.contains(&format!("output-line-{OUTPUT_LINES}"))); + assert!(stdout.contains(&format!("input-line-{:05}", INPUT_LINES - 1))); + assert_eq!( + stdout.matches("input-line-").count(), + usize::try_from(INPUT_LINES).expect("u32 fits into usize") + ); + } + + /// Dropping an attached execution must cancel a synchronous stdin write blocked by a child that does not read. + // oracle: mutant dropping an execution leaves its synchronous stdin writer blocked. + #[tokio::test(flavor = "multi_thread")] + async fn dropping_process_stops_a_blocked_stdin_writer() { + const SESSION_ID: u32 = 1; + + let (io_notification_tx, _io_notification_rx) = tokio::sync::mpsc::channel(100); + let job = ProcessJob::new().expect("create job"); + let process = WinApiProcessBuilder::new("cmd.exe") + .with_command_line("/d /c ping -n 600 127.0.0.1 > nul") + .with_io_redirection(true) + .run(SESSION_ID, io_notification_tx, &job) + .expect("start cmd.exe"); + + process.send_stdin(vec![b'x'; 1024 * 1024], false).await; + std::thread::sleep(Duration::from_millis(100)); + + let started_at = Instant::now(); + drop(process); + assert!( + started_at.elapsed() < STDIN_WRITER_STOP_TIMEOUT / 2, + "dropping the process did not stop the blocked stdin writer promptly" + ); + + drop(job); + } +} diff --git a/devolutions-session/src/dvc/task.rs b/devolutions-session/src/dvc/task.rs index f346602ac1..2a5b1618f4 100644 --- a/devolutions-session/src/dvc/task.rs +++ b/devolutions-session/src/dvc/task.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use anyhow::{Context, bail}; use async_trait::async_trait; use devolutions_agent_shared::temp_file::{BATCH_UTF8_PREAMBLE, POWERSHELL_UTF8_ENCODING_PREAMBLE, TmpFileGuard}; -use devolutions_gateway_task::{ChildTask, Task}; +use devolutions_gateway_task::{ChildTask, ShutdownSignal, Task}; use now_proto_pdu::ironrdp_core::IntoOwned; use now_proto_pdu::{ ComApartmentStateKind, NowChannelCapsetMsg, NowChannelCloseMsg, NowChannelHeartbeatMsg, NowChannelMessage, @@ -42,7 +42,7 @@ use windows::core::PCWSTR; use crate::dvc::channel::{WinapiSignaledSender, bounded_mpsc_channel, winapi_signaled_mpsc_channel}; use crate::dvc::encoding::DataEncoding; use crate::dvc::io::{DvcInitializationProgress, run_dvc_io}; -use crate::dvc::process::{ExecError, ServerChannelEvent, WinApiProcess, WinApiProcessBuilder}; +use crate::dvc::process::{ExecError, ProcessJob, ServerChannelEvent, WinApiProcess, WinApiProcessBuilder}; use crate::dvc::rdm::RdmMessageProcessor; use crate::dvc::window_monitor::{WindowMonitorConfig, run_window_monitor}; @@ -50,6 +50,14 @@ use crate::dvc::window_monitor::{WindowMonitorConfig, run_window_monitor}; const DEFAULT_HEARTBEAT_INTERVAL: core::time::Duration = core::time::Duration::from_secs(60); const HANDSHAKE_TIMEOUT: core::time::Duration = core::time::Duration::from_secs(5); +/// Delay before reopening a channel after the first failure; it doubles with each consecutive failed reopen. +const REOPEN_DELAY_FLOOR: core::time::Duration = core::time::Duration::from_secs(1); +const REOPEN_DELAY_CAP: core::time::Duration = core::time::Duration::from_secs(10); +/// Number of consecutive reopened channels that may fail to negotiate before the process gives up. +const MAX_CONSECUTIVE_FAILED_REOPENS: usize = 5; +/// Time the IO thread of a lost channel has to release it before the process gives up. +const IO_THREAD_STOP_TIMEOUT: core::time::Duration = core::time::Duration::from_secs(10); + const GENERIC_ERROR_CODE_ENCODING: u32 = 0x00000001; const GENERIC_ERROR_CODE_TOO_LONG_ERROR: u32 = 0x00000002; const GENERIC_ERROR_CODE_OTHER: u32 = 0xFFFFFFFF; @@ -63,62 +71,207 @@ impl Task for DvcIoTask { const NAME: &'static str = "DVC processing"; - async fn run(self, shutdown_signal: devolutions_gateway_task::ShutdownSignal) -> Self::Output { - let (write_tx, write_rx) = winapi_signaled_mpsc_channel()?; - let (read_tx, read_rx) = bounded_mpsc_channel()?; + async fn run(self, shutdown_signal: ShutdownSignal) -> Self::Output { + let channel_shutdown_signal = shutdown_signal.clone(); - // WinAPI event to terminate DVC IO thread. - let io_thread_shutdown_event = Event::new_unnamed()?; + serve_dvc(move || serve_channel(channel_shutdown_signal.clone()), shutdown_signal).await + } +} - let cloned_shutdown_event = io_thread_shutdown_event.clone(); +/// How serving one DVC channel ended. +#[derive(Debug)] +enum ChannelOutcome { + /// The channel was closed on purpose, by either peer or by shutdown, and must not be reopened. + Ended, + /// The channel failed before the client completed negotiation. + NotNegotiated(anyhow::Error), + /// The channel failed after the client completed negotiation. + Lost(anyhow::Error), +} - let init_progress = Arc::new(DvcInitializationProgress::default()); - let io_thread_progress = Arc::clone(&init_progress); +impl ChannelOutcome { + fn new(result: anyhow::Result<()>, init_progress: &DvcInitializationProgress) -> Self { + match result { + Ok(()) => Self::Ended, + Err(error) if init_progress.negotiated() => Self::Lost(error), + Err(error) => Self::NotNegotiated(error), + } + } +} - info!( - "Starting NowProto DVC v{}.{}", - NowProtoVersion::CURRENT.major, - NowProtoVersion::CURRENT.minor - ); +/// Serves DVC channels until one ends on purpose, reopening a channel lost after negotiation. +/// +/// A first channel that never negotiates is not reopened: the client likely has no NOW listener, and reopening would +/// keep a process alive that nobody can reach. +/// Once a client has negotiated, a failed channel is reopened after a delay that grows with each consecutive reopen +/// that fails to negotiate, until [`MAX_CONSECUTIVE_FAILED_REOPENS`] of them. +/// +/// Errors only when a negotiated client can no longer be served, which makes the process exit with a failure code +/// that the agent answers by relaunching it. +async fn serve_dvc(mut serve_channel: F, mut shutdown_signal: ShutdownSignal) -> anyhow::Result<()> +where + F: FnMut() -> Fut, + Fut: Future>, +{ + // `None` until a negotiated channel is lost. + let mut failed_reopens: Option = None; - // Spawning thread is relatively short operation, so it could be executed synchronously. - let io_thread = std::thread::spawn(move || { - let io_thread_result = run_dvc_io(write_rx, read_tx, cloned_shutdown_event, &io_thread_progress); + loop { + if let Some(failed_reopens) = failed_reopens { + select! { + () = tokio::time::sleep(reopen_delay(failed_reopens)) => {} + _ = shutdown_signal.wait() => return Ok(()), + } + } - if let Err(error) = io_thread_result { - error!(%error, "DVC IO thread failed"); + match serve_channel().await? { + ChannelOutcome::Ended => return Ok(()), + ChannelOutcome::Lost(error) => { + warn!( + error = format!("{error:#}"), + "Lost the negotiated DVC channel; reopening it" + ); + failed_reopens = Some(0); } - }); + ChannelOutcome::NotNegotiated(error) => { + let Some(previous_failures) = failed_reopens else { + // Not a failure: the client likely has no NOW listener, and the agent relaunches a process that + // exits with an error. + warn!( + error = format!("{error:#}"), + "The client did not negotiate the DVC channel" + ); + return Ok(()); + }; + + let failures = previous_failures + 1; - // Join thread some time in future. - tokio::task::spawn_blocking(move || { - if let Err(panic_payload) = io_thread.join() { - if let Some(s) = panic_payload.downcast_ref::<&dyn std::fmt::Debug>() { - error!("DVC IO thread panic: {:?}", s); + if failures < MAX_CONSECUTIVE_FAILED_REOPENS { + warn!( + error = format!("{error:#}"), + failures, "Failed to reopen the DVC channel" + ); + failed_reopens = Some(failures); } else { - error!("DVC IO thread join failed."); + return Err(error.context(format!("failed to reopen the DVC channel {failures} consecutive times"))); } } - }); + } + } +} - info!("Processing DVC messages..."); +fn reopen_delay(failed_reopens: usize) -> core::time::Duration { + let factor = u32::try_from(failed_reopens) + .ok() + .and_then(|exponent| 2u32.checked_pow(exponent)) + .unwrap_or(u32::MAX); - let process_result = process_messages(read_rx, write_tx, shutdown_signal, &init_progress).await; + REOPEN_DELAY_FLOOR.saturating_mul(factor).min(REOPEN_DELAY_CAP) +} - // Send shutdown signal to IO thread to release WTS channel resources. - info!("Shutting down DVC IO thread"); - let _ = io_thread_shutdown_event.set(); +/// Opens one DVC channel and serves it until it ends. +/// +/// Errors only when the process cannot go on serving channels at all. +async fn serve_channel(mut shutdown_signal: ShutdownSignal) -> anyhow::Result { + let (write_tx, write_rx) = winapi_signaled_mpsc_channel()?; + let (read_tx, read_rx) = bounded_mpsc_channel()?; - process_result?; + // WinAPI event to terminate DVC IO thread. + let io_thread_shutdown_event = Event::new_unnamed()?; - Ok(()) + let cloned_shutdown_event = io_thread_shutdown_event.clone(); + + // Each channel gets its own progress, and therefore its own negotiation deadline. + let init_progress = Arc::new(DvcInitializationProgress::default()); + let io_thread_progress = Arc::clone(&init_progress); + let (io_loop_ended_tx, io_loop_ended_rx) = tokio::sync::oneshot::channel(); + + info!( + "Starting NowProto DVC v{}.{}", + NowProtoVersion::CURRENT.major, + NowProtoVersion::CURRENT.minor + ); + + // Spawning thread is relatively short operation, so it could be executed synchronously. + let io_thread = std::thread::spawn(move || { + let io_thread_result = run_dvc_io( + write_rx, + read_tx, + cloned_shutdown_event, + &io_thread_progress, + io_loop_ended_tx, + ); + + if let Err(error) = io_thread_result { + error!(%error, "DVC IO thread failed"); + } + }); + + // Joined on a plain thread rather than the blocking pool, because the runtime waits for its blocking tasks when it + // shuts down, so an IO thread that never stops would keep the process alive. + let (io_thread_stopped_tx, io_thread_stopped_rx) = tokio::sync::oneshot::channel::<()>(); + std::thread::spawn(move || { + if let Err(panic_payload) = io_thread.join() { + if let Some(s) = panic_payload.downcast_ref::<&dyn std::fmt::Debug>() { + error!("DVC IO thread panic: {:?}", s); + } else { + error!("DVC IO thread join failed."); + } + } + + let _ = io_thread_stopped_tx.send(()); + }); + + info!("Processing DVC messages..."); + + let process_result = process_until_io_ends( + process_messages(read_rx, write_tx, shutdown_signal.clone(), &init_progress), + io_loop_ended_rx, + ) + .await; + + // Send shutdown signal to IO thread to release WTS channel resources. + info!("Shutting down DVC IO thread"); + let _ = io_thread_shutdown_event.set(); + + let outcome = ChannelOutcome::new(process_result, &init_progress); + + // UNTESTED: a failed channel may be reopened, so its IO thread must release the WTS handle before the next open, + // and a thread that does not stop must end the process, because the agent does not replace a process that is + // still running. + // An ended channel is never reopened, so the process exits without depending on the thread. + // The thread waits for its cancelled overlapped IO to finish before it stops; when changing the IO thread or its + // join, check that the timeout still bounds that wait and that nothing on the runtime still waits for the thread. + if let ChannelOutcome::Lost(error) | ChannelOutcome::NotNegotiated(error) = &outcome { + select! { + stopped = tokio::time::timeout(IO_THREAD_STOP_TIMEOUT, io_thread_stopped_rx) => { + if stopped.is_err() { + error!(error = format!("{error:#}"), "DVC channel failed"); + bail!("DVC IO thread did not stop within {IO_THREAD_STOP_TIMEOUT:?}"); + } + } + _ = shutdown_signal.wait() => return Ok(ChannelOutcome::Ended), + } + } + + Ok(outcome) +} + +async fn process_until_io_ends( + processing: impl Future>, + io_loop_ended: tokio::sync::oneshot::Receiver<()>, +) -> anyhow::Result<()> { + select! { + biased; + result = processing => result, + _ = io_loop_ended => bail!("DVC IO loop ended"), } } async fn process_messages( mut read_rx: Receiver, dvc_tx: WinapiSignaledSender, - mut shutdown_signal: devolutions_gateway_task::ShutdownSignal, + mut shutdown_signal: ShutdownSignal, init_progress: &DvcInitializationProgress, ) -> anyhow::Result<()> { let (io_notification_tx, mut task_rx) = mpsc::channel(100); @@ -189,6 +342,7 @@ async fn process_messages( reason = "the task handle is a drop guard; it is assigned, never inspected" )] let mut _heartbeat = spawn_heartbeat(dvc_tx.clone(), downgraded_caps.heartbeat_interval()); + init_progress.set_heartbeats_enabled(heartbeats_enabled(downgraded_caps.heartbeat_interval())); // Reply with the negotiated capset so the client sees the capabilities and heartbeat interval // the agent applies, not the server default that omits the interval. @@ -201,15 +355,31 @@ async fn process_messages( pwsh_executable.clone(), ); + init_progress.mark_negotiated(); info!("DVC negotiation completed"); + // Process output keeps being served while stdin waits, because the process may not read stdin until its output + // is drained. + let mut pending_stdin: Option<(u32, StdinSend)> = None; + loop { select! { - read_result = read_rx.recv() => { + () = async { + match pending_stdin.as_mut() { + Some((_, send)) => send.await, + None => std::future::pending().await, + } + } => { + pending_stdin = None; + } + read_result = read_rx.recv(), if pending_stdin.is_none() => { match read_result { Some(message) => { match processor.process_message(message).await { Ok(ProcessMessageAction::Continue) => {} + Ok(ProcessMessageAction::SendStdin { session_id, send }) => { + pending_stdin = Some((session_id, send)); + } Ok(ProcessMessageAction::Shutdown) => { info!("Received channel shutdown message..."); return Ok(()); @@ -222,9 +392,9 @@ async fn process_messages( let server_caps = default_server_caps(pwsh_executable.is_some()); let downgraded_caps = client_caps.downgrade(&server_caps); _heartbeat = spawn_heartbeat(dvc_tx.clone(), downgraded_caps.heartbeat_interval()); + init_progress.set_heartbeats_enabled(heartbeats_enabled(downgraded_caps.heartbeat_interval())); - // Old exec sessions will be abandoned (IO loops will terminate - // on old processor `Drop`). + // Dropping the old processor ends its exec sessions and the processes they started. processor = MessageProcessor::new( downgraded_caps.clone(), dvc_tx.clone(), @@ -271,6 +441,7 @@ async fn process_messages( ServerChannelEvent::SessionExited { session_id, exit_code } => { info!(session_id, %exit_code, "Session exited"); processor.remove_session(session_id); + discard_pending_stdin(&mut pending_stdin, session_id); let message = NowExecResultMsg::new_success(session_id, exit_code); dvc_tx.send(message.into()).await?; @@ -278,6 +449,7 @@ async fn process_messages( ServerChannelEvent::SessionFailed { session_id, error } => { error!(session_id, %error, "Session error"); processor.remove_session(session_id); + discard_pending_stdin(&mut pending_stdin, session_id); handle_exec_error(&dvc_tx, session_id, error).await; } @@ -343,18 +515,24 @@ fn spawn_heartbeat( /// `NOW_CHANNEL_SET_HEARTBEAT`. /// Without a negotiated interval, [`DEFAULT_HEARTBEAT_INTERVAL`] applies. fn heartbeat_timer(negotiated_interval: Option) -> Option { - let period = negotiated_interval.unwrap_or(DEFAULT_HEARTBEAT_INTERVAL); - - if period.is_zero() { + if !heartbeats_enabled(negotiated_interval) { return None; } + let period = negotiated_interval.unwrap_or(DEFAULT_HEARTBEAT_INTERVAL); let mut timer = tokio::time::interval_at(Instant::now() + period, period); timer.set_missed_tick_behavior(MissedTickBehavior::Delay); Some(timer) } +/// Whether the server may send heartbeats for a negotiated interval. +/// +/// Only a zero interval disables them; without a negotiated interval, [`DEFAULT_HEARTBEAT_INTERVAL`] applies. +fn heartbeats_enabled(negotiated_interval: Option) -> bool { + !negotiated_interval.unwrap_or(DEFAULT_HEARTBEAT_INTERVAL).is_zero() +} + fn default_server_caps(pwsh_available: bool) -> NowChannelCapsetMsg { let mut exec_flags = NowExecCapsetFlags::STYLE_RUN | NowExecCapsetFlags::STYLE_PROCESS @@ -416,10 +594,27 @@ fn find_pwsh_executable() -> Option { String::from_utf16(&buffer[..len]).ok() } +type StdinSend = std::pin::Pin + Send>>; + +/// Drops a stdin send for a session that has ended, so its input no longer holds back other messages. +fn discard_pending_stdin(pending_stdin: &mut Option<(u32, StdinSend)>, ended_session_id: u32) { + if pending_stdin + .as_ref() + .is_some_and(|(session_id, _)| *session_id == ended_session_id) + { + *pending_stdin = None; + } +} + enum ProcessMessageAction { Continue, Shutdown, Restart(NowChannelCapsetMsg), + /// Stdin for a session; later messages wait for it so input stays ordered and a full stdin queue pushes back. + SendStdin { + session_id: u32, + send: StdinSend, + }, } struct MessageProcessor { @@ -427,7 +622,15 @@ struct MessageProcessor { io_notification_tx: Sender, #[allow(dead_code)] // Not yet used. capabilities: NowChannelCapsetMsg, + /// Attached execution sessions of this channel. + /// + /// Ending the channel or renegotiating capabilities drops this processor, which stops their IO and kills them with + /// `process_job`; no later channel or negotiation could reach them, because each starts with no sessions. sessions: HashMap, + /// Job of the attached executions and the processes they start; created by the first attached execution. + /// + /// Processes left running by an attached execution that already exited are killed with it too. + process_job: Option, rdm_handler: RdmMessageProcessor, /// Resolved pwsh executable path; `None` when pwsh is not installed. pwsh_executable: Option, @@ -450,6 +653,7 @@ impl MessageProcessor { io_notification_tx, capabilities, sessions: HashMap::new(), + process_job: None, rdm_handler, pwsh_executable, window_monitor_shutdown_tx: None, @@ -457,6 +661,13 @@ impl MessageProcessor { } } + fn process_job(&mut self) -> Result<&ProcessJob, ExecError> { + Ok(match &mut self.process_job { + Some(job) => job, + slot @ None => slot.insert(ProcessJob::new()?), + }) + } + async fn ensure_session_id_free(&self, session_id: u32) -> Result<(), ExecError> { if self.sessions.contains_key(&session_id) { warn!(session_id, "Session ID is in use"); @@ -590,7 +801,12 @@ impl MessageProcessor { } }; - process.send_stdin(data_msg.data().to_vec(), data_msg.is_last()).await?; + let send = process.send_stdin(data_msg.data().to_vec(), data_msg.is_last()); + + return Ok(ProcessMessageAction::SendStdin { + session_id, + send: Box::pin(send), + }); } NowMessage::Session(NowSessionMessage::MsgBoxReq(request)) => { let tx = self.dvc_tx.clone(); @@ -716,14 +932,16 @@ impl MessageProcessor { info!(session_id, "Executing ShellExecuteW"); - // SAFETY: All buffers are valid, therefore `ShellExecuteW` is safe to call. + // UNTESTED: Keep the directory buffer alive through ShellExecuteW; consuming it creates a dangling pointer. + // Check the borrow when changing argument conversion; a successful launch cannot prove pointer validity. + // SAFETY: All buffers are borrowed from locals that remain alive through the call. let hinstance = unsafe { ShellExecuteW( None, operation.as_pcwstr(), command.as_pcwstr(), parameters.as_pcwstr(), - directory.map(|d| d.as_pcwstr()).unwrap_or(PCWSTR::null()), + directory.as_ref().map_or_else(PCWSTR::null, WideString::as_pcwstr), // Activate and show window SW_RESTORE, ) @@ -774,9 +992,11 @@ impl MessageProcessor { return Ok(()); } - let process = run_process - .with_io_redirection(exec_msg.is_with_io_redirection()) - .run(exec_msg.session_id(), self.io_notification_tx.clone())?; + let process = run_process.with_io_redirection(exec_msg.is_with_io_redirection()).run( + exec_msg.session_id(), + self.io_notification_tx.clone(), + self.process_job()?, + )?; self.sessions.insert(exec_msg.session_id(), process); @@ -821,9 +1041,11 @@ impl MessageProcessor { return Ok(()); } - let process = run_batch - .with_io_redirection(batch_msg.is_with_io_redirection()) - .run(batch_msg.session_id(), self.io_notification_tx.clone())?; + let process = run_batch.with_io_redirection(batch_msg.is_with_io_redirection()).run( + batch_msg.session_id(), + self.io_notification_tx.clone(), + self.process_job()?, + )?; self.sessions.insert(batch_msg.session_id(), process); @@ -883,7 +1105,11 @@ impl MessageProcessor { let process = run_process .with_io_redirection(winps_msg.is_with_io_redirection()) - .run(winps_msg.session_id(), self.io_notification_tx.clone())?; + .run( + winps_msg.session_id(), + self.io_notification_tx.clone(), + self.process_job()?, + )?; self.sessions.insert(winps_msg.session_id(), process); @@ -948,9 +1174,11 @@ impl MessageProcessor { return Ok(()); } - let process = run_process - .with_io_redirection(pwsh_msg.is_with_io_redirection()) - .run(pwsh_msg.session_id(), self.io_notification_tx.clone())?; + let process = run_process.with_io_redirection(pwsh_msg.is_with_io_redirection()).run( + pwsh_msg.session_id(), + self.io_notification_tx.clone(), + self.process_job()?, + )?; self.sessions.insert(pwsh_msg.session_id(), process); @@ -1324,19 +1552,432 @@ unsafe extern "C" { #[cfg(test)] mod tests { + use core::task::Poll; use core::time::Duration; + use std::collections::VecDeque; - use devolutions_gateway_task::ShutdownHandle; + use anyhow::anyhow; + use devolutions_gateway_task::{ShutdownHandle, ShutdownSignal}; use futures::poll; use now_proto_pdu::{ - NowChannelCapsetMsg, NowChannelCloseMsg, NowChannelMessage, NowExecCapsetFlags, NowMessage, - NowSessionCapsetFlags, NowSystemCapsetFlags, OwnedNowMessage, + NowChannelCapsetMsg, NowChannelCloseMsg, NowChannelMessage, NowExecCapsetFlags, NowExecMessage, + NowExecProcessMsg, NowExecResultMsg, NowMessage, NowSessionCapsetFlags, NowSystemCapsetFlags, OwnedNowMessage, + }; + use tokio::time::Instant; + use win_api_wrappers::process::Process; + use windows::Win32::Foundation::{WAIT_OBJECT_0, WAIT_TIMEOUT}; + use windows::Win32::System::Threading::{PROCESS_SYNCHRONIZE, PROCESS_TERMINATE}; + + use super::{ + ChannelOutcome, DEFAULT_HEARTBEAT_INTERVAL, MAX_CONSECUTIVE_FAILED_REOPENS, REOPEN_DELAY_CAP, + REOPEN_DELAY_FLOOR, heartbeat_timer, process_messages, process_until_io_ends, serve_dvc, spawn_heartbeat, }; - - use super::{DEFAULT_HEARTBEAT_INTERVAL, heartbeat_timer, process_messages, spawn_heartbeat}; use crate::dvc::channel::{bounded_mpsc_channel, winapi_signaled_mpsc_channel}; use crate::dvc::io::DvcInitializationProgress; + // oracle: regression #2156 detached Process must leave the channel alive for subsequent execution. + #[tokio::test] + async fn detached_process_keeps_channel_available_for_execution() { + let (write_tx, mut write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + read_tx + .send(NowChannelCapsetMsg::default().into()) + .await + .expect("queue client capset"); + read_tx + .send( + NowExecProcessMsg::new(1, "cmd.exe") + .expect("valid executable") + .with_parameters("/d /c exit 0") + .expect("valid parameters") + .with_detached() + .into(), + ) + .await + .expect("queue detached process"); + + let mut processing = Box::pin(process_messages(read_rx, write_tx, shutdown_signal, &init_progress)); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + tokio::select! { + result = &mut processing => panic!("processing ended before fresh execution completed: {result:?}"), + () = tokio::time::sleep(Duration::from_millis(10)) => {} + } + + while let Ok(message) = write_rx.try_recv() { + if let NowMessage::Exec(NowExecMessage::Result(result)) = message { + if result.session_id() == 1 { + assert_eq!(result, NowExecResultMsg::new_success(1, 0)); + read_tx + .send( + NowExecProcessMsg::new(2, "cmd.exe") + .expect("valid executable") + .with_parameters("/d /c exit 7") + .expect("valid parameters") + .into(), + ) + .await + .expect("queue fresh execution after detached acknowledgment"); + } else { + assert_eq!(result, NowExecResultMsg::new_success(2, 7)); + return; + } + } + } + } + }) + .await + .expect("fresh execution completes without the detached launch killing the host"); + + read_tx + .send(NowChannelCloseMsg::default().into()) + .await + .expect("queue channel close"); + processing.await.expect("channel still accepts a normal close"); + } + + /// An attached execution and the processes it started must not outlive their channel, because no other channel + /// can reach them. + // oracle: regression #2156 attached exec children survived the loss of their channel. + #[tokio::test] + async fn attached_exec_tree_ends_with_its_channel() { + let (write_tx, mut write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + read_tx + .send(NowChannelCapsetMsg::default().into()) + .await + .expect("queue client capset"); + read_tx + .send( + NowExecProcessMsg::new(1, "powershell.exe") + .expect("valid executable") + .with_parameters( + "-NoProfile -NonInteractive -Command \ + $g = Start-Process ping.exe -ArgumentList '-n 600 127.0.0.1' -WindowStyle Hidden -PassThru; \ + [Console]::Out.WriteLine([string]$PID + ' ' + $g.Id); Start-Sleep 600", + ) + .expect("valid parameters") + .with_io_redirection() + .into(), + ) + .await + .expect("queue attached process"); + + let mut processing = Box::pin(process_messages(read_rx, write_tx, shutdown_signal, &init_progress)); + let mut stdout = Vec::new(); + let tree = tokio::time::timeout(Duration::from_secs(30), async { + loop { + tokio::select! { + result = &mut processing => panic!("processing ended before the process tree started: {result:?}"), + () = tokio::time::sleep(Duration::from_millis(10)) => {} + } + + while let Ok(message) = write_rx.try_recv() { + if let NowMessage::Exec(NowExecMessage::Data(data)) = message { + stdout.extend_from_slice(data.data()); + } + } + + if let Some(line) = String::from_utf8_lossy(&stdout) + .lines() + .next() + .filter(|_| stdout.contains(&b'\n')) + { + // Opened right away so that the process IDs cannot be reused before the assertions. + return line + .split_whitespace() + .map(|pid| { + Process::get_by_pid(pid.parse().expect("process ID"), PROCESS_SYNCHRONIZE) + .expect("open process in the tree") + }) + .collect::>(); + } + } + }) + .await + .expect("attached process reports its tree"); + assert_eq!(tree.len(), 2, "child and grandchild"); + + drop(read_tx); + let result = processing.await; + assert!(result.is_err(), "channel loss ends processing with an error"); + + for process in &tree { + assert_eq!( + process.wait(Some(10_000)).expect("wait for process"), + WAIT_OBJECT_0, + "process in the attached tree outlived its channel" + ); + } + } + + /// A detached execution is fire-and-forget, so it must outlive the channel that started it. + // oracle: mutant attached-tree cleanup also kills detached launches. + #[tokio::test] + async fn detached_exec_outlives_its_channel() { + let pid_file = std::env::temp_dir().join(format!("now-detached-exec-{}.pid", std::process::id())); + let _ = std::fs::remove_file(&pid_file); + let (write_tx, mut write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + read_tx + .send(NowChannelCapsetMsg::default().into()) + .await + .expect("queue client capset"); + read_tx + .send( + NowExecProcessMsg::new(1, "powershell.exe") + .expect("valid executable") + .with_parameters(format!( + "-NoProfile -NonInteractive -Command \ + Set-Content -LiteralPath '{}' -Value $PID; Start-Sleep 600", + pid_file.display() + )) + .expect("valid parameters") + .with_detached() + .into(), + ) + .await + .expect("queue detached process"); + + let mut processing = Box::pin(process_messages(read_rx, write_tx, shutdown_signal, &init_progress)); + let process = tokio::time::timeout(Duration::from_secs(30), async { + loop { + tokio::select! { + result = &mut processing => panic!("processing ended before the detached process started: {result:?}"), + () = tokio::time::sleep(Duration::from_millis(10)) => {} + } + + while write_rx.try_recv().is_ok() {} + + let pid = std::fs::read_to_string(&pid_file).unwrap_or_default(); + if let Ok(pid) = pid.trim().parse() { + return Process::get_by_pid(pid, PROCESS_SYNCHRONIZE | PROCESS_TERMINATE) + .expect("open detached process"); + } + } + }) + .await + .expect("detached process reports its ID"); + + drop(read_tx); + let result = processing.await; + assert!(result.is_err(), "channel loss ends processing with an error"); + + let outlived = process.wait(Some(1_000)).expect("wait for process") == WAIT_TIMEOUT; + process.terminate(1).expect("clean up detached process"); + let _ = std::fs::remove_file(&pid_file); + assert!(outlived, "detached process ended with its channel"); + } + + // oracle: mutant a full stdin queue prevents a closed DVC from interrupting the processing loop. + #[tokio::test] + async fn io_loss_interrupts_processing_while_stdin_is_backpressured() { + let (stdin_tx, _stdin_rx) = tokio::sync::mpsc::channel(1); + stdin_tx.send(()).await.expect("fill stdin queue"); + let processing = async move { + stdin_tx.send(()).await.expect("stdin queue remains open"); + Ok(()) + }; + let (io_loop_ended_tx, io_loop_ended_rx) = tokio::sync::oneshot::channel(); + io_loop_ended_tx.send(()).expect("publish IO loop exit"); + + let result = tokio::time::timeout( + Duration::from_secs(1), + process_until_io_ends(processing, io_loop_ended_rx), + ) + .await + .expect("IO loss interrupts the blocked stdin send"); + + assert!(result.is_err(), "transport loss must enter channel recovery"); + } + + // oracle: mutant IO teardown changes an already-completed normal protocol close into a transport failure. + #[tokio::test] + async fn completed_protocol_close_stays_normal_when_io_also_ends() { + let (io_loop_ended_tx, io_loop_ended_rx) = tokio::sync::oneshot::channel(); + io_loop_ended_tx.send(()).expect("publish IO loop exit"); + + process_until_io_ends(async { Ok(()) }, io_loop_ended_rx) + .await + .expect("normal protocol close remains normal"); + } + + /// A negotiated channel whose IO ends, as when the client tears the DVC down while the RDP connection stays up, + /// must be told apart from one that never negotiated, because only the former is reopened. + // oracle: mutant a channel lost after negotiation is classified as not negotiated or ended + #[tokio::test] + async fn read_channel_closing_after_negotiation_loses_the_channel() { + let (write_tx, _write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + read_tx + .send(NowChannelCapsetMsg::default().into()) + .await + .expect("queue client capset"); + drop(read_tx); + + let result = process_messages(read_rx, write_tx, shutdown_signal, &init_progress).await; + + let outcome = ChannelOutcome::new(result, &init_progress); + assert!(matches!(outcome, ChannelOutcome::Lost(_)), "{outcome:?}"); + } + + /// A channel whose IO ends before the client negotiates, as with a client that has no NOW listener, must not be + /// classified as lost, so it is never reopened. + // oracle: mutant a channel ending before negotiation is classified as lost + #[tokio::test] + async fn read_channel_closing_before_negotiation_is_not_negotiated() { + let (write_tx, _write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + drop(read_tx); + + let result = process_messages(read_rx, write_tx, shutdown_signal, &init_progress).await; + + let outcome = ChannelOutcome::new(result, &init_progress); + assert!(matches!(outcome, ChannelOutcome::NotNegotiated(_)), "{outcome:?}"); + } + + /// A channel the client closes on purpose must end rather than be reopened against the client's will. + // oracle: spec Devolutions/now-proto@7519c6ff6cfa45bb70c3a3c67df2024d482eda53:protocols/docs/NOW-spec.md §NOW_CHANNEL_CLOSE_MSG + #[tokio::test] + async fn client_channel_close_ends_the_channel() { + let (write_tx, _write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + read_tx + .send(NowChannelCapsetMsg::default().into()) + .await + .expect("queue client capset"); + read_tx + .send(NowChannelCloseMsg::default().into()) + .await + .expect("queue channel close"); + + let result = process_messages(read_rx, write_tx, shutdown_signal, &init_progress).await; + + let outcome = ChannelOutcome::new(result, &init_progress); + assert!(matches!(outcome, ChannelOutcome::Ended), "{outcome:?}"); + } + + /// A first channel that never negotiates must end the loop, so a client without a NOW listener does not keep the + /// process alive reopening a channel nobody serves. + /// It ends without an error, because the agent relaunches a process that exits with one. + // oracle: mutant an unnegotiated first channel is reopened or ends with an error + #[tokio::test(start_paused = true)] + async fn unnegotiated_first_channel_is_not_reopened() { + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + + let (result, _) = serve_script([not_negotiated()], shutdown_signal).await; + + result.expect("a client without a NOW listener is not a failure"); + } + + // oracle: mutant an ended channel is reopened or ends with an error + #[tokio::test(start_paused = true)] + async fn ended_channel_is_not_reopened() { + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + + let (result, _) = serve_script([ChannelOutcome::Ended], shutdown_signal).await; + + result.expect("an ended channel ends the loop cleanly"); + } + + // oracle: mutant a lost channel is not reopened, or is reopened before the floor delay + #[tokio::test(start_paused = true)] + async fn lost_channel_is_reopened_after_a_delay() { + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + + let (result, served_at) = serve_script([lost(), ChannelOutcome::Ended], shutdown_signal).await; + + result.expect("the reopened channel ends the loop cleanly"); + assert!(REOPEN_DELAY_FLOOR <= served_at[1] - served_at[0]); + } + + /// A client that stopped serving NOW must not be retried forever, nor in a tight loop. + #[tokio::test(start_paused = true)] + async fn reopen_gives_up_after_consecutive_unnegotiated_channels() { + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let script = core::iter::once(lost()) + .chain(core::iter::repeat_with(not_negotiated).take(MAX_CONSECUTIVE_FAILED_REOPENS)); + + let (result, served_at) = serve_script(script, shutdown_signal).await; + + // oracle: mutant reopening never gives up, or gives up without an error + assert!(result.is_err()); + + let delays = served_at.windows(2).map(|w| w[1] - w[0]).collect::>(); + // oracle: property monotonic + assert!(delays.windows(2).all(|w| w[0] <= w[1]), "{delays:?}"); + // oracle: mutant a reopen delay falls below the floor or above the cap + assert!( + delays + .iter() + .all(|delay| REOPEN_DELAY_FLOOR <= *delay && *delay <= REOPEN_DELAY_CAP), + "{delays:?}" + ); + } + + /// A reopened channel that negotiates shows the client serves NOW again, so earlier failures must not count + /// toward giving up. + // oracle: mutant a negotiated reopen keeps earlier failures in the give-up budget + #[tokio::test(start_paused = true)] + async fn negotiated_reopen_resets_the_failure_budget() { + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let failures = || core::iter::repeat_with(not_negotiated).take(MAX_CONSECUTIVE_FAILED_REOPENS - 1); + let script = core::iter::once(lost()) + .chain(failures()) + .chain(core::iter::once(lost())) + .chain(failures()) + .chain(core::iter::once(ChannelOutcome::Ended)); + + let (result, _) = serve_script(script, shutdown_signal).await; + + result.expect("failures before a negotiated reopen do not end the loop"); + } + + /// A shutdown requested while waiting to reopen must end the loop at once, without opening another channel. + // oracle: mutant a shutdown during the reopen delay waits it out or opens another channel + #[tokio::test(start_paused = true)] + async fn shutdown_during_reopen_delay_stops_without_reopening() { + let (shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let mut served = 0; + + let mut serve = Box::pin(serve_dvc( + || { + served += 1; + async { Ok(lost()) } + }, + shutdown_signal, + )); + + assert!(poll!(serve.as_mut()).is_pending()); + shutdown_handle.signal(); + + let Poll::Ready(result) = poll!(serve.as_mut()) else { + panic!("the loop is still waiting after shutdown"); + }; + result.expect("shutdown ends the loop cleanly"); + + drop(serve); + assert_eq!(served, 1); + } + /// The negotiation reply must be the downgraded capset, so the client sees the capabilities /// and heartbeat interval the agent applies rather than the server default, which carries no /// interval at all. @@ -1529,4 +2170,94 @@ mod tests { "expected a heartbeat, got {message:?}" ); } + + /// The IO thread probes a held-back channel with heartbeats only while this flag is set, so it must follow every + /// negotiation: on by default, off for a zero interval, and on again when renegotiated to a nonzero one. + // oracle: mutant the heartbeat flag ignores renegotiation or a zero interval + #[tokio::test(start_paused = true)] + async fn processing_loop_publishes_whether_heartbeats_are_enabled() { + let (write_tx, mut write_rx) = winapi_signaled_mpsc_channel().expect("create outbound channel"); + let (read_tx, read_rx) = bounded_mpsc_channel().expect("create inbound channel"); + let (_shutdown_handle, shutdown_signal) = ShutdownHandle::new(); + let init_progress = DvcInitializationProgress::default(); + + let mut loop_future = Box::pin(process_messages(read_rx, write_tx, shutdown_signal, &init_progress)); + + let capsets = [ + (NowChannelCapsetMsg::default(), true), + (capset_with_zero_heartbeat_interval(), false), + ( + NowChannelCapsetMsg::default() + .with_heartbeat_interval(Duration::from_secs(30)) + .expect("30 s is a valid heartbeat interval"), + true, + ), + ]; + + for (capset, expected) in capsets { + read_tx.send(capset.into()).await.expect("queue client capset"); + assert!(poll!(loop_future.as_mut()).is_pending()); + write_rx.try_recv().expect("negotiation reply was sent"); + + assert_eq!(init_progress.heartbeats_enabled(), expected); + } + } + + /// A client capset that disables heartbeats. + /// + /// `with_heartbeat_interval` refuses zero, but the wire format still carries it, so it is decoded from bytes. + fn capset_with_zero_heartbeat_interval() -> NowChannelCapsetMsg { + const NOW_CHANNEL_SET_HEARTBEAT: u8 = 0x01; + const MSG_FLAGS_OFFSET: usize = 6; + + // The default capset already encodes a zero interval; only the flag announcing it is missing. + let mut bytes = now_proto_pdu::ironrdp_core::encode_vec(&OwnedNowMessage::from(NowChannelCapsetMsg::default())) + .expect("encode a capset"); + *bytes.get_mut(MSG_FLAGS_OFFSET).expect("the capset has a header") |= NOW_CHANNEL_SET_HEARTBEAT; + + let message: NowMessage<'_> = now_proto_pdu::ironrdp_core::decode(&bytes).expect("decode the capset"); + let NowMessage::Channel(NowChannelMessage::Capset(capset)) = message else { + panic!("expected a capset, got {message:?}"); + }; + assert_eq!(capset.heartbeat_interval(), Some(Duration::ZERO)); + + capset + } + + /// Serves `outcomes` in order and returns the loop result with the instant each channel was served. + /// + /// Panics unless the loop serves every outcome and no more. + async fn serve_script( + outcomes: impl IntoIterator, + shutdown_signal: ShutdownSignal, + ) -> (anyhow::Result<()>, Vec) { + let mut outcomes = outcomes.into_iter().collect::>(); + let mut served_at = Vec::new(); + + let result = serve_dvc( + || { + served_at.push(Instant::now()); + let outcome = outcomes.pop_front().expect("channel served more times than scripted"); + async { Ok(outcome) } + }, + shutdown_signal, + ) + .await; + + assert!( + outcomes.is_empty(), + "{} scripted channels were never served", + outcomes.len() + ); + + (result, served_at) + } + + fn lost() -> ChannelOutcome { + ChannelOutcome::Lost(anyhow!("channel lost")) + } + + fn not_negotiated() -> ChannelOutcome { + ChannelOutcome::NotNegotiated(anyhow!("channel not negotiated")) + } } diff --git a/devolutions-session/src/main.rs b/devolutions-session/src/main.rs index cfde821c90..ac9ca15e8f 100644 --- a/devolutions-session/src/main.rs +++ b/devolutions-session/src/main.rs @@ -14,6 +14,8 @@ use ::{ #[macro_use] extern crate tracing; +use std::process::ExitCode; + use anyhow::Context; use devolutions_gateway_task::{ChildTask, ShutdownHandle, ShutdownSignal}; #[cfg(all(windows, feature = "dvc"))] @@ -22,7 +24,7 @@ use devolutions_session::{ConfHandle, get_data_dir, init_log}; use tokio::runtime::Runtime; use tokio::task::JoinHandle; -fn main() -> anyhow::Result<()> { +fn main() -> anyhow::Result { // Ensure per-user data dir exists. std::fs::create_dir_all(get_data_dir()).context("failed to create data directory")?; @@ -49,11 +51,15 @@ fn main() -> anyhow::Result<()> { info!("Waiting for shutdown signal"); - runtime.block_on(join_handle)?; + let tasks_succeeded = runtime.block_on(join_handle)?; info!("Exiting Devolutions Session"); - Ok(()) + if tasks_succeeded { + Ok(ExitCode::SUCCESS) + } else { + Ok(ExitCode::FAILURE) + } } /// Returns the session ID the Devolutions Agent passed via `--session`, if any. @@ -74,7 +80,12 @@ fn session_id_from_args() -> Option { None } -pub fn start() -> anyhow::Result<(Runtime, ShutdownHandle, JoinHandle<()>)> { +/// Spawns the tasks and returns a handle that resolves to whether every task ended without an error. +// UNTESTED: the process exits with a failure code exactly when a task failed or panicked. +// The agent relaunches the process on a failure code and leaves it stopped on success, so a task that fails while this +// still reports success leaves the session without a NOW channel. +// When changing how tasks are joined, check that every failure still clears the flag. +pub fn start() -> anyhow::Result<(Runtime, ShutdownHandle, JoinHandle)> { let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build() @@ -90,13 +101,21 @@ pub fn start() -> anyhow::Result<(Runtime, ShutdownHandle, JoinHandle<()>)> { let mut join_all = futures::future::select_all(tasks.inner.into_iter().map(|child| Box::pin(child.join()))); let join_handle = runtime.spawn(async { + let mut succeeded = true; + loop { let (result, _, rest) = join_all.await; match result { Ok(Ok(())) => trace!("A task terminated gracefully"), - Ok(Err(error)) => error!(error = format!("{error:#}"), "A task failed"), - Err(error) => error!(%error, "Something went very wrong with a task"), + Ok(Err(error)) => { + error!(error = format!("{error:#}"), "A task failed"); + succeeded = false; + } + Err(error) => { + error!(%error, "Something went very wrong with a task"); + succeeded = false; + } } if rest.is_empty() { @@ -105,6 +124,8 @@ pub fn start() -> anyhow::Result<(Runtime, ShutdownHandle, JoinHandle<()>)> { join_all = futures::future::select_all(rest); } } + + succeeded }); Ok((runtime, shutdown_handle, join_handle))