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..e1f711b6c --- /dev/null +++ b/devolutions-agent/Dockerfile.psu-grpc-poc @@ -0,0 +1,38 @@ +FROM rust:1.90-bookworm AS build + +WORKDIR /workspace + +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 \ + && 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 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 + +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 new file mode 100644 index 000000000..974d798a4 --- /dev/null +++ b/devolutions-agent/Run-PsuGrpcAgentContainer.ps1 @@ -0,0 +1,46 @@ +[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] $AppToken, + [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 +} + +$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/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..ce6f49ea0 --- /dev/null +++ b/devolutions-agent/docker/psu-grpc-entrypoint.sh @@ -0,0 +1,56 @@ +#!/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 + +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" < 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..1dc90f73b 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,60 @@ 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, + + /// 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, + + /// 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, + app_token: None, + hubs: Vec::new(), + powershell: PsuPowerShellConf::default(), + } + } + } + #[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "PascalCase")] pub struct PsuPowerShellConf { @@ -685,6 +741,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 +776,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 +998,32 @@ 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", + "AppToken": "app-token", + "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.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/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..aa1819355 --- /dev/null +++ b/devolutions-agent/src/psu_grpc_agent/mod.rs @@ -0,0 +1,440 @@ +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; +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, clippy::clone_on_ref_ptr, clippy::similar_names)] +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, + app_token: Option, + powershell_executable: String, +} + +impl PsuGrpcAgent { + fn new(conf: dto::PsuGrpcAgentConf) -> anyhow::Result { + let server_url = conf + .server_url + .as_ref() + .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()); + 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(); + + Ok(Self { + conf, + server_url, + agent_id, + display_name, + machine_name, + app_token, + 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(connect_request( + ReceiverStream::new(outgoing_rx), + self.app_token.as_deref(), + )?) + .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) = mpsc::channel(8); + 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 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) + .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.to_ascii_lowercase().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(), + } +} + +#[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 new file mode 100644 index 000000000..cd61433cb --- /dev/null +++ b/devolutions-agent/src/psu_grpc_agent/process.rs @@ -0,0 +1,601 @@ +use std::collections::HashMap; +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::{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}; +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: mpsc::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(); + // 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; + } + } + } + + 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 = { + 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.send(kill_process).await; + } + } + + 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, + 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 + } 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 = 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); + + 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 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"); + + let (exit_status, killed) = wait_for_graceful_child_exit(&mut child, process_id).await?; + canceled |= killed; + break exit_status; + } + kill_process = control_rx.recv() => { + match kill_process { + 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")?; + } + 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; + } + None => {} + } + } + } + }; + + 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); + 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" + ); + } else { + 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(), stream_reason, stream_error)), + )) + .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, + AgentPayload::ProcessCompleted(ProcessCompleted { + correlation_id: correlation_id.to_owned(), + exit_code, + canceled, + error_message, + }), + )) + .await + .context("failed to send PSU gRPC ProcessCompleted message") +} + +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") +} + +#[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:?}"), + } + } +} 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 {