From b8d689ea3e58c15b847775c90c86dad6ffd95664 Mon Sep 17 00:00:00 2001
From: Przemyslaw Walski <pwalski@users.noreply.github.com>
Date: Wed, 10 Jan 2024 18:32:08 +0100
Subject: [PATCH] Runtime process termination on Windows with timeout

---
 utils/process/src/lib.rs | 74 +++++++++++++++++++++-------------------
 1 file changed, 39 insertions(+), 35 deletions(-)

diff --git a/utils/process/src/lib.rs b/utils/process/src/lib.rs
index 76f5e06030..21110d5d19 100644
--- a/utils/process/src/lib.rs
+++ b/utils/process/src/lib.rs
@@ -1,6 +1,7 @@
 use anyhow::{anyhow, Result};
 use derive_more::Display;
 use futures::channel::oneshot::channel;
+use futures::future::{AbortHandle, Abortable};
 use shared_child::SharedChild;
 use std::process::Command;
 use std::sync::Arc;
@@ -11,17 +12,14 @@ use std::time::Duration;
 pub mod lock;
 
 #[cfg(unix)]
-use {
-    futures::future::{AbortHandle, Abortable},
-    shared_child::unix::SharedChildExt,
-};
+use shared_child::unix::SharedChildExt;
 
 #[cfg(windows)]
-use {
-    winapi::um::handleapi::CloseHandle,
-    winapi::um::processthreadsapi::OpenProcess,
-    winapi::um::wincon::{GenerateConsoleCtrlEvent, CTRL_BREAK_EVENT},
-    winapi::um::winnt::{PROCESS_QUERY_INFORMATION, PROCESS_TERMINATE, SYNCHRONIZE},
+use std::process::ExitStatus;
+#[cfg(windows)]
+use winapi::um::{
+    errhandlingapi,
+    wincon::{GenerateConsoleCtrlEvent, CTRL_BREAK_EVENT},
 };
 
 pub trait ProcessGroupExt<T> {
@@ -125,36 +123,28 @@ impl ProcessHandle {
     }
 
     #[cfg(windows)]
-    pub async fn terminate(&self, _timeout: Duration) -> Result<()> {
+    pub async fn terminate(&self, timeout: Duration) -> Result<()> {
+        use anyhow::bail;
+
         let process = self.process.clone();
-        let process_pid = process.id();
+        let (abort_handle, abort_registration) = AbortHandle::new_pair();
+        let _process_pid = process.id();
 
-        unsafe {
-            let process_handle = OpenProcess(
-                PROCESS_QUERY_INFORMATION | SYNCHRONIZE | PROCESS_TERMINATE,
-                0,
-                process_pid,
-            );
-
-            if process_handle.is_null() {
-                return Err(anyhow!(
-                    "Unable to open the process. Error code: {}",
-                    winapi::um::errhandlingapi::GetLastError()
-                ));
-            }
-
-            let event_result = GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, process_pid);
-
-            if event_result == 0 {
-                return Err(anyhow!(
-                    "Unable to send CTRL+C event to the process. Error code: {}",
-                    winapi::um::errhandlingapi::GetLastError()
-                ));
-            };
-            CloseHandle(process_handle);
+        tokio::task::spawn_local(async move {
+            tokio::time::sleep(timeout).await;
+            abort_handle.abort();
+        });
+
+        if let Ok(Err(err)) = Abortable::new(
+            async { unsafe { ctrl_break_process(process) } },
+            abort_registration,
+        )
+        .await
+        {
+            bail!(err.context("Failed to CTRL-BREAK process"));
         }
 
-        Ok(())
+        self.check_if_running()
     }
 
     pub fn check_if_running(&self) -> Result<()> {
@@ -200,3 +190,17 @@ impl ProcessHandle {
         receiver.await.unwrap()
     }
 }
+
+#[cfg(windows)]
+unsafe fn ctrl_break_process(process: Arc<SharedChild>) -> Result<ExitStatus> {
+    let process_pid = process.id();
+    let event_result = GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, process_pid);
+
+    if event_result == 0 {
+        return Err(anyhow!(
+            "Unable to send CTRL-BREAK event to the process. Error code: {}",
+            errhandlingapi::GetLastError()
+        ));
+    };
+    Ok(process.wait()?)
+}