fix(agent): recover the NOW channel host

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<u32>) -> Resul
 }
 
 pub fn terminate_process_by_name(process_name: &str) -> Result<bool> {
-    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<bool> {
-    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<u32>) -> Result<bool> {
+/// 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<bool> {
+    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<u32>,
+    timeout_ms: Option<u32>,
+) -> Result<bool> {
     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<u32>) -
         }
 
         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<u32>) -
     Ok(false)
 }
 
+fn termination_failure_is_complete(wait: WAIT_EVENT) -> Result<bool> {
+    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<u32> {
     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"));
+    }
+}

crates/win-api-wrappers/src/wts.rs

@@ -196,6 +196,37 @@ pub fn get_session_domain_name(session_id: u32) -> anyhow::Result<String> {
     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<bool> {
+    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,