agentmesh 3.7.0

// Copyright (c) Microsoft Corporation. Licensed under the MIT License.

//! Sandbox provider trait and Docker-based implementation.
//!
//! Defines the backend-agnostic API for sandboxed agent execution. Any
//! sandbox backend — Docker, Hyperlight micro-VMs, cloud sandbox services,
//! or custom providers — implements the [`SandboxProvider`] trait.
//!
//! The [`DockerSandboxProvider`] uses the Docker CLI (`std::process::Command`)
//! to manage hardened containers with dropped capabilities, read-only
//! filesystems, and network isolation.

use std::collections::HashMap;
use std::fmt;
use std::process::Command;
use std::time::Instant;

/// Validate that an environment variable name contains only safe characters.
fn validate_env_key(key: &str) -> Result<(), String> {
    if key.is_empty()
        || !key.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
    {
        return Err(format!("Invalid environment variable name: {:?}", key));
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Enums
// ---------------------------------------------------------------------------

/// Lifecycle state of a sandbox session.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionStatus {
    Provisioning,
    Ready,
    Executing,
    Destroying,
    Destroyed,
    Failed,
}

impl fmt::Display for SessionStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Self::Provisioning => "provisioning",
            Self::Ready => "ready",
            Self::Executing => "executing",
            Self::Destroying => "destroying",
            Self::Destroyed => "destroyed",
            Self::Failed => "failed",
        };
        write!(f, "{}", s)
    }
}

/// State of a single code execution within a session.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecutionStatus {
    Pending,
    Running,
    Completed,
    Cancelled,
    Failed,
}

impl fmt::Display for ExecutionStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Self::Pending => "pending",
            Self::Running => "running",
            Self::Completed => "completed",
            Self::Cancelled => "cancelled",
            Self::Failed => "failed",
        };
        write!(f, "{}", s)
    }
}

// ---------------------------------------------------------------------------
// Data types
// ---------------------------------------------------------------------------

/// Configuration for a sandbox environment.
#[derive(Debug, Clone)]
pub struct SandboxConfig {
    pub timeout_seconds: f64,
    pub memory_mb: u32,
    pub cpu_limit: f64,
    pub network_enabled: bool,
    pub read_only_fs: bool,
    pub env_vars: HashMap<String, String>,
}

impl Default for SandboxConfig {
    fn default() -> Self {
        Self {
            timeout_seconds: 60.0,
            memory_mb: 512,
            cpu_limit: 1.0,
            network_enabled: false,
            read_only_fs: true,
            env_vars: HashMap::new(),
        }
    }
}

/// Result from a sandbox execution.
#[derive(Debug, Clone)]
pub struct SandboxResult {
    pub success: bool,
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub duration_seconds: f64,
    pub killed: bool,
    pub kill_reason: String,
}

impl Default for SandboxResult {
    fn default() -> Self {
        Self {
            success: false,
            exit_code: 0,
            stdout: String::new(),
            stderr: String::new(),
            duration_seconds: 0.0,
            killed: false,
            kill_reason: String::new(),
        }
    }
}

/// Returned by [`SandboxProvider::create_session`] — identifies an active sandbox session.
#[derive(Debug, Clone)]
pub struct SessionHandle {
    pub agent_id: String,
    pub session_id: String,
    pub status: SessionStatus,
}

/// Returned by [`SandboxProvider::execute_code`] — wraps the result of a single execution.
#[derive(Debug, Clone)]
pub struct ExecutionHandle {
    pub execution_id: String,
    pub agent_id: String,
    pub session_id: String,
    pub status: ExecutionStatus,
    pub result: Option<SandboxResult>,
}

// ---------------------------------------------------------------------------
// Trait
// ---------------------------------------------------------------------------

/// Abstract interface for sandbox providers.
///
/// Defines session-based lifecycle methods (`create_session`, `execute_code`,
/// `destroy_session`) and an availability check.
pub trait SandboxProvider {
    /// Provision a sandbox with optional resource constraints.
    fn create_session(
        &mut self,
        agent_id: &str,
        config: Option<&SandboxConfig>,
    ) -> Result<SessionHandle, String>;

    /// Run code inside an existing session.
    fn execute_code(
        &mut self,
        agent_id: &str,
        session_id: &str,
        code: &str,
    ) -> Result<ExecutionHandle, String>;

    /// Tear down the sandbox and release resources.
    fn destroy_session(&mut self, agent_id: &str, session_id: &str) -> Result<(), String>;

    /// Check if this sandbox provider is available.
    fn is_available(&self) -> bool;

    /// Run a raw command in the sandbox (low-level helper).
    ///
    /// The default implementation returns a failure result so that providers
    /// that do not support raw commands behave predictably.
    fn run(
        &mut self,
        agent_id: &str,
        command: &[&str],
        config: Option<&SandboxConfig>,
    ) -> SandboxResult {
        let _ = (agent_id, command, config);
        SandboxResult {
            success: false,
            exit_code: -1,
            stderr: format!(
                "{} run() is not implemented for this provider",
                std::any::type_name::<Self>()
            ),
            ..Default::default()
        }
    }
}

// ---------------------------------------------------------------------------
// Docker implementation
// ---------------------------------------------------------------------------

/// Container name prefix used to namespace sandbox containers.
const CONTAINER_PREFIX: &str = "agentmesh-sandbox";

/// Generate a unique 16-hex-char session/execution ID.
///
/// Uses `rand::random::<u64>()` (OS-seeded thread RNG) rather than mixing
/// the current nanosecond timestamp with `ThreadId` via FNV-1a — that older
/// approach could collide when two threads called within the same nanosecond
/// happened to produce the same FNV mix of `nanos || ThreadId`. The IDs are
/// not security-sensitive (they're used to namespace Docker container names
/// and reference in-process session state), so a non-CSPRNG `random::<u64>()`
/// is sufficient for the uniqueness guarantee.
fn generate_id() -> String {
    format!("{:016x}", rand::random::<u64>())
}

/// Format a Docker-safe container name from agent and session IDs.
/// Strips non-alphanumeric characters to prevent injection via agent_id.
fn container_name(agent_id: &str, session_id: &str) -> String {
    let safe_agent: String = agent_id
        .chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '-' || *c == '_')
        .take(64)
        .collect();
    let safe_session: String = session_id
        .chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '-' || *c == '_')
        .take(64)
        .collect();
    format!("{}-{}-{}", CONTAINER_PREFIX, safe_agent, safe_session)
}

/// `SandboxProvider` backed by hardened Docker containers.
///
/// Uses the Docker CLI via `std::process::Command` — no external Docker
/// crate dependency required.
pub struct DockerSandboxProvider {
    image: String,
    available: bool,
    /// Maps `(agent_id, session_id)` → Docker container name.
    containers: HashMap<(String, String), String>,
}

impl DockerSandboxProvider {
    /// Create a new provider, checking Docker CLI availability via `docker info`.
    pub fn new(image: &str) -> Self {
        let available = Command::new("docker")
            .args(["info"])
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .status()
            .map(|s| s.success())
            .unwrap_or(false);

        Self {
            image: image.to_string(),
            available,
            containers: HashMap::new(),
        }
    }

    /// Return the configured Docker image.
    pub fn image(&self) -> &str {
        &self.image
    }
}

impl SandboxProvider for DockerSandboxProvider {
    fn is_available(&self) -> bool {
        self.available
    }

    fn create_session(
        &mut self,
        agent_id: &str,
        config: Option<&SandboxConfig>,
    ) -> Result<SessionHandle, String> {
        if !self.available {
            return Err("Docker daemon is not available".into());
        }

        let cfg = config.cloned().unwrap_or_default();
        let session_id = generate_id();
        let name = container_name(agent_id, &session_id);

        let mut args = vec![
            "run".to_string(),
            "-d".to_string(),
            "--name".to_string(),
            name.clone(),
            format!("--memory={}m", cfg.memory_mb),
            format!("--cpus={}", cfg.cpu_limit),
            "--cap-drop=ALL".to_string(),
            "--security-opt=no-new-privileges".to_string(),
        ];

        if cfg.read_only_fs {
            args.push("--read-only".to_string());
        }

        if !cfg.network_enabled {
            args.push("--network=none".to_string());
        }

        for (k, v) in &cfg.env_vars {
            validate_env_key(k)?;
            args.push("-e".to_string());
            args.push(format!("{}={}", k, v));
        }

        args.push(self.image.clone());
        args.push("sleep".to_string());
        args.push("infinity".to_string());

        let output = Command::new("docker")
            .args(&args)
            .output()
            .map_err(|e| format!("Failed to run docker: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("docker run failed: {}", stderr.trim()));
        }

        self.containers
            .insert((agent_id.to_string(), session_id.clone()), name);

        Ok(SessionHandle {
            agent_id: agent_id.to_string(),
            session_id,
            status: SessionStatus::Ready,
        })
    }

    fn execute_code(
        &mut self,
        agent_id: &str,
        session_id: &str,
        code: &str,
    ) -> Result<ExecutionHandle, String> {
        let key = (agent_id.to_string(), session_id.to_string());
        let name = self
            .containers
            .get(&key)
            .ok_or_else(|| {
                format!(
                    "No active session for agent '{}' with session_id '{}'. \
                     Call create_session() first.",
                    agent_id, session_id
                )
            })?
            .clone();

        let execution_id = generate_id();
        let start = Instant::now();

        // Avoid shell interpolation: pipe code via stdin instead of sh -c.
        let output = Command::new("docker")
            .args(["exec", "-i", &name, "sh"])
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn()
            .and_then(|mut child| {
                use std::io::Write;
                if let Some(ref mut stdin) = child.stdin {
                    stdin.write_all(code.as_bytes())?;
                }
                drop(child.stdin.take());
                child.wait_with_output()
            })
            .map_err(|e| format!("Failed to run docker exec: {}", e))?;

        let duration = start.elapsed().as_secs_f64();
        let exit_code = output.status.code().unwrap_or(-1);
        let success = output.status.success();

        let result = SandboxResult {
            success,
            exit_code,
            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
            stderr: String::from_utf8_lossy(&output.stderr).to_string(),
            duration_seconds: duration,
            killed: false,
            kill_reason: String::new(),
        };

        let status = if success {
            ExecutionStatus::Completed
        } else {
            ExecutionStatus::Failed
        };

        Ok(ExecutionHandle {
            execution_id,
            agent_id: agent_id.to_string(),
            session_id: session_id.to_string(),
            status,
            result: Some(result),
        })
    }

    fn destroy_session(&mut self, agent_id: &str, session_id: &str) -> Result<(), String> {
        let key = (agent_id.to_string(), session_id.to_string());
        let name = match self.containers.remove(&key) {
            Some(n) => n,
            None => return Err(format!("No active session '{}'", session_id)),
        };

        let output = Command::new("docker")
            .args(["rm", "-f", &name])
            .output()
            .map_err(|e| format!("Failed to run docker rm: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("docker rm failed: {}", stderr.trim()));
        }

        Ok(())
    }

    fn run(
        &mut self,
        agent_id: &str,
        command: &[&str],
        config: Option<&SandboxConfig>,
    ) -> SandboxResult {
        let cfg = config.cloned().unwrap_or_default();

        // Create a one-shot container, run the command, remove on exit.
        let mut args = vec![
            "run".to_string(),
            "--rm".to_string(),
            format!("--memory={}m", cfg.memory_mb),
            format!("--cpus={}", cfg.cpu_limit),
            "--cap-drop=ALL".to_string(),
            "--security-opt=no-new-privileges".to_string(),
        ];

        if cfg.read_only_fs {
            args.push("--read-only".to_string());
        }
        if !cfg.network_enabled {
            args.push("--network=none".to_string());
        }
        for (k, v) in &cfg.env_vars {
            if let Err(e) = validate_env_key(k) {
                return SandboxResult {
                    success: false,
                    exit_code: -1,
                    stderr: e,
                    ..Default::default()
                };
            }
            args.push("-e".to_string());
            args.push(format!("{}={}", k, v));
        }

        args.push(self.image.clone());
        for part in command {
            args.push(part.to_string());
        }

        let _ = agent_id; // reserved for future per-agent auditing
        let start = Instant::now();

        match Command::new("docker").args(&args).output() {
            Ok(output) => {
                let duration = start.elapsed().as_secs_f64();
                let exit_code = output.status.code().unwrap_or(-1);
                SandboxResult {
                    success: output.status.success(),
                    exit_code,
                    stdout: String::from_utf8_lossy(&output.stdout).to_string(),
                    stderr: String::from_utf8_lossy(&output.stderr).to_string(),
                    duration_seconds: duration,
                    killed: false,
                    kill_reason: String::new(),
                }
            }
            Err(e) => SandboxResult {
                success: false,
                exit_code: -1,
                stderr: format!("Failed to execute docker run: {}", e),
                ..Default::default()
            },
        }
    }
}

#[cfg(test)]
#[path = "sandbox_test.rs"]
mod sandbox_test;