secrets/crates/desktop-daemon/src/exec.rs

use std::collections::BTreeMap;
use std::time::Duration;

use anyhow::{Context, Result, anyhow};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use tokio::process::Command;

use crate::target::{ExecutionTarget, ResolvedTarget};

const MAX_OUTPUT_CHARS: usize = 64 * 1024;

#[derive(Clone, Debug, Deserialize)]
pub struct TargetExecInput {
    pub target_ref: Option<String>,
    pub command: String,
    pub timeout_secs: Option<u64>,
    pub working_dir: Option<String>,
    pub env_overrides: Option<Map<String, Value>>,
}

#[derive(Clone, Debug, Serialize)]
pub struct ExecResult {
    pub resolved_target: ResolvedTarget,
    pub resolved_env_keys: Vec<String>,
    pub command: String,
    pub exit_code: Option<i32>,
    pub stdout: String,
    pub stderr: String,
    pub timed_out: bool,
    pub duration_ms: u128,
    pub stdout_truncated: bool,
    pub stderr_truncated: bool,
}

fn truncate_output(text: String) -> (String, bool) {
    if text.chars().count() <= MAX_OUTPUT_CHARS {
        return (text, false);
    }
    let truncated = text.chars().take(MAX_OUTPUT_CHARS).collect::<String>();
    (truncated, true)
}

fn stringify_env_override(value: &Value) -> Option<String> {
    match value {
        Value::Null => None,
        Value::String(s) => Some(s.clone()),
        Value::Bool(v) => Some(v.to_string()),
        Value::Number(v) => Some(v.to_string()),
        other => serde_json::to_string(other).ok(),
    }
}

fn apply_env_overrides(
    env: &mut BTreeMap<String, String>,
    overrides: Option<&Map<String, Value>>,
) -> Result<()> {
    let Some(overrides) = overrides else {
        return Ok(());
    };
    for (key, value) in overrides {
        if key.is_empty() || key.contains('=') {
            return Err(anyhow!("invalid env override key: {key}"));
        }
        if key.starts_with("TARGET_") {
            return Err(anyhow!(
                "env override `{key}` cannot override reserved TARGET_* variables"
            ));
        }
        if let Some(value) = stringify_env_override(value) {
            env.insert(key.clone(), value);
        }
    }
    Ok(())
}

pub async fn execute_command(
    input: &TargetExecInput,
    target: &ExecutionTarget,
    timeout_secs: u64,
) -> Result<ExecResult> {
    let mut env = target.env.clone();
    apply_env_overrides(&mut env, input.env_overrides.as_ref())?;

    let started = std::time::Instant::now();
    let mut command = Command::new("/bin/sh");
    command
        .arg("-lc")
        .arg(&input.command)
        .kill_on_drop(true)
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped());

    if let Some(dir) = input.working_dir.as_ref().filter(|dir| !dir.is_empty()) {
        command.current_dir(dir);
    }
    for (key, value) in &env {
        command.env(key, value);
    }

    let child = command
        .spawn()
        .with_context(|| format!("failed to spawn command: {}", input.command))?;

    let timed = tokio::time::timeout(
        Duration::from_secs(timeout_secs.clamp(1, 86400)),
        child.wait_with_output(),
    )
    .await;

    let (exit_code, stdout, stderr, timed_out) = match timed {
        Ok(output) => {
            let output = output.context("failed waiting for command output")?;
            (
                output.status.code(),
                String::from_utf8_lossy(&output.stdout).to_string(),
                String::from_utf8_lossy(&output.stderr).to_string(),
                false,
            )
        }
        Err(_) => (None, String::new(), "command timed out".to_string(), true),
    };

    let (stdout, stdout_truncated) = truncate_output(stdout);
    let (stderr, stderr_truncated) = truncate_output(stderr);

    Ok(ExecResult {
        resolved_target: target.resolved.clone(),
        resolved_env_keys: target.resolved_env_keys(),
        command: input.command.clone(),
        exit_code,
        stdout,
        stderr,
        timed_out,
        duration_ms: started.elapsed().as_millis(),
        stdout_truncated,
        stderr_truncated,
    })
}