clusterflux-public/crates/clusterflux-node/src/assignment_runner/process_runner.rs

use super::*;
use std::sync::mpsc::{self, Receiver, SyncSender};

struct LiveLogChunk {
    stream: &'static str,
    offset: u64,
    source_bytes: u64,
    bytes: Vec<u8>,
    truncated: bool,
}

fn redact_safe_live_log_prefix(
    pending: &[u8],
    configured_secrets: &[String],
    final_chunk: bool,
) -> Option<(usize, String)> {
    if pending.is_empty() {
        return None;
    }
    let secret_bytes = configured_secrets
        .iter()
        .filter(|secret| secret.len() >= 4)
        .map(String::as_bytes)
        .collect::<Vec<_>>();
    let maximum_secret_bytes = secret_bytes
        .iter()
        .map(|secret| secret.len())
        .max()
        .unwrap_or(0);
    if !final_chunk && pending.len() <= maximum_secret_bytes {
        return None;
    }
    let mut consumed = if final_chunk {
        pending.len()
    } else {
        pending.len() - maximum_secret_bytes
    };
    loop {
        let previous = consumed;
        for secret in &secret_bytes {
            for start in 0..=pending.len().saturating_sub(secret.len()) {
                let end = start + secret.len();
                if start < consumed && end > consumed && &pending[start..end] == *secret {
                    consumed = end;
                }
            }
        }
        if consumed == previous {
            break;
        }
    }
    if consumed == 0 {
        return None;
    }
    let text = redact_configured_values(
        String::from_utf8_lossy(&pending[..consumed]).into_owned(),
        configured_secrets,
    );
    Some((consumed, text))
}

pub(super) struct CoordinatorControlledProcessRunner {
    pub(super) args: Args,
    pub(super) process: String,
    pub(super) task: String,
    pub(super) node_private_key: String,
    pub(super) debug_control: Arc<WasmDebugControl>,
    pub(super) command_status: Arc<Mutex<Option<String>>>,
    pub(super) stdout_source_bytes: Arc<AtomicU64>,
    pub(super) stderr_source_bytes: Arc<AtomicU64>,
    pub(super) timeout: Duration,
    pub(super) configured_secrets: Vec<String>,
}

impl CoordinatorControlledProcessRunner {
    const MAX_CAPTURE_BYTES: usize = 256 * 1024 + 1;

    pub(super) fn new(
        host: &CoordinatorWasmTaskHost,
        timeout: Duration,
        configured_secrets: Vec<String>,
    ) -> Self {
        Self {
            args: host.args.clone(),
            process: host.process.clone(),
            task: host.parent_task.clone(),
            node_private_key: host.node_private_key.clone(),
            debug_control: Arc::clone(&host.debug_control),
            command_status: Arc::clone(&host.command_status),
            stdout_source_bytes: Arc::clone(&host.command_stdout_source_bytes),
            stderr_source_bytes: Arc::clone(&host.command_stderr_source_bytes),
            timeout,
            configured_secrets,
        }
    }

    fn set_command_status(&self, status: impl Into<String>) {
        if let Ok(mut current) = self.command_status.lock() {
            *current = Some(status.into());
        }
    }

    fn abort_requested(&self, session: &mut CoordinatorSession) -> Result<bool, BackendError> {
        let request = signed_node_request_json(
            &self.args,
            &self.node_private_key,
            "poll_task_control",
            serde_json::json!({
                "type": "poll_task_control",
                "tenant": &self.args.tenant,
                "project": &self.args.project,
                "process": &self.process,
                "node": &self.args.node,
                "task": &self.task,
            }),
        )
        .map_err(|error| BackendError::Command(error.to_string()))?;
        let response = session
            .request(request)
            .map_err(|error| BackendError::Command(format!("poll task control: {error}")))?;
        Ok(response
            .get("abort_requested")
            .and_then(serde_json::Value::as_bool)
            .unwrap_or(false))
    }

    fn terminate_process_group(child: &mut std::process::Child) {
        #[cfg(unix)]
        {
            let process_group = -(child.id() as i32);
            // The child is placed in a new process group before exec, so this reaches
            // all native descendants. Podman containers are removed separately.
            unsafe {
                libc::kill(process_group, libc::SIGKILL);
            }
        }
        let _ = child.kill();
        let _ = child.wait();
    }

    fn podman_container_name(command: &PodmanCommand) -> Option<String> {
        if command.program != "podman" || command.args.first().map(String::as_str) != Some("run") {
            return None;
        }
        command
            .args
            .windows(2)
            .find(|arguments| arguments[0] == "--name")
            .map(|arguments| arguments[1].clone())
    }

    fn set_podman_paused(container: &str, paused: bool) -> Result<(), BackendError> {
        let action = if paused { "pause" } else { "unpause" };
        let output = std::process::Command::new("podman")
            .args([action, container])
            .output()
            .map_err(|error| {
                BackendError::Command(format!(
                    "failed to invoke `podman {action}` for container `{container}`: {error}"
                ))
            })?;
        if !output.status.success() {
            return Err(BackendError::Command(format!(
                "`podman {action}` failed for container `{container}`: {}",
                String::from_utf8_lossy(&output.stderr).trim()
            )));
        }

        let inspection = std::process::Command::new("podman")
            .args(["inspect", "--format", "{{.State.Paused}}", container])
            .output()
            .map_err(|error| {
                BackendError::Command(format!(
                    "failed to inspect Podman pause state for container `{container}`: {error}"
                ))
            })?;
        let observed = String::from_utf8_lossy(&inspection.stdout);
        let expected = if paused { "true" } else { "false" };
        if !inspection.status.success() || observed.trim() != expected {
            return Err(BackendError::Command(format!(
                "container `{container}` did not verify as {} after `podman {action}`: status={:?} stdout={} stderr={}",
                if paused { "paused" } else { "running" },
                inspection.status.code(),
                observed.trim(),
                String::from_utf8_lossy(&inspection.stderr).trim()
            )));
        }
        Ok(())
    }

    fn terminate_execution(child: &mut std::process::Child, container: Option<&str>) {
        if let Some(container) = container {
            let _ = std::process::Command::new("podman")
                .args(["rm", "--force", container])
                .output();
        }
        Self::terminate_process_group(child);
    }

    fn freeze_execution(
        child: &std::process::Child,
        container: Option<&str>,
    ) -> Result<(), BackendError> {
        if let Some(container) = container {
            return Self::set_podman_paused(container, true);
        }
        #[cfg(unix)]
        {
            Self::freeze_process_group(child)
        }
        #[cfg(not(unix))]
        {
            let _ = child;
            Err(BackendError::Command(
                "native debug freeze requires Unix process groups".to_owned(),
            ))
        }
    }

    fn resume_execution(
        child: &std::process::Child,
        container: Option<&str>,
    ) -> Result<(), BackendError> {
        if let Some(container) = container {
            return Self::set_podman_paused(container, false);
        }
        #[cfg(unix)]
        {
            Self::resume_process_group(child)
        }
        #[cfg(not(unix))]
        {
            let _ = child;
            Err(BackendError::Command(
                "native debug resume requires Unix process groups".to_owned(),
            ))
        }
    }

    #[cfg(unix)]
    fn freeze_process_group(child: &std::process::Child) -> Result<(), BackendError> {
        let process_group = -(child.id() as i32);
        let result = unsafe { libc::kill(process_group, libc::SIGSTOP) };
        if result == 0 {
            Ok(())
        } else {
            Err(BackendError::Command(format!(
                "failed to freeze native process group {}: {}",
                child.id(),
                std::io::Error::last_os_error()
            )))
        }
    }

    #[cfg(unix)]
    fn resume_process_group(child: &std::process::Child) -> Result<(), BackendError> {
        let process_group = -(child.id() as i32);
        let result = unsafe { libc::kill(process_group, libc::SIGCONT) };
        if result == 0 {
            Ok(())
        } else {
            Err(BackendError::Command(format!(
                "failed to resume native process group {}: {}",
                child.id(),
                std::io::Error::last_os_error()
            )))
        }
    }

    fn drain_bounded(
        mut reader: impl Read + Send + 'static,
        maximum: usize,
        stream: &'static str,
        sender: SyncSender<LiveLogChunk>,
        configured_secrets: Vec<String>,
        source_bytes_total: Arc<AtomicU64>,
    ) -> thread::JoinHandle<Result<Vec<u8>, String>> {
        thread::spawn(move || {
            let mut captured = Vec::new();
            let mut buffer = [0_u8; 16 * 1024];
            let stream_base = source_bytes_total.load(Ordering::Relaxed);
            let mut pending_offset = stream_base;
            let mut pending = Vec::new();
            let mut discarded = false;
            loop {
                let count = reader
                    .read(&mut buffer)
                    .map_err(|error| error.to_string())?;
                if count == 0 {
                    break;
                }
                let _ = source_bytes_total.fetch_update(
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                    |current| Some(current.saturating_add(count as u64)),
                );
                let remaining = maximum.saturating_sub(captured.len());
                let retained = count.min(remaining);
                if retained > 0 {
                    captured.extend_from_slice(&buffer[..retained]);
                    pending.extend_from_slice(&buffer[..retained]);
                    if let Some((consumed, text)) =
                        redact_safe_live_log_prefix(&pending, &configured_secrets, false)
                    {
                        let _ = sender.try_send(LiveLogChunk {
                            stream,
                            offset: pending_offset,
                            source_bytes: consumed as u64,
                            bytes: text.into_bytes(),
                            truncated: false,
                        });
                        pending.drain(..consumed);
                        pending_offset = pending_offset.saturating_add(consumed as u64);
                    }
                }
                if retained < count {
                    discarded = true;
                }
            }
            if let Some((consumed, text)) =
                redact_safe_live_log_prefix(&pending, &configured_secrets, true)
            {
                let _ = sender.try_send(LiveLogChunk {
                    stream,
                    offset: pending_offset,
                    source_bytes: consumed as u64,
                    bytes: text.into_bytes(),
                    truncated: false,
                });
            }
            if discarded {
                let _ = sender.try_send(LiveLogChunk {
                    stream,
                    offset: stream_base.saturating_add(maximum as u64),
                    source_bytes: 0,
                    bytes: b"[log output truncated at node capture limit]".to_vec(),
                    truncated: true,
                });
            }
            Ok(captured)
        })
    }

    fn spawn_live_log_reporter(&self, receiver: Receiver<LiveLogChunk>) -> thread::JoinHandle<()> {
        let args = self.args.clone();
        let process = self.process.clone();
        let task = self.task.clone();
        let node_private_key = self.node_private_key.clone();
        let configured_secrets = self.configured_secrets.clone();
        let command_status = Arc::clone(&self.command_status);
        thread::spawn(move || {
            let mut log_session = None;
            let mut delivery_available = true;
            while let Ok(chunk) = receiver.recv() {
                if !delivery_available {
                    continue;
                }
                let mut text = String::from_utf8_lossy(&chunk.bytes).into_owned();
                text = redact_configured_values(text, &configured_secrets);
                let mut delivered = false;
                for _ in 0..2 {
                    if log_session.is_none() {
                        log_session = CoordinatorSession::connect_with_timeouts(
                            &args.coordinator,
                            Duration::from_millis(500),
                            Duration::from_millis(500),
                        )
                        .ok();
                    }
                    let Some(session) = log_session.as_mut() else {
                        continue;
                    };
                    let request = signed_node_request_json(
                        &args,
                        &node_private_key,
                        "report_task_log_chunk",
                        serde_json::json!({
                            "type": "report_task_log_chunk",
                            "tenant": &args.tenant,
                            "project": &args.project,
                            "process": &process,
                            "node": &args.node,
                            "task": &task,
                            "stream": chunk.stream,
                            "offset": chunk.offset,
                            "source_bytes": chunk.source_bytes,
                            "text": &text,
                            "truncated": chunk.truncated,
                        }),
                    );
                    let result = match request {
                        Ok(request) => session
                            .request(request)
                            .map(|_| ())
                            .map_err(|error| error.to_string()),
                        Err(error) => Err(error.to_string()),
                    };
                    match result {
                        Ok(()) => {
                            delivered = true;
                            break;
                        }
                        Err(_) => {
                            log_session = None;
                        }
                    }
                }
                if !delivered {
                    delivery_available = false;
                    if let Ok(mut current) = command_status.lock() {
                        *current = Some(
                            "live log delivery was interrupted; final bounded output remains available"
                                .to_owned(),
                        );
                    }
                }
            }
        })
    }
}

impl ProcessRunner for CoordinatorControlledProcessRunner {
    fn run(&mut self, command: &PodmanCommand) -> Result<ProcessOutput, BackendError> {
        let podman_container = Self::podman_container_name(command);
        self.set_command_status(format!(
            "starting native command: {} {}",
            command.program,
            command.args.join(" ")
        ));
        let mut process = std::process::Command::new(&command.program);
        process
            .args(&command.args)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        #[cfg(unix)]
        process.process_group(0);
        let mut child = process
            .spawn()
            .map_err(|error| BackendError::Command(error.to_string()))?;
        self.set_command_status(format!(
            "running native command pid {}: {} {}",
            child.id(),
            command.program,
            command.args.join(" ")
        ));
        let (live_log_sender, live_log_receiver) = mpsc::sync_channel(64);
        let stdout = Self::drain_bounded(
            child
                .stdout
                .take()
                .ok_or_else(|| BackendError::Command("command stdout pipe missing".to_owned()))?,
            Self::MAX_CAPTURE_BYTES,
            "stdout",
            live_log_sender.clone(),
            self.configured_secrets.clone(),
            Arc::clone(&self.stdout_source_bytes),
        );
        let stderr = Self::drain_bounded(
            child
                .stderr
                .take()
                .ok_or_else(|| BackendError::Command("command stderr pipe missing".to_owned()))?,
            Self::MAX_CAPTURE_BYTES,
            "stderr",
            live_log_sender,
            self.configured_secrets.clone(),
            Arc::clone(&self.stderr_source_bytes),
        );
        let live_log_reporter = self.spawn_live_log_reporter(live_log_receiver);
        let mut session = match CoordinatorSession::connect(&self.args.coordinator) {
            Ok(session) => session,
            Err(error) => {
                Self::terminate_execution(&mut child, podman_container.as_deref());
                let _ = stdout.join();
                let _ = stderr.join();
                let _ = live_log_reporter.join();
                return Err(BackendError::Command(format!(
                    "establish execution control channel: {error}"
                )));
            }
        };

        let mut frozen_epoch = None;
        let started = Instant::now();
        let status = loop {
            match child.try_wait() {
                Ok(Some(status)) => break status,
                Ok(None) => {}
                Err(error) => {
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    let _ = stdout.join();
                    let _ = stderr.join();
                    let _ = live_log_reporter.join();
                    return Err(BackendError::Command(error.to_string()));
                }
            }
            if started.elapsed() >= self.timeout {
                self.set_command_status(format!(
                    "native command pid {} exceeded wall-clock timeout of {} ms",
                    child.id(),
                    self.timeout.as_millis()
                ));
                Self::terminate_execution(&mut child, podman_container.as_deref());
                let _ = stdout.join();
                let _ = stderr.join();
                let _ = live_log_reporter.join();
                return Err(BackendError::Command(format!(
                    "native command exceeded wall-clock timeout of {} ms",
                    self.timeout.as_millis()
                )));
            }
            match self.abort_requested(&mut session) {
                Ok(true) => {
                    self.set_command_status(format!(
                        "aborting native command pid {} at coordinator request",
                        child.id()
                    ));
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    let _ = stdout.join();
                    let _ = stderr.join();
                    let _ = live_log_reporter.join();
                    return Err(BackendError::Cancelled(
                        "coordinator requested cancellation or abort".to_owned(),
                    ));
                }
                Ok(false) => {}
                Err(error) => {
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    let _ = stdout.join();
                    let _ = stderr.join();
                    let _ = live_log_reporter.join();
                    return Err(error);
                }
            }
            if let Some(epoch) = self.debug_control.requested_epoch() {
                if self.debug_control.resume_requested(epoch) {
                    if frozen_epoch == Some(epoch) {
                        match Self::resume_execution(&child, podman_container.as_deref()) {
                            Ok(()) => {
                                self.set_command_status(format!(
                                    "running command pid {} after debug epoch {epoch} resumed",
                                    child.id()
                                ));
                                self.debug_control.mark_running(epoch);
                                frozen_epoch = None;
                            }
                            Err(error) => self.set_command_status(format!(
                                "debug epoch {epoch} resume is pending: {error}"
                            )),
                        }
                    }
                } else if frozen_epoch != Some(epoch)
                    && self.debug_control.frozen_epoch() != Some(epoch)
                {
                    match Self::freeze_execution(&child, podman_container.as_deref()) {
                        Ok(()) => {
                            self.set_command_status(format!(
                                "frozen command pid {} for debug epoch {epoch}",
                                child.id()
                            ));
                            self.debug_control.mark_frozen(epoch);
                            frozen_epoch = Some(epoch);
                        }
                        Err(error) => self.set_command_status(format!(
                            "debug epoch {epoch} freeze is pending: {error}"
                        )),
                    }
                }
            }
            thread::sleep(Duration::from_millis(50));
        };
        let stdout = stdout
            .join()
            .map_err(|_| BackendError::Command("stdout reader panicked".to_owned()))?
            .map_err(BackendError::Command)?;
        let stderr = stderr
            .join()
            .map_err(|_| BackendError::Command("stderr reader panicked".to_owned()))?
            .map_err(BackendError::Command)?;
        live_log_reporter
            .join()
            .map_err(|_| BackendError::Command("live log reporter panicked".to_owned()))?;
        self.set_command_status(format!(
            "native command exited with status {:?}",
            status.code()
        ));
        Ok(ProcessOutput {
            status_code: status.code(),
            stdout,
            stderr,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{redact_safe_live_log_prefix, CoordinatorControlledProcessRunner};
    use std::io::Cursor;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::sync::{mpsc, Arc};

    #[test]
    fn live_log_redaction_holds_boundaries_until_split_secrets_are_complete() {
        let secrets = vec!["correct-horse".to_owned()];
        let mut pending = b"prefix correct-".to_vec();
        let (consumed, first) = redact_safe_live_log_prefix(&pending, &secrets, false).unwrap();
        pending.drain(..consumed);
        pending.extend_from_slice(b"horse suffix");
        let (_, second) = redact_safe_live_log_prefix(&pending, &secrets, true).unwrap();

        let combined = format!("{first}{second}");
        assert_eq!(combined, "prefix [REDACTED] suffix");
        assert!(!combined.contains("correct-"));
        assert!(!combined.contains("horse"));
    }

    #[test]
    fn bounded_capture_retains_the_complete_source_byte_count() {
        let source_bytes = Arc::new(AtomicU64::new(5));
        let (sender, receiver) = mpsc::sync_channel(8);
        let reader = CoordinatorControlledProcessRunner::drain_bounded(
            Cursor::new(vec![b'x'; 32]),
            8,
            "stdout",
            sender,
            Vec::new(),
            Arc::clone(&source_bytes),
        );
        let captured = reader.join().unwrap().unwrap();
        let chunks = receiver.into_iter().collect::<Vec<_>>();

        assert_eq!(captured, vec![b'x'; 8]);
        assert_eq!(source_bytes.load(Ordering::Relaxed), 37);
        assert_eq!(
            chunks.iter().map(|chunk| chunk.source_bytes).sum::<u64>(),
            8
        );
        assert_eq!(chunks[0].offset, 5);
        assert_eq!(chunks.last().unwrap().offset, 13);
        assert!(chunks.iter().any(|chunk| chunk.truncated));
    }
}