feat(agent): support canceling package broker operations

Cargo.toml

@@ -33,6 +33,10 @@ lto = true
 
 [patch.crates-io]
 tracing-appender = { git = "https://github.com/CBenoit/tracing.git", rev = "42097daf92e683cf18da7639ddccb056721a796c" }
+# TODO: remove these path patches once now-policy 0.2 and now-policy-api / now-policy-server-template 0.3 are published to crates.io.
+now-policy = { path = "D:/now-proto/.worktrees/broker-cancel/policies/rust/now-policy" }
+now-policy-api = { path = "D:/now-proto/.worktrees/broker-cancel/policies/rust/now-policy-api" }
+now-policy-server-template = { path = "D:/now-proto/.worktrees/broker-cancel/policies/rust/now-policy-server-template" }
 
 [workspace.lints.rust]
 # Declare the custom cfgs.

devolutions-agent/Cargo.toml

@@ -42,8 +42,8 @@ http-client-proxy = { path = "../crates/http-client-proxy" }
 ipnetwork = "0.20"
 notify = { version = "7", default-features = false, features = ["macos_kqueue"] }
 now-policy = "0.2"
-now-policy-api = { version = "0.2", features = ["policy-compat"] }
-now-policy-server-template = { version = "0.2", features = ["policy-compat"] }
+now-policy-api = { version = "0.3", features = ["policy-compat"] }
+now-policy-server-template = { version = "0.3", features = ["policy-compat"] }
 parking_lot = "0.12"
 prost = "0.13"
 prost-types = "0.13"

devolutions-agent/src/broker/auth.rs

@@ -3,7 +3,7 @@
 use std::path::{Path, PathBuf};
 
 use anyhow::{Context as _, bail};
-use now_policy_api::{ClientContext, PackageRequest, StatusRequest};
+use now_policy_api::{CancelRequest, ClientContext, PackageRequest, StatusRequest};
 use tokio::net::windows::named_pipe::NamedPipeServer;
 use tracing::{debug, warn};
 use widestring::U16CString;
@@ -74,6 +74,15 @@ impl PipeClient {
         self.validate_signature(skip_signature_validation)
     }
 
+    pub(crate) fn validate_cancel_request(
+        &self,
+        request: &CancelRequest,
+        skip_signature_validation: bool,
+    ) -> anyhow::Result<()> {
+        self.validate_client_context(&request.client)?;
+        self.validate_signature(skip_signature_validation)
+    }
+
     fn validate_client_context(&self, client: &ClientContext) -> anyhow::Result<()> {
         self.validate_effective_user(&client.effective_user)?;
         self.validate_executable_path(&client.client_executable_path)

devolutions-agent/src/broker/executor/mod.rs

@@ -5,6 +5,7 @@
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use now_policy_api::{Elevation, ManagerName, Scope};
+use tokio_util::sync::CancellationToken;
 use tracing::info;
 use win_api_wrappers::identity::sid::Sid;
 
@@ -53,8 +54,22 @@ pub struct ExecutionContext {
     pub scope: Option<Scope>,
     /// When true, capture the main command's combined stdout+stderr.
     pub capture_output: bool,
+    /// Cancelation signal for the operation; the executor terminates the running process when triggered.
+    pub cancel_token: CancellationToken,
 }
 
+/// Marker error returned when execution is canceled by the client.
+#[derive(Debug)]
+pub struct ExecutionCanceled;
+
+impl std::fmt::Display for ExecutionCanceled {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter.write_str("operation canceled")
+    }
+}
+
+impl std::error::Error for ExecutionCanceled {}
+
 pub type ProcessStartedCallback = std::sync::Arc<dyn Fn(DateTime<Utc>) + Send + Sync>;
 
 /// All package managers the broker knows how to drive.

devolutions-agent/src/broker/executor/windows/mod.rs

@@ -10,6 +10,7 @@ use anyhow::{Context as _, bail};
 use async_trait::async_trait;
 use devolutions_agent_shared::temp_file::{BATCH_UTF8_PREAMBLE, POWERSHELL_UTF8_ENCODING_PREAMBLE, TmpFileGuard};
 use now_policy_api::{Elevation, ManagerName, Scope};
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, warn};
 use win_api_wrappers::identity::sid::Sid;
 use win_api_wrappers::process::Process;
@@ -18,7 +19,10 @@ use win_api_wrappers::token::Token;
 use win_api_wrappers::utils::WideString;
 use windows::Win32::Security::TOKEN_ALL_ACCESS;
 
-use super::{BROKER_SUPPORTED_MANAGERS, CommandExecutor, ExecutionContext, ExecutionOutput, ProcessStartedCallback};
+use super::{
+    BROKER_SUPPORTED_MANAGERS, CommandExecutor, ExecutionCanceled, ExecutionContext, ExecutionOutput,
+    ProcessStartedCallback,
+};
 use crate::broker::policy_security;
 
 mod privileges;
@@ -329,6 +333,10 @@ fn run_plan(
     process_started: Option<ProcessStartedCallback>,
 ) -> anyhow::Result<ExecutionOutput> {
     let requires_elevation = ctx.elevation == Elevation::Elevated || ctx.scope == Some(Scope::Machine);
+    if ctx.cancel_token.is_cancelled() {
+        return Err(anyhow::Error::new(ExecutionCanceled)).context("operation canceled by client request");
+    }
+
     if requires_elevation && command_is_bun(&ctx.command) {
         bail!("elevated Bun package operations are not supported by the broker");
     }
@@ -348,14 +356,23 @@ fn run_plan(
     }
 
     // 1. Kill requested processes (best-effort; a missing process is not an error).
+    let cleanup_cancel_token = CancellationToken::new();
     for process_name in &ctx.kill_processes {
         let kill_cmd = vec![
             trusted_system32_executable("taskkill.exe"),
             "/F".to_owned(),
             "/IM".to_owned(),
             process_name.clone(),
         ];
-        match create_process(token, &kill_cmd, session_id, false, requires_elevation, None) {
+        match create_process(
+            token,
+            &kill_cmd,
+            session_id,
+            false,
+            requires_elevation,
+            &cleanup_cancel_token,
+            None,
+        ) {
             Ok(out) => info!(%process_name, exit_code = out.exit_code, "Kill-before-operation completed"),
             Err(error) => warn!(%process_name, %error, "Kill-before-operation failed (ignored)"),
         }
@@ -371,6 +388,7 @@ fn run_plan(
             session_id,
             ctx.capture_output,
             requires_elevation,
+            &ctx.cancel_token,
             None,
         )
         .context("failed to run pre-operation command")?;
@@ -391,6 +409,7 @@ fn run_plan(
         session_id,
         ctx.capture_output,
         requires_elevation,
+        &ctx.cancel_token,
         process_started,
     )?;
 
@@ -399,7 +418,15 @@ fn run_plan(
     if let Some(post) = &ctx.post_command {
         info!("Running post-operation command");
         match prepare_shell_command(token, post) {
-            Ok(command) => match create_process(token, command.args(), session_id, false, requires_elevation, None) {
+            Ok(command) => match create_process(
+                token,
+                command.args(),
+                session_id,
+                false,
+                requires_elevation,
+                &cleanup_cancel_token,
+                None,
+            ) {
                 Ok(out) if out.exit_code == 0 => {}
                 Ok(out) => warn!(exit_code = out.exit_code, "Post-operation command exited non-zero"),
                 Err(error) => warn!(%error, "Post-operation command failed"),
@@ -1317,7 +1344,7 @@ mod tests {
     use windows::Win32::Security::{NO_INHERITANCE, WinWorldSid};
 
     use super::{
-        POWERSHELL_UTF8_ENCODING_PREAMBLE, WindowsExecutor, execute_as_current_user,
+        CancellationToken, POWERSHELL_UTF8_ENCODING_PREAMBLE, WindowsExecutor, execute_as_current_user,
         prepare_chocolatey_script_in_with_default_install_root, prepare_main_command_in, prepare_shell_command_in,
         reject_unsupported_vcpkg_elevation, resolve_trusted_chocolatey_executable, resolve_winget_executable,
     };
@@ -1918,6 +1945,7 @@ mod tests {
             elevation: Elevation::Elevated,
             scope: Some(Scope::User),
             capture_output: false,
+            cancel_token: CancellationToken::new(),
         };
 
         let error = reject_unsupported_vcpkg_elevation(&ctx).expect_err("elevated vcpkg should fail");
@@ -1941,6 +1969,7 @@ mod tests {
             elevation: Elevation::Elevated,
             scope: Some(Scope::User),
             capture_output: false,
+            cancel_token: CancellationToken::new(),
         };
 
         let executor = WindowsExecutor { is_system: true };
@@ -1968,6 +1997,7 @@ mod tests {
             elevation: Elevation::Standard,
             scope: Some(Scope::User),
             capture_output: false,
+            cancel_token: CancellationToken::new(),
         };
 
         let error = execute_as_current_user(&ctx, None).expect_err("mismatched client SID should fail");

devolutions-agent/src/broker/executor/windows/process.rs

@@ -2,9 +2,11 @@
 
 use std::io::Read as _;
 use std::path::{Path, PathBuf};
+use std::time::Instant;
 
 use anyhow::{Context as _, bail};
 use chrono::Utc;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, warn};
 use win_api_wrappers::process::{self, StartupInfo};
 use win_api_wrappers::security::attributes::SecurityAttributesInit;
@@ -16,10 +18,14 @@ use windows::Win32::System::Threading::{
 };
 use windows::Win32::UI::WindowsAndMessaging::SW_HIDE;
 
-use crate::broker::executor::{ExecutionOutput, MAX_CAPTURED_OUTPUT_BYTES, ProcessStartedCallback, tail_utf8};
+use crate::broker::executor::{
+    ExecutionCanceled, ExecutionOutput, MAX_CAPTURED_OUTPUT_BYTES, ProcessStartedCallback, tail_utf8,
+};
 use crate::broker::operation_tracker::OperationTracker;
 use crate::broker::policy_security;
 
+const WAIT_SLICE_MS: u32 = 500;
+
 /// Create a process under the given token and wait for exit.
 ///
 /// This is the unified process-creation path used by both SYSTEM and current-user modes.
@@ -36,6 +42,7 @@ pub(super) fn create_process(
     session_id: u32,
     capture: bool,
     requires_elevation: bool,
+    cancel_token: &CancellationToken,
     process_started: Option<ProcessStartedCallback>,
 ) -> anyhow::Result<ExecutionOutput> {
     let cmd_line = CommandLine::new(command.to_vec());
@@ -166,28 +173,51 @@ pub(super) fn create_process(
         })
     });
 
-    let timeout_ms = operation_timeout_ms();
-    if process_info
-        .process
-        .wait(Some(timeout_ms))
-        .context("failed to wait for process")?
-        == WAIT_TIMEOUT
-    {
-        warn!(
-            session_id,
-            pid = process_info.process_id,
-            timeout_ms,
-            "Process timed out; terminating"
-        );
-        process_info
+    let deadline = Instant::now() + OperationTracker::operation_timeout();
+    loop {
+        if process_info
             .process
-            .terminate(1)
-            .context("failed to terminate timed-out process")?;
-        let _ = process_info.process.wait(None);
-        bail!(
-            "operation timed out after {} seconds",
-            OperationTracker::operation_timeout().as_secs()
-        );
+            .wait(Some(WAIT_SLICE_MS))
+            .context("failed to wait for process")?
+            != WAIT_TIMEOUT
+        {
+            break;
+        }
+
+        if cancel_token.is_cancelled() {
+            warn!(
+                session_id,
+                pid = process_info.process_id,
+                "Process canceled; terminating"
+            );
+            process_info
+                .process
+                .terminate(1)
+                .context("failed to terminate canceled process")?;
+            let _ = process_info.process.wait(None);
+            if let Some(handle) = reader {
+                let _ = handle.join();
+            }
+            return Err(anyhow::Error::new(ExecutionCanceled)).context("operation canceled by client request");
+        }
+
+        if Instant::now() >= deadline {
+            warn!(
+                session_id,
+                pid = process_info.process_id,
+                timeout_ms = operation_timeout_ms(),
+                "Process timed out; terminating"
+            );
+            process_info
+                .process
+                .terminate(1)
+                .context("failed to terminate timed-out process")?;
+            let _ = process_info.process.wait(None);
+            bail!(
+                "operation timed out after {} seconds",
+                OperationTracker::operation_timeout().as_secs()
+            );
+        }
     }
 
     let exit_code = process_info