a3s-box-runtime 3.1.0

//! Generation-fenced command, PTY, and file sessions.

use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;

use a3s_box_core::pty::PtyRequest;
use a3s_box_core::{
    BoxError, ExecEvent, ExecOutput, ExecRequest, ExecutionGeneration, ExecutionId,
    ExecutionManagerError, ExecutionManagerResult, ExecutionProcess, ExecutionProcessInput,
    ExecutionProcessSignal, ExecutionProcessStream, ExecutionSessionManager, FileRequest,
    FileResponse, FilesystemRequest, FilesystemResponse,
};
use async_trait::async_trait;

use super::LocalExecutionManager;
use crate::{
    BoxRecord, ExecClient, PtyClient, StreamingExec, StreamingExecInput, StreamingPty,
    StreamingPtyInput,
};

#[async_trait]
impl ExecutionSessionManager for LocalExecutionManager {
    async fn execute(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
        mut request: ExecRequest,
    ) -> ExecutionManagerResult<ExecOutput> {
        request.streaming = false;
        let (record, client, stream) = self.bind_exec(execution_id, generation).await?;
        inherit_container_environment(&record.env, &mut request.env);
        debug_session_environment(
            execution_id,
            generation,
            "execute",
            &record.env,
            &request.env,
        );
        client
            .exec_command_on_stream(stream, &request)
            .await
            .map_err(|error| session_error(execution_id, "execute command", error))
    }

    async fn start_process(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
        mut request: ExecRequest,
    ) -> ExecutionManagerResult<ExecutionProcess> {
        let (record, client, stream) = self.bind_exec(execution_id, generation).await?;
        inherit_container_environment(&record.env, &mut request.env);
        debug_session_environment(
            execution_id,
            generation,
            "start_process",
            &record.env,
            &request.env,
        );
        let stream = client
            .exec_stream_on_stream(stream, &request)
            .await
            .map_err(|error| session_error(execution_id, "start command", error))?;
        let input: Arc<dyn ExecutionProcessInput> = Arc::new(ExecInput {
            execution_id: execution_id.clone(),
            input: stream.input(),
        });
        Ok(Box::new(ExecStream { stream, input }))
    }

    async fn start_pty(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
        mut request: PtyRequest,
    ) -> ExecutionManagerResult<ExecutionProcess> {
        let record = self
            .require_running_record(execution_id, generation)
            .await?;
        inherit_container_environment(&record.env, &mut request.env);
        debug_session_environment(
            execution_id,
            generation,
            "start_pty",
            &record.env,
            &request.env,
        );
        let socket_path = record.exec_socket_path.with_file_name("pty.sock");
        let client = PtyClient::connect(&socket_path)
            .await
            .map_err(|error| session_error(execution_id, "connect PTY", error))?;
        self.require_same_runtime(&record, execution_id, generation)
            .await?;
        let stream = client
            .start_stream(&request)
            .await
            .map_err(|error| session_error(execution_id, "start PTY", error))?;
        let input: Arc<dyn ExecutionProcessInput> = Arc::new(PtyInput {
            execution_id: execution_id.clone(),
            input: stream.input(),
        });
        Ok(Box::new(PtyStream { stream, input }))
    }

    async fn transfer_file(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
        request: FileRequest,
    ) -> ExecutionManagerResult<FileResponse> {
        let (_record, client, stream) = self.bind_exec(execution_id, generation).await?;
        client
            .file_transfer_on_stream(stream, &request)
            .await
            .map_err(|error| session_error(execution_id, "transfer file", error))
    }

    async fn filesystem(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
        request: FilesystemRequest,
    ) -> ExecutionManagerResult<FilesystemResponse> {
        let (_record, client, stream) = self.bind_exec(execution_id, generation).await?;
        client
            .filesystem_on_stream(stream, &request)
            .await
            .map_err(|error| session_error(execution_id, "access filesystem", error))
    }
}

impl LocalExecutionManager {
    async fn bind_exec(
        &self,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
    ) -> ExecutionManagerResult<(BoxRecord, ExecClient, tokio::net::UnixStream)> {
        let record = self
            .require_running_record(execution_id, generation)
            .await?;
        let client = ExecClient::for_socket(&record.exec_socket_path);
        let stream = client
            .open_stream()
            .await
            .map_err(|error| session_error(execution_id, "connect exec", error))?;
        self.require_same_runtime(&record, execution_id, generation)
            .await?;
        Ok((record, client, stream))
    }

    async fn require_same_runtime(
        &self,
        bound: &BoxRecord,
        execution_id: &ExecutionId,
        generation: ExecutionGeneration,
    ) -> ExecutionManagerResult<()> {
        let current = self
            .require_running_record(execution_id, generation)
            .await?;
        if current.pid != bound.pid
            || current.pid_start_time != bound.pid_start_time
            || current.exec_socket_path != bound.exec_socket_path
        {
            return Err(ExecutionManagerError::Conflict {
                execution_id: execution_id.clone(),
                message: "runtime generation changed while binding its execution session"
                    .to_string(),
            });
        }
        Ok(())
    }
}

/// Merge the environment fixed at Sandbox creation into an exec/PTY request.
///
/// The guest normally inherits these values from guest-init, but the execution
/// contract must not depend on which user a later request selects. Per-request
/// values remain authoritative, matching OCI/Docker exec environment semantics.
fn inherit_container_environment(container: &HashMap<String, String>, request: &mut Vec<String>) {
    let mut merged: BTreeMap<String, String> = container
        .iter()
        .map(|(key, value)| (key.clone(), value.clone()))
        .collect();
    let mut malformed = Vec::new();
    for entry in std::mem::take(request) {
        if let Some((key, value)) = entry.split_once('=') {
            merged.insert(key.to_string(), value.to_string());
        } else {
            malformed.push(entry);
        }
    }
    request.extend(
        merged
            .into_iter()
            .map(|(key, value)| format!("{key}={value}")),
    );
    request.extend(malformed);
}

/// Record only environment key names at the final host-to-runtime boundary.
///
/// Values are deliberately never included because creation and per-request
/// environments may contain credentials. This log is useful when diagnosing a
/// persisted lifecycle record that reaches a later exec or PTY generation.
fn debug_session_environment(
    execution_id: &ExecutionId,
    generation: ExecutionGeneration,
    operation: &str,
    container: &HashMap<String, String>,
    request: &[String],
) {
    let mut container_keys: Vec<&str> = container.keys().map(String::as_str).collect();
    container_keys.sort_unstable();
    let mut request_keys: Vec<&str> = request
        .iter()
        .filter_map(|entry| entry.split_once('=').map(|(key, _)| key))
        .collect();
    request_keys.sort_unstable();
    request_keys.dedup();
    let malformed_request_entries = request.iter().filter(|entry| !entry.contains('=')).count();

    tracing::debug!(
        %execution_id,
        generation = generation.get(),
        operation,
        container_env_count = container.len(),
        container_env_keys = ?container_keys,
        merged_request_env_count = request.len(),
        merged_request_env_keys = ?request_keys,
        malformed_request_entries,
        "Prepared managed execution session environment"
    );
}

struct ExecInput {
    execution_id: ExecutionId,
    input: StreamingExecInput,
}

#[async_trait]
impl ExecutionProcessInput for ExecInput {
    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
        self.input
            .write_stdin(data)
            .await
            .map_err(|error| session_error(&self.execution_id, "write command stdin", error))
    }

    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
        self.input
            .close_stdin()
            .await
            .map_err(|error| session_error(&self.execution_id, "close command stdin", error))
    }

    async fn cancel(&self) -> ExecutionManagerResult<()> {
        self.input
            .cancel()
            .await
            .map_err(|error| session_error(&self.execution_id, "cancel command", error))
    }

    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
        self.input
            .send_signal(signal)
            .await
            .map_err(|error| session_error(&self.execution_id, "signal command", error))
    }
}

struct ExecStream {
    stream: StreamingExec,
    input: Arc<dyn ExecutionProcessInput>,
}

#[async_trait]
impl ExecutionProcessStream for ExecStream {
    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
        self.input.clone()
    }

    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
        self.stream
            .next_event()
            .await
            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
    }
}

struct PtyInput {
    execution_id: ExecutionId,
    input: StreamingPtyInput,
}

#[async_trait]
impl ExecutionProcessInput for PtyInput {
    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
        self.input
            .write_stdin(data)
            .await
            .map_err(|error| session_error(&self.execution_id, "write PTY stdin", error))
    }

    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
        self.cancel().await
    }

    async fn cancel(&self) -> ExecutionManagerResult<()> {
        self.input
            .close()
            .await
            .map_err(|error| session_error(&self.execution_id, "close PTY", error))
    }

    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
        self.input
            .send_signal(signal)
            .await
            .map_err(|error| session_error(&self.execution_id, "signal PTY", error))
    }

    async fn resize_pty(&self, cols: u16, rows: u16) -> ExecutionManagerResult<()> {
        self.input
            .resize(cols, rows)
            .await
            .map_err(|error| session_error(&self.execution_id, "resize PTY", error))
    }
}

struct PtyStream {
    stream: StreamingPty,
    input: Arc<dyn ExecutionProcessInput>,
}

#[async_trait]
impl ExecutionProcessStream for PtyStream {
    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
        self.input.clone()
    }

    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
        self.stream
            .next_event()
            .await
            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
    }
}

fn session_error(
    execution_id: &ExecutionId,
    operation: &str,
    error: BoxError,
) -> ExecutionManagerError {
    ExecutionManagerError::Unavailable(format!(
        "failed to {operation} for execution {execution_id}: {error}"
    ))
}

#[cfg(test)]
mod tests {
    use super::inherit_container_environment;
    use std::collections::HashMap;

    #[test]
    fn request_environment_overrides_inherited_container_values() {
        let container = HashMap::from([
            ("ALPHA".to_string(), "container".to_string()),
            ("BETA".to_string(), "container".to_string()),
        ]);
        let mut request = vec!["BETA=request".to_string(), "GAMMA=request".to_string()];

        inherit_container_environment(&container, &mut request);

        assert_eq!(
            request,
            ["ALPHA=container", "BETA=request", "GAMMA=request"]
        );
    }

    #[test]
    fn malformed_request_entries_are_preserved_after_inherited_values() {
        let container = HashMap::from([("ALPHA".to_string(), "container".to_string())]);
        let mut request = vec!["MALFORMED".to_string()];

        inherit_container_environment(&container, &mut request);

        assert_eq!(request, ["ALPHA=container", "MALFORMED"]);
    }
}