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:?}"), + } + } +}