From 74ce99b479c9da34644083314d78b2d4fd6fd578 Mon Sep 17 00:00:00 2001 From: llervikdevo Date: Thu, 25 Jun 2026 10:50:57 -0500 Subject: [PATCH 1/4] adding PSU agent functionality (POC) --- .dockerignore | 2 + Cargo.lock | 264 +++++++++++- devolutions-agent/Cargo.toml | 8 + devolutions-agent/Dockerfile.psu-grpc-poc | 34 ++ .../Run-PsuGrpcAgentContainer.ps1 | 34 ++ devolutions-agent/build.rs | 13 + .../docker/psu-grpc-entrypoint.sh | 50 +++ devolutions-agent/proto/psu_agent.proto | 131 ++++++ devolutions-agent/src/config.rs | 82 ++++ devolutions-agent/src/lib.rs | 1 + devolutions-agent/src/main.rs | 4 + devolutions-agent/src/psu_grpc_agent/mod.rs | 392 ++++++++++++++++++ .../src/psu_grpc_agent/process.rs | 381 +++++++++++++++++ devolutions-agent/src/service.rs | 7 +- 14 files changed, 1394 insertions(+), 9 deletions(-) create mode 100644 devolutions-agent/Dockerfile.psu-grpc-poc create mode 100644 devolutions-agent/Run-PsuGrpcAgentContainer.ps1 create mode 100644 devolutions-agent/docker/psu-grpc-entrypoint.sh create mode 100644 devolutions-agent/proto/psu_agent.proto create mode 100644 devolutions-agent/src/psu_grpc_agent/mod.rs create mode 100644 devolutions-agent/src/psu_grpc_agent/process.rs diff --git a/.dockerignore b/.dockerignore index bf06978ba..0346d8c66 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,4 +1,6 @@ **/target +**/node_modules +.vs/ .git/ .dockerignore Dockerfile diff --git a/Cargo.lock b/Cargo.lock index 2068eaf62..207e52585 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -456,6 +456,33 @@ dependencies = [ "tower-service", ] +[[package]] +name = "axum" +version = "0.7.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +dependencies = [ + "async-trait", + "axum-core 0.4.5", + "bytes 1.11.1", + "futures-util", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "itoa", + "matchit 0.7.3", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite 0.2.17", + "rustversion", + "serde", + "sync_wrapper 1.0.2", + "tower 0.5.3", + "tower-layer", + "tower-service", +] + [[package]] name = "axum" version = "0.8.9" @@ -509,6 +536,26 @@ dependencies = [ "tower-service", ] +[[package]] +name = "axum-core" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +dependencies = [ + "async-trait", + "bytes 1.11.1", + "futures-util", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "mime", + "pin-project-lite 0.2.17", + "rustversion", + "sync_wrapper 1.0.2", + "tower-layer", + "tower-service", +] + [[package]] name = "axum-core" version = "0.5.6" @@ -1560,6 +1607,9 @@ dependencies = [ "ironrdp", "notify-debouncer-mini", "parking_lot", + "prost 0.13.5", + "prost-types", + "protoc-bin-vendored", "quinn", "rand 0.8.6", "rcgen", @@ -1576,7 +1626,10 @@ dependencies = [ "time", "tokio 1.52.3", "tokio-rustls", + "tokio-stream", "tokio-tungstenite 0.26.2", + "tonic 0.12.3", + "tonic-build", "tracing", "url", "uuid", @@ -2250,6 +2303,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + [[package]] name = "flagset" version = "0.4.7" @@ -2987,6 +3046,19 @@ dependencies = [ "tokio-io-timeout", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.9.0", + "hyper-util", + "pin-project-lite 0.2.17", + "tokio 1.52.3", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.5.0" @@ -3948,7 +4020,7 @@ dependencies = [ "tokio 1.52.3", "tokio-stream", "tokio-util", - "tonic", + "tonic 0.11.0", "tonic-web", "tower 0.4.13", "tower-http 0.4.4", @@ -3977,7 +4049,7 @@ checksum = "d3358538b52cfcf9af4fe7aeb57d6843aafed2e8af80807bd636fd1448e94ea7" dependencies = [ "base64 0.21.7", "bytes 1.11.1", - "prost", + "prost 0.12.6", "serde", ] @@ -4041,13 +4113,13 @@ dependencies = [ "libsql-rusqlite", "libsql-sys", "parking_lot", - "prost", + "prost 0.12.6", "serde", "thiserror 1.0.69", "tokio 1.52.3", "tokio-stream", "tokio-util", - "tonic", + "tonic 0.11.0", "tracing", "uuid", "zerocopy 0.7.35", @@ -4316,6 +4388,12 @@ dependencies = [ "data-encoding-macro", ] +[[package]] +name = "multimap" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" + [[package]] name = "muxado" version = "0.4.0" @@ -4963,6 +5041,16 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "petgraph" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" +dependencies = [ + "fixedbitset", + "indexmap 2.14.0", +] + [[package]] name = "phf" version = "0.11.3" @@ -5540,7 +5628,37 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "deb1435c188b76130da55f17a466d252ff7b1418b2ad3e037d127b94e3411f29" dependencies = [ "bytes 1.11.1", - "prost-derive", + "prost-derive 0.12.6", +] + +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes 1.11.1", + "prost-derive 0.13.5", +] + +[[package]] +name = "prost-build" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" +dependencies = [ + "heck", + "itertools 0.12.1", + "log", + "multimap", + "once_cell", + "petgraph", + "prettyplease", + "prost 0.13.5", + "prost-types", + "regex", + "syn 2.0.117", + "tempfile", ] [[package]] @@ -5556,6 +5674,92 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools 0.12.1", + "proc-macro2 1.0.106", + "quote 1.0.45", + "syn 2.0.117", +] + +[[package]] +name = "prost-types" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +dependencies = [ + "prost 0.13.5", +] + +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + [[package]] name = "proxy-generators" version = "0.0.0" @@ -7502,10 +7706,40 @@ dependencies = [ "http 0.2.12", "http-body 0.4.6", "hyper 0.14.32", - "hyper-timeout", + "hyper-timeout 0.4.1", + "percent-encoding", + "pin-project 1.1.13", + "prost 0.12.6", + "tokio 1.52.3", + "tokio-stream", + "tower 0.4.13", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52" +dependencies = [ + "async-stream", + "async-trait", + "axum 0.7.9", + "base64 0.22.1", + "bytes 1.11.1", + "h2 0.4.13", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "hyper 1.9.0", + "hyper-timeout 0.5.2", + "hyper-util", "percent-encoding", "pin-project 1.1.13", - "prost", + "prost 0.13.5", + "socket2 0.5.10", "tokio 1.52.3", "tokio-stream", "tower 0.4.13", @@ -7514,6 +7748,20 @@ dependencies = [ "tracing", ] +[[package]] +name = "tonic-build" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9557ce109ea773b399c9b9e5dca39294110b74f1f342cb347a80d1fce8c26a11" +dependencies = [ + "prettyplease", + "proc-macro2 1.0.106", + "prost-build", + "prost-types", + "quote 1.0.45", + "syn 2.0.117", +] + [[package]] name = "tonic-web" version = "0.11.0" @@ -7527,7 +7775,7 @@ dependencies = [ "hyper 0.14.32", "pin-project 1.1.13", "tokio-stream", - "tonic", + "tonic 0.11.0", "tower-http 0.4.4", "tower-layer", "tower-service", diff --git a/devolutions-agent/Cargo.toml b/devolutions-agent/Cargo.toml index 27ce4e7fd..3e19a8017 100644 --- a/devolutions-agent/Cargo.toml +++ b/devolutions-agent/Cargo.toml @@ -30,6 +30,8 @@ hostname = "0.4" http-client-proxy = { path = "../crates/http-client-proxy" } ipnetwork = "0.20" parking_lot = "0.12" +prost = "0.13" +prost-types = "0.13" quinn = "0.11" rand = "0.8" # FIXME(@CBenoit): maybe we don't need this crate rcgen = { version = "0.13", features = ["pem"] } @@ -44,6 +46,8 @@ tap = "1.0" tempfile = "3" tokio-tungstenite = { version = "0.26", features = ["rustls-tls-native-roots"] } tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "tls12", "ring"] } +tokio-stream = "0.1" +tonic = { version = "0.12", features = ["transport"] } tracing = "0.1" url = { version = "2.5", features = ["serde"] } x509-parser = "0.16" @@ -99,5 +103,9 @@ features = [ [target.'cfg(windows)'.build-dependencies] embed-resource = "3.0" +[build-dependencies] +tonic-build = "0.12" +protoc-bin-vendored = "3" + [target.'cfg(windows)'.dev-dependencies] expect-test = "1.5" diff --git a/devolutions-agent/Dockerfile.psu-grpc-poc b/devolutions-agent/Dockerfile.psu-grpc-poc new file mode 100644 index 000000000..f1c2c0c9f --- /dev/null +++ b/devolutions-agent/Dockerfile.psu-grpc-poc @@ -0,0 +1,34 @@ +FROM rust:1.90-bookworm AS build + +WORKDIR /workspace + +ENV CARGO_NET_GIT_FETCH_WITH_CLI=true \ + GIT_SSL_NO_VERIFY=true + +RUN apt-get update \ + && apt-get install -y --no-install-recommends build-essential clang cmake git libclang-dev pkg-config \ + && rm -rf /var/lib/apt/lists/* + +COPY . . + +RUN rm -f rust-toolchain.toml \ + && rustc --version \ + && cargo build -p devolutions-agent --release + +FROM mcr.microsoft.com/powershell:7.5-debian-12 + +WORKDIR /opt/devolutions/agent + +ENV DAGENT_CONFIG_PATH=/etc/devolutions-agent \ + POWERSHELL_TELEMETRY_OPTOUT=1 \ + PSU_SERVER_URL=http://host.docker.internal:5006 \ + PSU_AGENT_ID=devo-agent-linux \ + PSU_DISPLAY_NAME="Devolutions Agent Linux" \ + PSU_HUBS=default + +COPY --from=build /workspace/target/release/devolutions-agent /opt/devolutions/agent/devolutions-agent +COPY devolutions-agent/docker/psu-grpc-entrypoint.sh /usr/local/bin/psu-grpc-entrypoint.sh + +RUN chmod +x /opt/devolutions/agent/devolutions-agent /usr/local/bin/psu-grpc-entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/psu-grpc-entrypoint.sh"] \ No newline at end of file diff --git a/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 b/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 new file mode 100644 index 000000000..65a536c57 --- /dev/null +++ b/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 @@ -0,0 +1,34 @@ +[CmdletBinding()] +param( + [string] $ImageName = 'devolutions-agent-psu-grpc-poc', + [string] $ContainerName = 'devo-agent-linux-poc', + [string] $ServerUrl = 'http://host.docker.internal:5006', + [string] $AgentId = 'devo-agent-linux', + [string] $DisplayName = 'Devolutions Agent Linux', + [string[]] $Hubs = @('default'), + [switch] $NoBuild +) + +$ErrorActionPreference = 'Stop' +$repoRoot = Resolve-Path (Join-Path $PSScriptRoot '..') + +if (-not $NoBuild) { + docker build ` + --file (Join-Path $PSScriptRoot 'Dockerfile.psu-grpc-poc') ` + --tag $ImageName ` + $repoRoot +} + +$existing = docker ps -a --filter "name=^/$ContainerName$" --format '{{.ID}}' +if ($existing) { + docker rm -f $ContainerName | Out-Null +} + +docker run --rm -it ` + --name $ContainerName ` + --add-host host.docker.internal:host-gateway ` + -e "PSU_SERVER_URL=$ServerUrl" ` + -e "PSU_AGENT_ID=$AgentId" ` + -e "PSU_DISPLAY_NAME=$DisplayName" ` + -e "PSU_HUBS=$($Hubs -join ',')" ` + $ImageName \ No newline at end of file diff --git a/devolutions-agent/build.rs b/devolutions-agent/build.rs index ca339ac59..b8d9ad669 100644 --- a/devolutions-agent/build.rs +++ b/devolutions-agent/build.rs @@ -1,8 +1,21 @@ fn main() { + generate_psu_agent_proto(); + #[cfg(target_os = "windows")] win::embed_version_rc(); } +fn generate_psu_agent_proto() { + let protoc = protoc_bin_vendored::protoc_bin_path().expect("failed to locate vendored protoc"); + // SAFETY: Build scripts run single-threaded for this crate before prost-build reads PROTOC. + unsafe { std::env::set_var("PROTOC", protoc) }; + + tonic_build::configure() + .build_transport(false) + .compile_protos(&["proto/psu_agent.proto"], &["proto"]) + .expect("failed to compile PSU agent proto"); +} + #[cfg(target_os = "windows")] mod win { use std::{env, fs}; diff --git a/devolutions-agent/docker/psu-grpc-entrypoint.sh b/devolutions-agent/docker/psu-grpc-entrypoint.sh new file mode 100644 index 000000000..5d16a26e4 --- /dev/null +++ b/devolutions-agent/docker/psu-grpc-entrypoint.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +mkdir -p "${DAGENT_CONFIG_PATH}" + +json_escape() { + printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g' +} + +IFS=',' read -r -a hubs <<< "${PSU_HUBS:-default}" +hubs_json="" +for hub in "${hubs[@]}"; do + hub="${hub#${hub%%[![:space:]]*}}" + hub="${hub%${hub##*[![:space:]]}}" + if [ -z "${hub}" ]; then + continue + fi + + if [ -n "${hubs_json}" ]; then + hubs_json="${hubs_json}, " + fi + hubs_json="${hubs_json}\"$(json_escape "${hub}")\"" +done + +if [ -z "${hubs_json}" ]; then + hubs_json='"default"' +fi + +cat > "${DAGENT_CONFIG_PATH}/agent.json" < environment = 7; + map metadata = 8; +} + +message StopProcess { + string correlation_id = 1; + bool kill_process = 3; +} + +message ProcessStarted { + string correlation_id = 1; + int32 process_id = 3; +} + +message ProcessCompleted { + string correlation_id = 1; + int32 exit_code = 3; + bool canceled = 4; + string error_message = 5; +} + +message StreamData { + string stream_id = 1; + uint64 sequence = 2; + bytes data = 3; + bool end_of_stream = 4; +} + +message StreamClosed { + string stream_id = 1; + string reason = 2; + bool error = 3; +} + +message AgentDiagnostic { + string level = 1; + string message = 2; + map properties = 3; +} \ No newline at end of file diff --git a/devolutions-agent/src/config.rs b/devolutions-agent/src/config.rs index e50c184ff..2220b13db 100644 --- a/devolutions-agent/src/config.rs +++ b/devolutions-agent/src/config.rs @@ -22,6 +22,7 @@ pub struct Conf { pub session: dto::SessionConf, pub tunnel: TunnelConf, pub psu_event_hub: dto::PsuEventHubConf, + pub psu_grpc_agent: dto::PsuGrpcAgentConf, pub proxy: dto::ProxyConf, pub debug: dto::DebugConf, } @@ -124,6 +125,7 @@ impl Conf { pedm: conf_file.pedm.clone().unwrap_or_default(), session: conf_file.session.clone().unwrap_or_default(), psu_event_hub: conf_file.psu_event_hub.clone().unwrap_or_default(), + psu_grpc_agent: conf_file.psu_grpc_agent.clone().unwrap_or_default(), tunnel: conf_file .tunnel .clone() @@ -580,6 +582,55 @@ pub mod dto { pub description: Option, } + /// PowerShell Universal gRPC remoting agent configuration. + #[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)] + #[serde(rename_all = "PascalCase")] + pub struct PsuGrpcAgentConf { + /// Enable the PSU gRPC agent transport. + pub enabled: bool, + + /// PSU gRPC endpoint, for example http://localhost:5000. + #[serde(skip_serializing_if = "Option::is_none")] + pub server_url: Option, + + /// Stable agent identifier presented to PSU. Defaults to the local machine name. + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + + /// Friendly display name shown in PSU. Defaults to the agent identifier. + #[serde(skip_serializing_if = "Option::is_none")] + pub display_name: Option, + + /// Hubs/queues advertised during registration. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub hubs: Vec, + + /// PowerShell child process configuration. + #[serde( + default, + rename = "PowerShell", + skip_serializing_if = "PsuPowerShellConf::is_default" + )] + pub powershell: PsuPowerShellConf, + } + + #[expect( + clippy::derivable_impls, + reason = "manually implemented so we are being explicit about the defaults" + )] + impl Default for PsuGrpcAgentConf { + fn default() -> Self { + Self { + enabled: false, + server_url: None, + agent_id: None, + display_name: None, + hubs: Vec::new(), + powershell: PsuPowerShellConf::default(), + } + } + } + #[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "PascalCase")] pub struct PsuPowerShellConf { @@ -685,6 +736,10 @@ pub mod dto { #[serde(skip_serializing_if = "Option::is_none")] pub psu_event_hub: Option, + /// PowerShell Universal gRPC remoting agent. + #[serde(skip_serializing_if = "Option::is_none")] + pub psu_grpc_agent: Option, + /// HTTP/SOCKS proxy configuration for outbound requests #[serde(skip_serializing_if = "Option::is_none")] pub proxy: Option, @@ -716,6 +771,7 @@ pub mod dto { session: Some(SessionConf { enabled: false }), tunnel: None, psu_event_hub: None, + psu_grpc_agent: None, rest: serde_json::Map::new(), } } @@ -937,4 +993,30 @@ mod tests { assert_eq!(conf.psu_event_hub.connections[0].hub, "Hub"); assert_eq!(conf.psu_event_hub.powershell.version_selector.as_deref(), Some("7.4")); } + + #[test] + fn psu_grpc_agent_config_deserializes() { + let conf_file: dto::ConfFile = serde_json::from_value(serde_json::json!({ + "PsuGrpcAgent": { + "Enabled": true, + "ServerUrl": "http://localhost:5000", + "AgentId": "agent-01", + "DisplayName": "Agent 01", + "Hubs": ["default"], + "PowerShell": { + "VersionSelector": "7.5" + } + } + })) + .expect("deserialize config"); + + let conf = Conf::from_conf_file(&conf_file).expect("load config"); + assert!(conf.psu_grpc_agent.enabled); + assert_eq!( + conf.psu_grpc_agent.server_url.as_ref().expect("server url").as_str(), + "http://localhost:5000/" + ); + assert_eq!(conf.psu_grpc_agent.agent_id.as_deref(), Some("agent-01")); + assert_eq!(conf.psu_grpc_agent.powershell.version_selector.as_deref(), Some("7.5")); + } } diff --git a/devolutions-agent/src/lib.rs b/devolutions-agent/src/lib.rs index 894ed3a89..5fb601919 100644 --- a/devolutions-agent/src/lib.rs +++ b/devolutions-agent/src/lib.rs @@ -10,6 +10,7 @@ pub mod domain_detect; pub mod enrollment; pub mod log; pub mod psu_event_hub; +pub mod psu_grpc_agent; pub mod remote_desktop; pub mod tunnel; mod tunnel_helpers; diff --git a/devolutions-agent/src/main.rs b/devolutions-agent/src/main.rs index 485c6db2e..02941753e 100644 --- a/devolutions-agent/src/main.rs +++ b/devolutions-agent/src/main.rs @@ -14,6 +14,8 @@ use http_client_proxy as _; use ipnetwork as _; use ironrdp as _; use parking_lot as _; +use prost as _; +use prost_types as _; use quinn as _; use rand as _; use reqwest as _; @@ -22,6 +24,8 @@ use serde_json as _; use tap as _; use tokio as _; use tokio_rustls as _; +use tokio_stream as _; +use tonic as _; use url as _; use uuid as _; #[cfg(windows)] diff --git a/devolutions-agent/src/psu_grpc_agent/mod.rs b/devolutions-agent/src/psu_grpc_agent/mod.rs new file mode 100644 index 000000000..abff550b1 --- /dev/null +++ b/devolutions-agent/src/psu_grpc_agent/mod.rs @@ -0,0 +1,392 @@ +mod process; + +use std::collections::HashMap; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context as _, bail}; +use async_trait::async_trait; +use backoff::backoff::Backoff as _; +use devolutions_gateway_task::{ShutdownSignal, Task}; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinSet; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Request; +use tonic::transport::Endpoint; +use uuid::Uuid; + +use crate::config::{ConfHandle, dto}; +use crate::psu_grpc_agent::process::{ProcessControl, ProcessRegistry}; + +#[allow(unused_qualifications)] +pub mod protocol { + tonic::include_proto!("devolutions.psu.agent.poc.v1"); +} + +use protocol::agent_control_client::AgentControlClient; +use protocol::agent_message::Payload as AgentPayload; +use protocol::server_message::Payload as ServerPayload; +use protocol::{ + AgentCapability, AgentDiagnostic, AgentMessage, PowerShellRuntime, RegisterAgent, StreamClosed, StreamData, +}; + +const PROTOCOL_VERSION: &str = "poc.v1"; +const CAPABILITY_JOB_EXECUTION: &str = "job_execution"; +const CAPABILITY_PSREMOTING_TUNNEL: &str = "psremoting_grpc_tunnel"; + +pub struct PsuGrpcAgentTask { + conf_handle: ConfHandle, +} + +impl PsuGrpcAgentTask { + pub fn new(conf_handle: ConfHandle) -> Self { + Self { conf_handle } + } +} + +#[async_trait] +impl Task for PsuGrpcAgentTask { + type Output = anyhow::Result<()>; + + const NAME: &'static str = "psu grpc agent"; + + async fn run(self, shutdown_signal: ShutdownSignal) -> anyhow::Result<()> { + let conf = self.conf_handle.get_conf().psu_grpc_agent.clone(); + let agent = PsuGrpcAgent::new(conf).context("failed to initialize PSU gRPC agent")?; + agent.run(shutdown_signal).await + } +} + +#[derive(Debug, Clone)] +struct PsuGrpcAgent { + conf: dto::PsuGrpcAgentConf, + server_url: String, + agent_id: String, + display_name: String, + machine_name: String, + powershell_executable: String, +} + +impl PsuGrpcAgent { + fn new(conf: dto::PsuGrpcAgentConf) -> anyhow::Result { + let server_url = conf + .server_url + .as_ref() + .context("PsuGrpcAgent is enabled but ServerUrl is not configured")? + .to_string(); + let machine_name = machine_name(); + let agent_id = conf.agent_id.clone().unwrap_or_else(|| machine_name.clone()); + let display_name = conf.display_name.clone().unwrap_or_else(|| agent_id.clone()); + let powershell_executable = resolve_powershell_executable(&conf.powershell) + .to_string_lossy() + .into_owned(); + + Ok(Self { + conf, + server_url, + agent_id, + display_name, + machine_name, + powershell_executable, + }) + } + + async fn run(self, mut shutdown_signal: ShutdownSignal) -> anyhow::Result<()> { + const RETRY_INITIAL_INTERVAL: Duration = Duration::from_secs(1); + const RETRY_MAX_INTERVAL: Duration = Duration::from_secs(60); + const RETRY_MULTIPLIER: f64 = 2.0; + const CONNECTED_THRESHOLD: Duration = Duration::from_secs(30); + + let mut backoff = backoff::ExponentialBackoffBuilder::default() + .with_initial_interval(RETRY_INITIAL_INTERVAL) + .with_max_interval(RETRY_MAX_INTERVAL) + .with_multiplier(RETRY_MULTIPLIER) + .with_max_elapsed_time(None) + .build(); + + loop { + let start = Instant::now(); + + match self.run_single_connection(&mut shutdown_signal).await { + Ok(()) => return Ok(()), + Err(error) => { + warn!(url = %self.server_url, error = format!("{error:#}"), "PSU gRPC agent connection failed") + } + } + + if start.elapsed() > CONNECTED_THRESHOLD { + backoff.reset(); + } + + let wait = match backoff.next_backoff() { + Some(wait) => wait, + None => { + warn!("PSU gRPC agent reconnect backoff exhausted, resetting"); + backoff.reset(); + RETRY_INITIAL_INTERVAL + } + }; + + info!(?wait, "Reconnecting PSU gRPC agent after backoff"); + + tokio::select! { + _ = shutdown_signal.wait() => return Ok(()), + _ = tokio::time::sleep(wait) => {} + } + } + } + + async fn run_single_connection(&self, shutdown_signal: &mut ShutdownSignal) -> anyhow::Result<()> { + let endpoint = Endpoint::from_shared(self.server_url.clone())?; + let channel = endpoint + .connect() + .await + .with_context(|| format!("failed to connect PSU gRPC endpoint at {}", self.server_url))?; + let mut client = AgentControlClient::new(channel); + + let (outgoing_tx, outgoing_rx) = mpsc::channel(256); + let powershell_version = get_powershell_version(&self.powershell_executable).await; + outgoing_tx + .send(self.create_registration_message(powershell_version)) + .await + .context("failed to queue PSU gRPC agent registration")?; + + let mut response_stream = client + .connect(Request::new(ReceiverStream::new(outgoing_rx))) + .await + .context("failed to start PSU gRPC agent stream")? + .into_inner(); + + info!(agent_id = %self.agent_id, url = %self.server_url, "Connected PSU gRPC agent"); + + let registry = ProcessRegistry::default(); + let mut process_tasks = JoinSet::new(); + let mut connection_id = String::new(); + + loop { + tokio::select! { + _ = shutdown_signal.wait() => { + process_tasks.shutdown().await; + return Ok(()); + } + message = response_stream.message() => { + let Some(message) = message.context("failed to read PSU gRPC server message")? else { + bail!("PSU gRPC server closed the agent stream"); + }; + + if !message.connection_id.trim().is_empty() { + connection_id.clone_from(&message.connection_id); + } + + self.handle_server_message( + message, + &outgoing_tx, + ®istry, + &mut process_tasks, + &mut connection_id, + ).await?; + } + Some(result) = process_tasks.join_next(), if !process_tasks.is_empty() => { + match result { + Ok(Ok(())) => trace!("PSU gRPC child process task completed"), + Ok(Err(error)) => warn!(error = format!("{error:#}"), "PSU gRPC child process task failed"), + Err(error) => warn!(%error, "PSU gRPC child process task panicked"), + } + } + } + } + } + + async fn handle_server_message( + &self, + message: protocol::ServerMessage, + outgoing_tx: &mpsc::Sender, + registry: &ProcessRegistry, + process_tasks: &mut JoinSet>, + connection_id: &mut String, + ) -> anyhow::Result<()> { + match message.payload { + Some(ServerPayload::RegisterAccepted(accepted)) => { + connection_id.clone_from(&accepted.connection_id); + info!(connection_id = %accepted.connection_id, "PSU gRPC agent registration accepted"); + } + Some(ServerPayload::StartProcess(start_process)) => { + let incoming_rx = registry.register_stream(&start_process.stream_id).await; + let (control_tx, control_rx) = oneshot::channel(); + registry + .register_process( + start_process.correlation_id.clone(), + ProcessControl { stop: control_tx }, + ) + .await; + + let agent_id = self.agent_id.clone(); + let connection_id = connection_id.clone(); + let default_executable = self.powershell_executable.clone(); + let outgoing_tx = outgoing_tx.clone(); + let registry = registry.clone(); + + process_tasks.spawn(async move { + process::run_process( + start_process, + incoming_rx, + control_rx, + outgoing_tx, + registry, + agent_id, + connection_id, + default_executable, + ) + .await + }); + } + Some(ServerPayload::StreamData(stream_data)) => registry.dispatch_stream_data(stream_data).await, + Some(ServerPayload::StreamClosed(stream_closed)) => registry.close_stream(&stream_closed.stream_id).await, + Some(ServerPayload::StopProcess(stop_process)) => { + registry + .stop_process(&stop_process.correlation_id, stop_process.kill_process) + .await; + } + Some(ServerPayload::Heartbeat(_)) | None => {} + } + + Ok(()) + } + + fn create_registration_message(&self, powershell_version: String) -> AgentMessage { + AgentMessage { + request_id: Uuid::new_v4().simple().to_string(), + agent_id: self.agent_id.clone(), + connection_id: String::new(), + timestamp: Some(timestamp_now()), + payload: Some(AgentPayload::RegisterAgent(RegisterAgent { + agent_id: self.agent_id.clone(), + instance_id: Uuid::new_v4().simple().to_string(), + display_name: self.display_name.clone(), + machine_name: self.machine_name.clone(), + os: std::env::consts::OS.to_owned(), + architecture: std::env::consts::ARCH.to_owned(), + agent_version: env!("CARGO_PKG_VERSION").to_owned(), + protocol_version: PROTOCOL_VERSION.to_owned(), + hubs: self.conf.hubs.clone(), + capabilities: vec![ + AgentCapability { + name: CAPABILITY_JOB_EXECUTION.to_owned(), + version: PROTOCOL_VERSION.to_owned(), + }, + AgentCapability { + name: CAPABILITY_PSREMOTING_TUNNEL.to_owned(), + version: PROTOCOL_VERSION.to_owned(), + }, + ], + powershell_runtimes: vec![PowerShellRuntime { + runtime_id: "pwsh-default".to_owned(), + kind: "pwsh".to_owned(), + version: powershell_version, + executable_path: self.powershell_executable.clone(), + }], + })), + } + } +} + +pub(crate) fn agent_message(agent_id: &str, connection_id: &str, payload: AgentPayload) -> AgentMessage { + AgentMessage { + request_id: Uuid::new_v4().simple().to_string(), + agent_id: agent_id.to_owned(), + connection_id: connection_id.to_owned(), + timestamp: Some(timestamp_now()), + payload: Some(payload), + } +} + +pub(crate) fn stream_data(stream_id: String, sequence: u64, data: Vec, end_of_stream: bool) -> StreamData { + StreamData { + stream_id, + sequence, + data, + end_of_stream, + } +} + +pub(crate) fn stream_closed(stream_id: String, reason: String, error: bool) -> StreamClosed { + StreamClosed { + stream_id, + reason, + error, + } +} + +pub(crate) fn diagnostic(level: &str, message: String) -> AgentDiagnostic { + AgentDiagnostic { + level: level.to_owned(), + message, + properties: HashMap::new(), + } +} + +fn timestamp_now() -> prost_types::Timestamp { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)); + prost_types::Timestamp { + seconds: i64::try_from(now.as_secs()).unwrap_or(i64::MAX), + nanos: i32::try_from(now.subsec_nanos()).unwrap_or(0), + } +} + +fn machine_name() -> String { + hostname::get() + .ok() + .map(|name| name.to_string_lossy().into_owned()) + .filter(|name| !name.trim().is_empty()) + .unwrap_or_else(|| "devolutions-agent".to_owned()) +} + +fn resolve_powershell_executable(conf: &dto::PsuPowerShellConf) -> std::ffi::OsString { + if let Some(path) = &conf.executable_path { + return path.as_str().into(); + } + + if let Some(selector) = &conf.version_selector { + if selector.eq_ignore_ascii_case("pwsh") + || selector.eq_ignore_ascii_case("pwsh-preview") + || selector.eq_ignore_ascii_case("pwsh-lts") + || selector.starts_with("pwsh-") + { + return selector.into(); + } + + return format!("pwsh-{selector}").into(); + } + + if conf.use_windows_power_shell { + "powershell.exe".into() + } else { + "pwsh".into() + } +} + +async fn get_powershell_version(executable: &str) -> String { + let output = tokio::time::timeout( + Duration::from_secs(5), + tokio::process::Command::new(executable) + .arg("-NoLogo") + .arg("-NoProfile") + .arg("-Command") + .arg("$PSVersionTable.PSVersion.ToString()") + .output(), + ) + .await; + + match output { + Ok(Ok(output)) if output.status.success() => { + let version = String::from_utf8_lossy(&output.stdout).trim().to_owned(); + if version.is_empty() { + "unknown".to_owned() + } else { + version + } + } + _ => "unknown".to_owned(), + } +} diff --git a/devolutions-agent/src/psu_grpc_agent/process.rs b/devolutions-agent/src/psu_grpc_agent/process.rs new file mode 100644 index 000000000..6a3c0c327 --- /dev/null +++ b/devolutions-agent/src/psu_grpc_agent/process.rs @@ -0,0 +1,381 @@ +use std::collections::HashMap; +use std::process::Stdio; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context as _; +use tokio::io::{AsyncBufReadExt as _, AsyncRead, AsyncReadExt as _, AsyncWriteExt as _, BufReader}; +use tokio::process::Command; +use tokio::sync::{Mutex, mpsc, oneshot}; + +use crate::psu_grpc_agent::protocol::agent_message::Payload as AgentPayload; +use crate::psu_grpc_agent::protocol::{AgentMessage, ProcessCompleted, ProcessStarted, StartProcess, StreamData}; +use crate::psu_grpc_agent::{agent_message, diagnostic, stream_closed, stream_data}; + +const PWSH_STDIN_CLOSED_EXIT_CODE: i32 = 160; + +#[derive(Debug)] +pub(super) struct ProcessControl { + pub(super) stop: oneshot::Sender, +} + +#[derive(Debug, Default, Clone)] +pub(super) struct ProcessRegistry { + inner: Arc>, +} + +#[derive(Debug, Default)] +struct ProcessRegistryInner { + streams: HashMap>, + processes: HashMap, +} + +impl ProcessRegistry { + pub(super) async fn register_stream(&self, stream_id: &str) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(256); + self.inner.lock().await.streams.insert(stream_id.to_owned(), tx); + rx + } + + pub(super) async fn dispatch_stream_data(&self, stream_data: StreamData) { + let sender = self.inner.lock().await.streams.get(&stream_data.stream_id).cloned(); + if let Some(sender) = sender { + let end_of_stream = stream_data.end_of_stream; + let stream_id = stream_data.stream_id.clone(); + if sender.send(stream_data).await.is_ok() && end_of_stream { + self.close_stream(&stream_id).await; + } + } + } + + pub(super) async fn close_stream(&self, stream_id: &str) { + self.inner.lock().await.streams.remove(stream_id); + } + + pub(super) async fn register_process(&self, correlation_id: String, control: ProcessControl) { + self.inner.lock().await.processes.insert(correlation_id, control); + } + + pub(super) async fn stop_process(&self, correlation_id: &str, kill_process: bool) { + let control = self.inner.lock().await.processes.remove(correlation_id); + if let Some(control) = control { + let _ = control.stop.send(kill_process); + } + } + + async fn remove_process(&self, correlation_id: &str) { + self.inner.lock().await.processes.remove(correlation_id); + } +} + +#[allow(clippy::too_many_arguments)] +pub(super) async fn run_process( + request: StartProcess, + incoming_rx: mpsc::Receiver, + mut control_rx: oneshot::Receiver, + outgoing_tx: mpsc::Sender, + registry: ProcessRegistry, + agent_id: String, + connection_id: String, + default_executable: String, +) -> anyhow::Result<()> { + let executable = if request.executable.trim().is_empty() { + default_executable + } else { + request.executable.clone() + }; + + info!(correlation_id = %request.correlation_id, executable = %executable, arguments = ?request.arguments, "Starting PSU gRPC child process"); + + let mut command = Command::new(&executable); + command + .args(&request.arguments) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + if !request.working_directory.trim().is_empty() && std::path::Path::new(&request.working_directory).is_dir() { + command.current_dir(&request.working_directory); + } + + for (key, value) in &request.environment { + command.env(key, value); + } + + let mut child = command + .spawn() + .with_context(|| format!("failed to start PSU gRPC child process using {executable}"))?; + let process_id_u32 = child.id().unwrap_or(0); + let process_id = i32::try_from(process_id_u32).unwrap_or(i32::MAX); + + outgoing_tx + .send(agent_message( + &agent_id, + &connection_id, + AgentPayload::ProcessStarted(ProcessStarted { + correlation_id: request.correlation_id.clone(), + process_id, + }), + )) + .await + .context("failed to send PSU gRPC ProcessStarted message")?; + + let stdin = child.stdin.take().context("child process stdin was not piped")?; + let stdout = child.stdout.take().context("child process stdout was not piped")?; + let stderr = child.stderr.take().context("child process stderr was not piped")?; + + let stdout_task = tokio::spawn(pump_stdout_to_server( + stdout, + request.stream_id.clone(), + outgoing_tx.clone(), + agent_id.clone(), + connection_id.clone(), + process_id, + )); + let stderr_task = tokio::spawn(pump_stderr_diagnostics( + stderr, + outgoing_tx.clone(), + agent_id.clone(), + connection_id.clone(), + process_id, + )); + let stdin_task = tokio::spawn(pump_server_to_stdin(incoming_rx, stdin, process_id)); + tokio::pin!(stdin_task); + + let mut stdin_closed_from_end_of_stream = false; + let mut canceled = false; + + let status = loop { + tokio::select! { + status = child.wait() => break status.context("failed to wait for PSU gRPC child process")?, + stdin_result = &mut stdin_task => { + stdin_closed_from_end_of_stream = stdin_result.unwrap_or(false); + info!(process_id, "Finished receiving PSU gRPC stdin data; waiting for graceful child process exit"); + + match tokio::time::timeout(Duration::from_secs(5), child.wait()).await { + Ok(status) => break status.context("failed to wait for PSU gRPC child process")?, + Err(_) => { + warn!(process_id, "PSU gRPC child process did not exit after stdin closed; killing process tree"); + child.start_kill().context("failed to kill PSU gRPC child process")?; + canceled = true; + } + } + } + kill_process = &mut control_rx => { + match kill_process { + Ok(true) => { + info!(process_id, correlation_id = %request.correlation_id, "Killing PSU gRPC child process on server request"); + child.start_kill().context("failed to kill PSU gRPC child process")?; + canceled = true; + break child.wait().await.context("failed to wait for killed PSU gRPC child process")?; + } + Ok(false) => { + info!(process_id, correlation_id = %request.correlation_id, "Graceful stop requested for PSU gRPC child process; waiting for stream shutdown"); + } + Err(_) => {} + } + } + } + }; + + stdout_task.abort(); + stderr_task.abort(); + + let exit_code = status.code().unwrap_or(-1); + if stdin_closed_from_end_of_stream && exit_code == PWSH_STDIN_CLOSED_EXIT_CODE { + info!( + process_id, + exit_code, "PSU gRPC child process exited with expected code after stdin EOF for pwsh -s" + ); + } else { + info!(process_id, exit_code, "PSU gRPC child process exited"); + } + + let _ = outgoing_tx + .send(agent_message( + &agent_id, + &connection_id, + AgentPayload::StreamClosed(stream_closed( + request.stream_id.clone(), + "child process completed".to_owned(), + false, + )), + )) + .await; + + outgoing_tx + .send(agent_message( + &agent_id, + &connection_id, + AgentPayload::ProcessCompleted(ProcessCompleted { + correlation_id: request.correlation_id.clone(), + exit_code, + canceled, + error_message: String::new(), + }), + )) + .await + .context("failed to send PSU gRPC ProcessCompleted message")?; + + registry.close_stream(&request.stream_id).await; + registry.remove_process(&request.correlation_id).await; + + Ok(()) +} + +async fn pump_stdout_to_server( + mut stdout: R, + stream_id: String, + outgoing_tx: mpsc::Sender, + agent_id: String, + connection_id: String, + process_id: i32, +) -> anyhow::Result<()> +where + R: AsyncRead + Unpin, +{ + let mut buffer = [0u8; 4096]; + let mut line = Vec::new(); + let mut sequence = 0; + + loop { + let read = stdout.read(&mut buffer).await.context("failed to read child stdout")?; + if read == 0 { + break; + } + + for byte in &buffer[..read] { + match *byte { + b'\r' => {} + b'\n' => { + send_stream_frame( + &outgoing_tx, + &agent_id, + &connection_id, + &stream_id, + sequence, + std::mem::take(&mut line), + false, + ) + .await?; + sequence += 1; + } + byte => line.push(byte), + } + } + } + + if !line.is_empty() { + send_stream_frame( + &outgoing_tx, + &agent_id, + &connection_id, + &stream_id, + sequence, + line, + false, + ) + .await?; + sequence += 1; + } + + send_stream_frame( + &outgoing_tx, + &agent_id, + &connection_id, + &stream_id, + sequence, + Vec::new(), + true, + ) + .await?; + info!(process_id, stream_id = %stream_id, sequence, "Finished sending PSU gRPC stdout frames"); + Ok(()) +} + +async fn send_stream_frame( + outgoing_tx: &mpsc::Sender, + agent_id: &str, + connection_id: &str, + stream_id: &str, + sequence: u64, + data: Vec, + end_of_stream: bool, +) -> anyhow::Result<()> { + outgoing_tx + .send(agent_message( + agent_id, + connection_id, + AgentPayload::StreamData(stream_data(stream_id.to_owned(), sequence, data, end_of_stream)), + )) + .await + .context("failed to send PSU gRPC stdout frame") +} + +async fn pump_server_to_stdin( + mut incoming_rx: mpsc::Receiver, + mut stdin: tokio::process::ChildStdin, + process_id: i32, +) -> bool { + let mut closed_from_end_of_stream = false; + + while let Some(frame) = incoming_rx.recv().await { + if frame.end_of_stream { + info!(process_id, "Received PSU gRPC stdin end-of-stream; closing child stdin"); + closed_from_end_of_stream = true; + break; + } + + let mut data = frame.data; + if !ends_with_line_ending(&data) { + data.push(b'\n'); + } + + if let Err(error) = stdin.write_all(&data).await { + warn!(process_id, %error, "Failed to write PSU gRPC frame to child stdin"); + break; + } + + if let Err(error) = stdin.flush().await { + warn!(process_id, %error, "Failed to flush child stdin"); + break; + } + } + + let _ = stdin.shutdown().await; + closed_from_end_of_stream +} + +async fn pump_stderr_diagnostics( + stderr: R, + outgoing_tx: mpsc::Sender, + agent_id: String, + connection_id: String, + process_id: i32, +) -> anyhow::Result<()> +where + R: AsyncRead + Unpin, +{ + let mut lines = BufReader::new(stderr).lines(); + while let Some(line) = lines.next_line().await.context("failed to read child stderr")? { + if line.trim().is_empty() { + continue; + } + + outgoing_tx + .send(agent_message( + &agent_id, + &connection_id, + AgentPayload::Diagnostic(diagnostic("warning", format!("pwsh[{process_id}] {line}"))), + )) + .await + .context("failed to send PSU gRPC stderr diagnostic")?; + } + + Ok(()) +} + +fn ends_with_line_ending(data: &[u8]) -> bool { + data.ends_with(b"\n") || data.ends_with(b"\r") +} diff --git a/devolutions-agent/src/service.rs b/devolutions-agent/src/service.rs index d1ab5e6cc..9eeb178b8 100644 --- a/devolutions-agent/src/service.rs +++ b/devolutions-agent/src/service.rs @@ -5,6 +5,7 @@ use devolutions_agent::AgentServiceEvent; use devolutions_agent::config::ConfHandle; use devolutions_agent::log::AgentLog; use devolutions_agent::psu_event_hub::PsuEventHubTask; +use devolutions_agent::psu_grpc_agent::PsuGrpcAgentTask; use devolutions_agent::remote_desktop::RemoteDesktopTask; #[cfg(windows)] use devolutions_agent::session_manager::SessionManager; @@ -237,7 +238,11 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { } if conf.psu_event_hub.enabled { - tasks.register(PsuEventHubTask::new(conf_handle)); + tasks.register(PsuEventHubTask::new(conf_handle.clone())); + } + + if conf.psu_grpc_agent.enabled { + tasks.register(PsuGrpcAgentTask::new(conf_handle)); } Ok(TasksCtx { From 54770bcb74a72e128fc439c18f8c7a40b8fb3e76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Moreau?= Date: Tue, 30 Jun 2026 15:58:30 -0400 Subject: [PATCH 2/4] feat(agent): harden PSU gRPC POC support (#1843) Hardens the PSU gRPC agent POC and adds optional AppToken support for authenticating the agent to PowerShell Universal. The POC container also avoids running as root and no longer disables TLS verification during Docker builds. This keeps unauthenticated local development possible when no AppToken is configured, while allowing PSU deployments to pass an application token through agent configuration or the container helper. Issue: N/A --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- devolutions-agent/Dockerfile.psu-grpc-poc | 12 +- .../Run-PsuGrpcAgentContainer.ps1 | 28 +- .../docker/psu-grpc-entrypoint.sh | 6 + devolutions-agent/src/config.rs | 7 + devolutions-agent/src/psu_grpc_agent/mod.rs | 56 +++- .../src/psu_grpc_agent/process.rs | 283 +++++++++++++++--- 6 files changed, 339 insertions(+), 53 deletions(-) diff --git a/devolutions-agent/Dockerfile.psu-grpc-poc b/devolutions-agent/Dockerfile.psu-grpc-poc index f1c2c0c9f..e1f711b6c 100644 --- a/devolutions-agent/Dockerfile.psu-grpc-poc +++ b/devolutions-agent/Dockerfile.psu-grpc-poc @@ -2,8 +2,7 @@ FROM rust:1.90-bookworm AS build WORKDIR /workspace -ENV CARGO_NET_GIT_FETCH_WITH_CLI=true \ - GIT_SSL_NO_VERIFY=true +ENV CARGO_NET_GIT_FETCH_WITH_CLI=true RUN apt-get update \ && apt-get install -y --no-install-recommends build-essential clang cmake git libclang-dev pkg-config \ @@ -29,6 +28,11 @@ ENV DAGENT_CONFIG_PATH=/etc/devolutions-agent \ COPY --from=build /workspace/target/release/devolutions-agent /opt/devolutions/agent/devolutions-agent COPY devolutions-agent/docker/psu-grpc-entrypoint.sh /usr/local/bin/psu-grpc-entrypoint.sh -RUN chmod +x /opt/devolutions/agent/devolutions-agent /usr/local/bin/psu-grpc-entrypoint.sh +RUN useradd --system --create-home --home-dir /var/lib/devolutions-agent --shell /usr/sbin/nologin devolutions-agent \ + && mkdir -p /etc/devolutions-agent \ + && chown -R devolutions-agent:devolutions-agent /etc/devolutions-agent /opt/devolutions/agent \ + && chmod +x /opt/devolutions/agent/devolutions-agent /usr/local/bin/psu-grpc-entrypoint.sh -ENTRYPOINT ["/usr/local/bin/psu-grpc-entrypoint.sh"] \ No newline at end of file +USER devolutions-agent + +ENTRYPOINT ["/usr/local/bin/psu-grpc-entrypoint.sh"] diff --git a/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 b/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 index 65a536c57..974d798a4 100644 --- a/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 +++ b/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 @@ -5,6 +5,7 @@ param( [string] $ServerUrl = 'http://host.docker.internal:5006', [string] $AgentId = 'devo-agent-linux', [string] $DisplayName = 'Devolutions Agent Linux', + [string] $AppToken, [string[]] $Hubs = @('default'), [switch] $NoBuild ) @@ -24,11 +25,22 @@ if ($existing) { docker rm -f $ContainerName | Out-Null } -docker run --rm -it ` - --name $ContainerName ` - --add-host host.docker.internal:host-gateway ` - -e "PSU_SERVER_URL=$ServerUrl" ` - -e "PSU_AGENT_ID=$AgentId" ` - -e "PSU_DISPLAY_NAME=$DisplayName" ` - -e "PSU_HUBS=$($Hubs -join ',')" ` - $ImageName \ No newline at end of file +$dockerArgs = @( + 'run', + '--rm', + '-it', + '--name', $ContainerName, + '--add-host', 'host.docker.internal:host-gateway', + '-e', "PSU_SERVER_URL=$ServerUrl", + '-e', "PSU_AGENT_ID=$AgentId", + '-e', "PSU_DISPLAY_NAME=$DisplayName", + '-e', "PSU_HUBS=$($Hubs -join ',')" +) + +if (-not [string]::IsNullOrWhiteSpace($AppToken)) { + $dockerArgs += @('-e', "PSU_APP_TOKEN=$AppToken") +} + +$dockerArgs += $ImageName + +docker @dockerArgs diff --git a/devolutions-agent/docker/psu-grpc-entrypoint.sh b/devolutions-agent/docker/psu-grpc-entrypoint.sh index 5d16a26e4..ce6f49ea0 100644 --- a/devolutions-agent/docker/psu-grpc-entrypoint.sh +++ b/devolutions-agent/docker/psu-grpc-entrypoint.sh @@ -26,6 +26,11 @@ if [ -z "${hubs_json}" ]; then hubs_json='"default"' fi +app_token_property="" +if [ -n "${PSU_APP_TOKEN:-}" ]; then + app_token_property=" \"AppToken\": \"$(json_escape "${PSU_APP_TOKEN}")\"," +fi + cat > "${DAGENT_CONFIG_PATH}/agent.json" < "${DAGENT_CONFIG_PATH}/agent.json" <, + /// PSU application token used to authenticate the gRPC agent. + #[serde(skip_serializing_if = "Option::is_none")] + pub app_token: Option, + /// Hubs/queues advertised during registration. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub hubs: Vec, @@ -625,6 +629,7 @@ pub mod dto { server_url: None, agent_id: None, display_name: None, + app_token: None, hubs: Vec::new(), powershell: PsuPowerShellConf::default(), } @@ -1002,6 +1007,7 @@ mod tests { "ServerUrl": "http://localhost:5000", "AgentId": "agent-01", "DisplayName": "Agent 01", + "AppToken": "app-token", "Hubs": ["default"], "PowerShell": { "VersionSelector": "7.5" @@ -1017,6 +1023,7 @@ mod tests { "http://localhost:5000/" ); assert_eq!(conf.psu_grpc_agent.agent_id.as_deref(), Some("agent-01")); + assert_eq!(conf.psu_grpc_agent.app_token.as_deref(), Some("app-token")); assert_eq!(conf.psu_grpc_agent.powershell.version_selector.as_deref(), Some("7.5")); } } diff --git a/devolutions-agent/src/psu_grpc_agent/mod.rs b/devolutions-agent/src/psu_grpc_agent/mod.rs index abff550b1..c9aa01dda 100644 --- a/devolutions-agent/src/psu_grpc_agent/mod.rs +++ b/devolutions-agent/src/psu_grpc_agent/mod.rs @@ -7,17 +7,18 @@ use anyhow::{Context as _, bail}; use async_trait::async_trait; use backoff::backoff::Backoff as _; use devolutions_gateway_task::{ShutdownSignal, Task}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::mpsc; use tokio::task::JoinSet; use tokio_stream::wrappers::ReceiverStream; use tonic::Request; +use tonic::metadata::MetadataValue; use tonic::transport::Endpoint; use uuid::Uuid; use crate::config::{ConfHandle, dto}; use crate::psu_grpc_agent::process::{ProcessControl, ProcessRegistry}; -#[allow(unused_qualifications)] +#[allow(unused_qualifications, clippy::clone_on_ref_ptr, clippy::similar_names)] pub mod protocol { tonic::include_proto!("devolutions.psu.agent.poc.v1"); } @@ -63,6 +64,7 @@ struct PsuGrpcAgent { agent_id: String, display_name: String, machine_name: String, + app_token: Option, powershell_executable: String, } @@ -76,6 +78,7 @@ impl PsuGrpcAgent { let machine_name = machine_name(); let agent_id = conf.agent_id.clone().unwrap_or_else(|| machine_name.clone()); let display_name = conf.display_name.clone().unwrap_or_else(|| agent_id.clone()); + let app_token = conf.app_token.clone().filter(|token| !token.trim().is_empty()); let powershell_executable = resolve_powershell_executable(&conf.powershell) .to_string_lossy() .into_owned(); @@ -86,6 +89,7 @@ impl PsuGrpcAgent { agent_id, display_name, machine_name, + app_token, powershell_executable, }) } @@ -151,7 +155,10 @@ impl PsuGrpcAgent { .context("failed to queue PSU gRPC agent registration")?; let mut response_stream = client - .connect(Request::new(ReceiverStream::new(outgoing_rx))) + .connect(connect_request( + ReceiverStream::new(outgoing_rx), + self.app_token.as_deref(), + )?) .await .context("failed to start PSU gRPC agent stream")? .into_inner(); @@ -211,7 +218,7 @@ impl PsuGrpcAgent { } Some(ServerPayload::StartProcess(start_process)) => { let incoming_rx = registry.register_stream(&start_process.stream_id).await; - let (control_tx, control_rx) = oneshot::channel(); + let (control_tx, control_rx) = mpsc::channel(8); registry .register_process( start_process.correlation_id.clone(), @@ -324,6 +331,20 @@ pub(crate) fn diagnostic(level: &str, message: String) -> AgentDiagnostic { } } +fn connect_request(stream: T, app_token: Option<&str>) -> anyhow::Result> { + let mut request = Request::new(stream); + + if let Some(token) = app_token { + let authorization = format!("Bearer {token}"); + request.metadata_mut().insert( + "authorization", + MetadataValue::try_from(authorization).context("invalid PSU gRPC AppToken metadata")?, + ); + } + + Ok(request) +} + fn timestamp_now() -> prost_types::Timestamp { let now = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -390,3 +411,30 @@ async fn get_powershell_version(executable: &str) -> String { _ => "unknown".to_owned(), } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn connect_request_omits_authorization_without_app_token() { + let request = connect_request((), None).expect("create request"); + + assert!(!request.metadata().contains_key("authorization")); + } + + #[test] + fn connect_request_adds_authorization_with_app_token() { + let request = connect_request((), Some("token")).expect("create request"); + + assert_eq!( + request + .metadata() + .get("authorization") + .expect("authorization metadata") + .to_str() + .expect("metadata string"), + "Bearer token" + ); + } +} diff --git a/devolutions-agent/src/psu_grpc_agent/process.rs b/devolutions-agent/src/psu_grpc_agent/process.rs index 6a3c0c327..5aed82673 100644 --- a/devolutions-agent/src/psu_grpc_agent/process.rs +++ b/devolutions-agent/src/psu_grpc_agent/process.rs @@ -1,12 +1,13 @@ use std::collections::HashMap; -use std::process::Stdio; +use std::process::{ExitStatus, Stdio}; use std::sync::Arc; use std::time::Duration; use anyhow::Context as _; use tokio::io::{AsyncBufReadExt as _, AsyncRead, AsyncReadExt as _, AsyncWriteExt as _, BufReader}; -use tokio::process::Command; -use tokio::sync::{Mutex, mpsc, oneshot}; +use tokio::process::{Child, Command}; +use tokio::sync::{Mutex, mpsc}; +use tokio::task::JoinHandle; use crate::psu_grpc_agent::protocol::agent_message::Payload as AgentPayload; use crate::psu_grpc_agent::protocol::{AgentMessage, ProcessCompleted, ProcessStarted, StartProcess, StreamData}; @@ -16,7 +17,7 @@ const PWSH_STDIN_CLOSED_EXIT_CODE: i32 = 160; #[derive(Debug)] pub(super) struct ProcessControl { - pub(super) stop: oneshot::Sender, + pub(super) stop: mpsc::Sender, } #[derive(Debug, Default, Clone)] @@ -57,9 +58,17 @@ impl ProcessRegistry { } pub(super) async fn stop_process(&self, correlation_id: &str, kill_process: bool) { - let control = self.inner.lock().await.processes.remove(correlation_id); + let control = { + let mut inner = self.inner.lock().await; + if kill_process { + inner.processes.remove(correlation_id).map(|control| control.stop) + } else { + inner.processes.get(correlation_id).map(|control| control.stop.clone()) + } + }; + if let Some(control) = control { - let _ = control.stop.send(kill_process); + let _ = control.send(kill_process).await; } } @@ -72,12 +81,42 @@ impl ProcessRegistry { pub(super) async fn run_process( request: StartProcess, incoming_rx: mpsc::Receiver, - mut control_rx: oneshot::Receiver, + control_rx: mpsc::Receiver, outgoing_tx: mpsc::Sender, registry: ProcessRegistry, agent_id: String, connection_id: String, default_executable: String, +) -> anyhow::Result<()> { + let correlation_id = request.correlation_id.clone(); + let stream_id = request.stream_id.clone(); + + let result = run_process_inner( + request, + incoming_rx, + control_rx, + outgoing_tx, + agent_id, + connection_id, + default_executable, + ) + .await; + + registry.close_stream(&stream_id).await; + registry.remove_process(&correlation_id).await; + + result +} + +#[allow(clippy::too_many_arguments)] +async fn run_process_inner( + request: StartProcess, + incoming_rx: mpsc::Receiver, + mut control_rx: mpsc::Receiver, + outgoing_tx: mpsc::Sender, + agent_id: String, + connection_id: String, + default_executable: String, ) -> anyhow::Result<()> { let executable = if request.executable.trim().is_empty() { default_executable @@ -103,9 +142,32 @@ pub(super) async fn run_process( command.env(key, value); } - let mut child = command - .spawn() - .with_context(|| format!("failed to start PSU gRPC child process using {executable}"))?; + let mut child = match command.spawn() { + Ok(child) => child, + Err(error) => { + let error = + anyhow::Error::new(error).context(format!("failed to start PSU gRPC child process using {executable}")); + let error_message = format!("{error:#}"); + let _ = outgoing_tx + .send(agent_message( + &agent_id, + &connection_id, + AgentPayload::StreamClosed(stream_closed(request.stream_id.clone(), error_message.clone(), true)), + )) + .await; + let _ = send_process_completed( + &outgoing_tx, + &agent_id, + &connection_id, + &request.correlation_id, + -1, + false, + error_message, + ) + .await; + return Err(error); + } + }; let process_id_u32 = child.id().unwrap_or(0); let process_id = i32::try_from(process_id_u32).unwrap_or(i32::MAX); @@ -140,47 +202,55 @@ pub(super) async fn run_process( connection_id.clone(), process_id, )); - let stdin_task = tokio::spawn(pump_server_to_stdin(incoming_rx, stdin, process_id)); - tokio::pin!(stdin_task); + let mut stdin_task = tokio::spawn(pump_server_to_stdin(incoming_rx, stdin, process_id)); let mut stdin_closed_from_end_of_stream = false; + let mut stdin_task_completed = false; let mut canceled = false; let status = loop { tokio::select! { status = child.wait() => break status.context("failed to wait for PSU gRPC child process")?, stdin_result = &mut stdin_task => { + stdin_task_completed = true; stdin_closed_from_end_of_stream = stdin_result.unwrap_or(false); info!(process_id, "Finished receiving PSU gRPC stdin data; waiting for graceful child process exit"); - match tokio::time::timeout(Duration::from_secs(5), child.wait()).await { - Ok(status) => break status.context("failed to wait for PSU gRPC child process")?, - Err(_) => { - warn!(process_id, "PSU gRPC child process did not exit after stdin closed; killing process tree"); - child.start_kill().context("failed to kill PSU gRPC child process")?; - canceled = true; - } - } + let (exit_status, killed) = wait_for_graceful_child_exit(&mut child, process_id).await?; + canceled |= killed; + break exit_status; } - kill_process = &mut control_rx => { + kill_process = control_rx.recv() => { match kill_process { - Ok(true) => { + Some(true) => { info!(process_id, correlation_id = %request.correlation_id, "Killing PSU gRPC child process on server request"); child.start_kill().context("failed to kill PSU gRPC child process")?; canceled = true; break child.wait().await.context("failed to wait for killed PSU gRPC child process")?; } - Ok(false) => { - info!(process_id, correlation_id = %request.correlation_id, "Graceful stop requested for PSU gRPC child process; waiting for stream shutdown"); + Some(false) => { + info!(process_id, correlation_id = %request.correlation_id, "Gracefully stopping PSU gRPC child process by closing stdin"); + canceled = true; + stdin_task.abort(); + let _ = (&mut stdin_task).await; + stdin_task_completed = true; + let (exit_status, killed) = wait_for_graceful_child_exit(&mut child, process_id).await?; + canceled |= killed; + break exit_status; } - Err(_) => {} + None => {} } } } }; - stdout_task.abort(); - stderr_task.abort(); + if !stdin_task_completed { + stdin_task.abort(); + let _ = stdin_task.await; + } + + await_pump_task(stdout_task, process_id, "stdout").await; + await_pump_task(stderr_task, process_id, "stderr").await; let exit_code = status.code().unwrap_or(-1); if stdin_closed_from_end_of_stream && exit_code == PWSH_STDIN_CLOSED_EXIT_CODE { @@ -204,24 +274,76 @@ pub(super) async fn run_process( )) .await; + send_process_completed( + &outgoing_tx, + &agent_id, + &connection_id, + &request.correlation_id, + exit_code, + canceled, + String::new(), + ) + .await + .context("failed to send PSU gRPC ProcessCompleted message")?; + + Ok(()) +} + +async fn wait_for_graceful_child_exit(child: &mut Child, process_id: i32) -> anyhow::Result<(ExitStatus, bool)> { + match tokio::time::timeout(Duration::from_secs(5), child.wait()).await { + Ok(status) => Ok((status.context("failed to wait for PSU gRPC child process")?, false)), + Err(_) => { + warn!( + process_id, + "PSU gRPC child process did not exit after stdin closed; killing child process" + ); + child.start_kill().context("failed to kill PSU gRPC child process")?; + let status = child + .wait() + .await + .context("failed to wait for killed PSU gRPC child process")?; + Ok((status, true)) + } + } +} + +async fn await_pump_task(mut task: JoinHandle>, process_id: i32, stream_name: &'static str) { + tokio::select! { + result = &mut task => match result { + Ok(Ok(())) => {} + Ok(Err(error)) => warn!(process_id, stream_name, error = format!("{error:#}"), "PSU gRPC child stream pump failed"), + Err(error) => warn!(process_id, stream_name, %error, "PSU gRPC child stream pump panicked"), + }, + _ = tokio::time::sleep(Duration::from_secs(5)) => { + warn!(process_id, stream_name, "Timed out draining PSU gRPC child stream pump"); + task.abort(); + let _ = task.await; + } + } +} + +async fn send_process_completed( + outgoing_tx: &mpsc::Sender, + agent_id: &str, + connection_id: &str, + correlation_id: &str, + exit_code: i32, + canceled: bool, + error_message: String, +) -> anyhow::Result<()> { outgoing_tx .send(agent_message( - &agent_id, - &connection_id, + agent_id, + connection_id, AgentPayload::ProcessCompleted(ProcessCompleted { - correlation_id: request.correlation_id.clone(), + correlation_id: correlation_id.to_owned(), exit_code, canceled, - error_message: String::new(), + error_message, }), )) .await - .context("failed to send PSU gRPC ProcessCompleted message")?; - - registry.close_stream(&request.stream_id).await; - registry.remove_process(&request.correlation_id).await; - - Ok(()) + .context("failed to send PSU gRPC ProcessCompleted message") } async fn pump_stdout_to_server( @@ -379,3 +501,90 @@ where fn ends_with_line_ending(data: &[u8]) -> bool { data.ends_with(b"\n") || data.ends_with(b"\r") } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn graceful_stop_keeps_process_registered_for_later_kill() { + let registry = ProcessRegistry::default(); + let (control_tx, mut control_rx) = mpsc::channel(8); + + registry + .register_process("correlation-id".to_owned(), ProcessControl { stop: control_tx }) + .await; + + registry.stop_process("correlation-id", false).await; + assert_eq!(control_rx.recv().await, Some(false)); + assert!(registry.inner.lock().await.processes.contains_key("correlation-id")); + + registry.stop_process("correlation-id", true).await; + assert_eq!(control_rx.recv().await, Some(true)); + assert!(!registry.inner.lock().await.processes.contains_key("correlation-id")); + } + + #[tokio::test] + async fn run_process_cleans_registry_and_reports_spawn_failure() { + let registry = ProcessRegistry::default(); + let incoming_rx = registry.register_stream("stream-id").await; + let (control_tx, control_rx) = mpsc::channel(8); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(8); + + registry + .register_process("correlation-id".to_owned(), ProcessControl { stop: control_tx }) + .await; + + let result = run_process( + StartProcess { + correlation_id: "correlation-id".to_owned(), + stream_id: "stream-id".to_owned(), + executable: "definitely-not-a-devolutions-agent-test-command".to_owned(), + arguments: Vec::new(), + working_directory: String::new(), + environment: HashMap::new(), + metadata: HashMap::new(), + }, + incoming_rx, + control_rx, + outgoing_tx, + registry.clone(), + "agent-id".to_owned(), + "connection-id".to_owned(), + "pwsh".to_owned(), + ) + .await; + + assert!(result.is_err()); + + let registry = registry.inner.lock().await; + assert!(registry.streams.is_empty()); + assert!(registry.processes.is_empty()); + drop(registry); + + let stream_message = outgoing_rx.recv().await.expect("stream closed message"); + match stream_message.payload { + Some(AgentPayload::StreamClosed(closed)) => { + assert_eq!(closed.stream_id, "stream-id"); + assert!(closed.error); + assert!(closed.reason.contains("failed to start PSU gRPC child process")); + } + payload => panic!("unexpected payload: {payload:?}"), + } + + let completed_message = outgoing_rx.recv().await.expect("process completed message"); + match completed_message.payload { + Some(AgentPayload::ProcessCompleted(completed)) => { + assert_eq!(completed.correlation_id, "correlation-id"); + assert_eq!(completed.exit_code, -1); + assert!(!completed.canceled); + assert!( + completed + .error_message + .contains("failed to start PSU gRPC child process") + ); + } + payload => panic!("unexpected payload: {payload:?}"), + } + } +} From a8c44ae02d9ddb20df24238756ed5578c2175523 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20CORTIER?= Date: Thu, 2 Jul 2026 20:50:00 +0900 Subject: [PATCH 3/4] =?UTF-8?q?Address=20Copilot=E2=80=99s=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- devolutions-agent/src/psu_grpc_agent/mod.rs | 4 ++-- .../src/psu_grpc_agent/process.rs | 23 +++++++++++++++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/devolutions-agent/src/psu_grpc_agent/mod.rs b/devolutions-agent/src/psu_grpc_agent/mod.rs index c9aa01dda..aa1819355 100644 --- a/devolutions-agent/src/psu_grpc_agent/mod.rs +++ b/devolutions-agent/src/psu_grpc_agent/mod.rs @@ -73,7 +73,7 @@ impl PsuGrpcAgent { let server_url = conf .server_url .as_ref() - .context("PsuGrpcAgent is enabled but ServerUrl is not configured")? + .context("server URL is not configured but the PSU gRPC agent is enabled")? .to_string(); let machine_name = machine_name(); let agent_id = conf.agent_id.clone().unwrap_or_else(|| machine_name.clone()); @@ -372,7 +372,7 @@ fn resolve_powershell_executable(conf: &dto::PsuPowerShellConf) -> std::ffi::OsS if selector.eq_ignore_ascii_case("pwsh") || selector.eq_ignore_ascii_case("pwsh-preview") || selector.eq_ignore_ascii_case("pwsh-lts") - || selector.starts_with("pwsh-") + || selector.to_ascii_lowercase().starts_with("pwsh-") { return selector.into(); } diff --git a/devolutions-agent/src/psu_grpc_agent/process.rs b/devolutions-agent/src/psu_grpc_agent/process.rs index 5aed82673..7e84ab031 100644 --- a/devolutions-agent/src/psu_grpc_agent/process.rs +++ b/devolutions-agent/src/psu_grpc_agent/process.rs @@ -43,7 +43,10 @@ impl ProcessRegistry { if let Some(sender) = sender { let end_of_stream = stream_data.end_of_stream; let stream_id = stream_data.stream_id.clone(); - if sender.send(stream_data).await.is_ok() && end_of_stream { + // Close the stream when it is the last frame, or when the receiver is + // gone (send failed), so the mapping is never leaked in the registry. + let send_failed = sender.send(stream_data).await.is_err(); + if end_of_stream || send_failed { self.close_stream(&stream_id).await; } } @@ -253,7 +256,8 @@ async fn run_process_inner( await_pump_task(stderr_task, process_id, "stderr").await; let exit_code = status.code().unwrap_or(-1); - if stdin_closed_from_end_of_stream && exit_code == PWSH_STDIN_CLOSED_EXIT_CODE { + let expected_pwsh_exit = stdin_closed_from_end_of_stream && exit_code == PWSH_STDIN_CLOSED_EXIT_CODE; + if expected_pwsh_exit { info!( process_id, exit_code, "PSU gRPC child process exited with expected code after stdin EOF for pwsh -s" @@ -262,14 +266,25 @@ async fn run_process_inner( info!(process_id, exit_code, "PSU gRPC child process exited"); } + // Reflect the actual outcome so the server can distinguish success from + // cancellation or a non-zero exit based on the StreamClosed message. + let stream_error = canceled || (exit_code != 0 && !expected_pwsh_exit); + let stream_reason = if canceled { + "child process canceled".to_owned() + } else if stream_error { + format!("child process exited with code {exit_code}") + } else { + "child process completed".to_owned() + }; + let _ = outgoing_tx .send(agent_message( &agent_id, &connection_id, AgentPayload::StreamClosed(stream_closed( request.stream_id.clone(), - "child process completed".to_owned(), - false, + stream_reason, + stream_error, )), )) .await; From 2f35d42f26d661eaa01d26f6aad59310bb4a7699 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20CORTIER?= Date: Thu, 2 Jul 2026 20:51:47 +0900 Subject: [PATCH 4/4] . --- devolutions-agent/src/psu_grpc_agent/process.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/devolutions-agent/src/psu_grpc_agent/process.rs b/devolutions-agent/src/psu_grpc_agent/process.rs index 7e84ab031..cd61433cb 100644 --- a/devolutions-agent/src/psu_grpc_agent/process.rs +++ b/devolutions-agent/src/psu_grpc_agent/process.rs @@ -281,11 +281,7 @@ async fn run_process_inner( .send(agent_message( &agent_id, &connection_id, - AgentPayload::StreamClosed(stream_closed( - request.stream_id.clone(), - stream_reason, - stream_error, - )), + AgentPayload::StreamClosed(stream_closed(request.stream_id.clone(), stream_reason, stream_error)), )) .await;