From 56c405904de9a47108764a8b7368819997a9584f Mon Sep 17 00:00:00 2001 From: Sam & Claude Date: Sat, 13 Jun 2026 23:20:34 +0200 Subject: [PATCH] fix(spawner): stage jailed env payloads (Sam & Codex) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace inherited env delivery for jailed agent and external MCP spawns with staged launcher/env files under the jail-visible root. Add JailConfig.root_path for named jails that need staged payload delivery. Tests: pass — cargo fmt --all; cargo test -p colibri-daemon jail_tests -- --nocapture; cargo test -p colibri-mcp -- --nocapture --- crates/colibri-daemon/src/spawner.rs | 275 +++++++++++++++++++++++++-- crates/colibri-mcp/src/external.rs | 90 +++++---- 2 files changed, 315 insertions(+), 50 deletions(-) diff --git a/crates/colibri-daemon/src/spawner.rs b/crates/colibri-daemon/src/spawner.rs index 4b59ccd..56ad104 100644 --- a/crates/colibri-daemon/src/spawner.rs +++ b/crates/colibri-daemon/src/spawner.rs @@ -13,6 +13,8 @@ //! Pi-JSONL agents and for future local-only tools. use std::collections::HashMap; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::Arc; use std::time::Duration; @@ -116,6 +118,7 @@ impl PrivMode { /// Default path to the setuid jail-spawn helper used in `PrivMode::Helper`. pub const DEFAULT_JAIL_HELPER: &str = "/usr/local/libexec/colibri-jail-spawn"; +const STAGED_JAIL_RUN_DIR: &str = "/var/run/colibri-stage"; /// Optional FreeBSD jail confinement for a spawned agent. /// @@ -135,6 +138,12 @@ pub struct JailConfig { /// Root path for an ephemeral jail (`jail -c path=…`). Used when `name` is None. #[serde(default)] pub path: Option, + /// Host-visible root path for a named jail. Required when staged payload + /// delivery is needed (env/system prompt/session context launcher files). + /// + /// If unset, `path` is also accepted as the staging root fallback. + #[serde(default)] + pub root_path: Option, /// IPv4 spec: `inherit` (default) → `ip4=inherit`; any other value → /// `ip4.addr=`. #[serde(default)] @@ -168,6 +177,146 @@ fn priv_wrap( } } +fn jail_stage_root(jail: &JailConfig) -> Option<&str> { + jail.root_path.as_deref().or(jail.path.as_deref()) +} + +fn shell_single_quote(value: &str) -> String { + format!("'{}'", value.replace('\'', "'\"'\"'")) +} + +fn valid_env_key(key: &str) -> bool { + let mut chars = key.chars(); + match chars.next() { + Some(c) if c == '_' || c.is_ascii_alphabetic() => {} + _ => return false, + } + chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) +} + +fn build_env_script(env: &HashMap) -> std::io::Result { + let mut entries: Vec<_> = env.iter().collect(); + entries.sort_by(|a, b| a.0.cmp(b.0)); + + let mut script = String::new(); + for (key, value) in entries { + if !valid_env_key(key) { + return Err(std::io::Error::other(format!( + "invalid env key for staged jail payload: {key}" + ))); + } + script.push_str("export "); + script.push_str(key); + script.push('='); + script.push_str(&shell_single_quote(value)); + script.push('\n'); + } + Ok(script) +} + +fn build_jail_launcher() -> &'static str { + r#"#!/bin/sh +set -eu + +ENV_FILE="$1" +shift +WORK_DIR="$1" +shift + +if [ "$ENV_FILE" != "-" ] && [ -f "$ENV_FILE" ]; then + . "$ENV_FILE" +fi + +if [ "$WORK_DIR" != "-" ]; then + cd "$WORK_DIR" +fi + +exec "$@" +"# +} + +#[derive(Debug)] +pub struct PreparedSpawnCommand { + pub program: String, + pub argv: Vec, + pub env: HashMap, + pub cleanup_dir: Option, +} + +async fn remove_staged_dir(path: &Path) { + let _ = tokio::fs::remove_dir_all(path).await; +} + +pub async fn prepare_spawn_command( + binary: &str, + args: &[String], + env: HashMap, + working_dir: Option<&str>, + jail: Option<&JailConfig>, + priv_mode: &PrivMode, + helper: &str, + stage_id: &str, +) -> std::io::Result { + let Some(jail) = jail else { + return Ok(PreparedSpawnCommand { + program: binary.to_string(), + argv: args.to_vec(), + env, + cleanup_dir: None, + }); + }; + + let needs_staging = !env.is_empty() || working_dir.is_some(); + if !needs_staging { + let (program, argv) = jail_wrap(binary, args, Some(jail), priv_mode, helper); + return Ok(PreparedSpawnCommand { + program, + argv, + env, + cleanup_dir: None, + }); + } + + let stage_root = jail_stage_root(jail).ok_or_else(|| { + std::io::Error::other( + "jailed spawn with staged payload requires JailConfig.root_path (or path for ephemeral jails)", + ) + })?; + + let jail_stage_dir = format!("{STAGED_JAIL_RUN_DIR}/{stage_id}"); + let host_stage_dir = Path::new(stage_root).join(jail_stage_dir.trim_start_matches('/')); + tokio::fs::create_dir_all(&host_stage_dir).await?; + + let launcher_host = host_stage_dir.join("launch.sh"); + tokio::fs::write(&launcher_host, build_jail_launcher()).await?; + std::fs::set_permissions(&launcher_host, std::fs::Permissions::from_mode(0o700))?; + + let env_rel = if env.is_empty() { + "-".to_string() + } else { + let env_host = host_stage_dir.join("env.sh"); + tokio::fs::write(&env_host, build_env_script(&env)?).await?; + std::fs::set_permissions(&env_host, std::fs::Permissions::from_mode(0o600))?; + format!("{jail_stage_dir}/env.sh") + }; + + let mut staged_args = vec![ + format!("{jail_stage_dir}/launch.sh"), + env_rel, + working_dir.unwrap_or("-").to_string(), + binary.to_string(), + ]; + staged_args.extend(args.iter().cloned()); + + let (program, argv) = jail_wrap("/bin/sh", &staged_args, Some(jail), priv_mode, helper); + Ok(PreparedSpawnCommand { + program, + argv, + env: HashMap::new(), + cleanup_dir: Some(host_stage_dir), + }) +} + /// Wrap an agent command for optional jail confinement plus privilege /// escalation, returning the `(program, argv)` to hand to `Command::new`. /// @@ -325,14 +474,16 @@ pub struct AgentHandle { stdout: Mutex>, /// Creation timestamp. pub created_at: String, + /// Optional staged jail payload directory to remove when the child exits. + staging_dir: Option, } impl AgentHandle { pub fn new(config: AgentSpawnConfig) -> Self { - Self::new_with_id(Uuid::new_v4().to_string(), config) + Self::new_with_id(Uuid::new_v4().to_string(), config, None) } - pub fn new_with_id(id: String, config: AgentSpawnConfig) -> Self { + pub fn new_with_id(id: String, config: AgentSpawnConfig, staging_dir: Option) -> Self { Self { id, config, @@ -340,6 +491,7 @@ impl AgentHandle { child: Mutex::new(None), stdout: Mutex::new(None), created_at: chrono::Utc::now().to_rfc3339(), + staging_dir, } } @@ -363,6 +515,9 @@ impl AgentHandle { } *child = None; *self.status.write().await = AgentStatus::Stopped; + if let Some(path) = self.staging_dir.as_deref() { + remove_staged_dir(path).await; + } Ok(()) } @@ -374,12 +529,18 @@ impl AgentHandle { Ok(Some(status)) => { *child_guard = None; *self.status.write().await = AgentStatus::Stopped; + if let Some(path) = self.staging_dir.as_deref() { + remove_staged_dir(path).await; + } Some(status) } Ok(None) => None, Err(_) => { *child_guard = None; *self.status.write().await = AgentStatus::Failed("poll error".to_string()); + if let Some(path) = self.staging_dir.as_deref() { + remove_staged_dir(path).await; + } None } } @@ -466,8 +627,6 @@ impl Spawner { continue; } - let handle = AgentHandle::new_with_id(agent_id.clone(), agent_config.clone()); - // Build environment for this provider let mut env_map = agent_config.env.clone(); if let Some(api_key) = provider.api_key(&self.config) { @@ -497,29 +656,40 @@ impl Spawner { // Spawn with retry/backoff let spawn_result = { - // Apply optional jail confinement + privilege escalation. With - // no jail this is the bare (binary, args); stdio is unaffected. - let (exe, argv) = jail_wrap( + // Prepare the command once per spawn attempt. For jailed runs, + // large/custom payload is staged into a launcher+env file + // instead of relying on inherited env across jexec/mdo. + let prepared = prepare_spawn_command( &agent_config.binary, &agent_config.args, + env_map.clone(), + agent_config.working_dir.as_deref(), agent_config.jail.as_ref(), &priv_mode, &jail_helper, + &agent_id, + ) + .await + .map_err(SpawnerError::Io)?; + let handle = AgentHandle::new_with_id( + agent_id.clone(), + agent_config.clone(), + prepared.cleanup_dir.clone(), ); - let env = env_map.clone(); - let working_dir = agent_config.working_dir.clone(); let _startup_timeout = agent_config.startup_timeout_secs; let op = || async { - let mut cmd = Command::new(&exe); - cmd.args(&argv) - .envs(&env) + let mut cmd = Command::new(&prepared.program); + cmd.args(&prepared.argv) + .envs(&prepared.env) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); - if let Some(ref dir) = working_dir { - cmd.current_dir(dir); + if agent_config.jail.is_none() { + if let Some(ref dir) = agent_config.working_dir { + cmd.current_dir(dir); + } } let mut child = cmd.spawn().map_err(SpawnerError::Io)?; @@ -538,6 +708,9 @@ impl Spawner { } else { String::new() }; + if let Some(path) = prepared.cleanup_dir.as_deref() { + remove_staged_dir(path).await; + } Err(SpawnerError::Io(std::io::Error::other(format!( "process exited immediately with status {status}: {stderr}" )))) @@ -555,11 +728,13 @@ impl Spawner { .with_max_delay(Duration::from_secs(10)) .with_max_times(agent_config.max_retries as usize); - op.retry(backoff).await + op.retry(backoff) + .await + .map(|(child, stdout)| (handle, child, stdout)) }; match spawn_result { - Ok((child, stdout)) => { + Ok((handle, child, stdout)) => { *handle.child.lock().await = Some(child); *handle.stdout.lock().await = stdout; *handle.status.write().await = AgentStatus::Running; @@ -630,6 +805,7 @@ mod jail_tests { fn named_jail_via_mdo() { let j = JailConfig { name: Some("pi0".into()), + root_path: Some("/usr/local/bastille/jails/pi0/root".into()), ..Default::default() }; let (exe, a) = jail_wrap( @@ -647,6 +823,7 @@ mod jail_tests { fn named_jail_with_user_no_priv() { let j = JailConfig { name: Some("pi0".into()), + root_path: Some("/usr/local/bastille/jails/pi0/root".into()), user: Some("clawdie".into()), ..Default::default() }; @@ -666,6 +843,7 @@ mod jail_tests { let j = JailConfig { name: Some("pi0".into()), path: Some("/var/jails/pi".into()), + root_path: Some("/usr/local/bastille/jails/pi0/root".into()), ..Default::default() }; let (exe, a) = jail_wrap("pi", &[], Some(&j), &PrivMode::None, DEFAULT_JAIL_HELPER); @@ -721,4 +899,69 @@ mod jail_tests { ]) ); } + + #[tokio::test] + async fn staged_named_jail_writes_launcher_and_env() { + let root = std::env::temp_dir().join(format!("colibri-jail-stage-{}", Uuid::new_v4())); + let j = JailConfig { + name: Some("pi0".into()), + root_path: Some(root.to_string_lossy().to_string()), + ..Default::default() + }; + let mut env = HashMap::new(); + env.insert("COLIBRI_SYSTEM_PROMPT".into(), "Sam's prompt".into()); + let prepared = prepare_spawn_command( + "pi", + &argv(&["--mode", "json"]), + env, + Some("/work"), + Some(&j), + &PrivMode::None, + DEFAULT_JAIL_HELPER, + "agent-test", + ) + .await + .unwrap(); + + assert_eq!(prepared.program, "jexec"); + assert_eq!(prepared.argv[0], "pi0"); + assert_eq!(prepared.argv[1], "/bin/sh"); + assert_eq!( + prepared.argv[2], + "/var/run/colibri-stage/agent-test/launch.sh" + ); + assert!(prepared.env.is_empty()); + + let launcher = root.join("var/run/colibri-stage/agent-test/launch.sh"); + let env_file = root.join("var/run/colibri-stage/agent-test/env.sh"); + assert!(launcher.exists()); + assert!(env_file.exists()); + let env_raw = std::fs::read_to_string(env_file).unwrap(); + assert!(env_raw.contains("COLIBRI_SYSTEM_PROMPT")); + assert!(env_raw.contains("Sam'\"'\"'s prompt")); + + remove_staged_dir(prepared.cleanup_dir.as_deref().unwrap()).await; + } + + #[tokio::test] + async fn staged_jail_requires_root_path() { + let mut env = HashMap::new(); + env.insert("COLIBRI_AGENT_ID".into(), "a1".into()); + let err = prepare_spawn_command( + "pi", + &[], + env, + None, + Some(&JailConfig { + name: Some("pi0".into()), + ..Default::default() + }), + &PrivMode::None, + DEFAULT_JAIL_HELPER, + "missing-root", + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("requires JailConfig.root_path")); + } } diff --git a/crates/colibri-mcp/src/external.rs b/crates/colibri-mcp/src/external.rs index 03b2dfc..4779449 100644 --- a/crates/colibri-mcp/src/external.rs +++ b/crates/colibri-mcp/src/external.rs @@ -3,7 +3,12 @@ //! This is intentionally thin: stdio transport only, one process per request, //! no long-lived registry, and no production-grade policy engine yet. -use std::{collections::BTreeMap, path::Path, process::Stdio, time::Duration}; +use std::{ + collections::{BTreeMap, HashMap}, + path::{Path, PathBuf}, + process::Stdio, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -13,7 +18,9 @@ use tokio::{ time::timeout, }; -use colibri_daemon::spawner::{jail_wrap, JailConfig, PrivMode, DEFAULT_JAIL_HELPER}; +#[cfg(test)] +use colibri_daemon::spawner::jail_wrap; +use colibri_daemon::spawner::{prepare_spawn_command, JailConfig, PrivMode, DEFAULT_JAIL_HELPER}; use crate::protocol::McpError; @@ -44,24 +51,6 @@ pub struct ExternalMcpServer { pub jail: Option, } -impl ExternalMcpServer { - /// Resolve the `(program, argv)` to exec, applying optional jail - /// confinement + privilege escalation via the shared spawner policy. - /// With `jail == None` this is just `(command, args)`. - fn resolved_command(&self) -> (String, Vec) { - let priv_mode = PrivMode::from_env(); - let helper = std::env::var("COLIBRI_JAIL_HELPER") - .unwrap_or_else(|_| DEFAULT_JAIL_HELPER.to_string()); - jail_wrap( - &self.command, - &self.args, - self.jail.as_ref(), - &priv_mode, - &helper, - ) - } -} - impl ExternalMcpRegistry { pub async fn load(path: &Path) -> Result { let raw = tokio::fs::read_to_string(path).await.map_err(|e| { @@ -143,22 +132,51 @@ struct ExternalMcpSession { stdin: tokio::process::ChildStdin, stdout: BufReader, next_id: i64, + cleanup_dir: Option, } impl ExternalMcpSession { async fn start(server: &ExternalMcpServer) -> Result { - // Apply optional jail confinement. With no jail this is (command, args); - // jexec/jail/mdo all inherit stdio, so the piped stdin/stdout below still - // reach the actual MCP server inside the jail. - let (program, argv) = server.resolved_command(); - let mut cmd = Command::new(&program); - cmd.args(&argv) - .envs(&server.env) + let env: HashMap = server + .env + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let stage_id = format!( + "mcp-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() + ); + let prepared = prepare_spawn_command( + &server.command, + &server.args, + env, + None, + server.jail.as_ref(), + &PrivMode::from_env(), + &std::env::var("COLIBRI_JAIL_HELPER") + .unwrap_or_else(|_| DEFAULT_JAIL_HELPER.to_string()), + &stage_id, + ) + .await + .map_err(|e| McpError::internal(format!("failed to prepare jailed MCP spawn: {e}")))?; + + let mut cmd = Command::new(&prepared.program); + cmd.args(&prepared.argv) + .envs(&prepared.env) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::null()); let mut child = cmd.spawn().map_err(|e| { + if let Some(path) = prepared.cleanup_dir.as_deref() { + let path = path.to_path_buf(); + tokio::spawn(async move { + let _ = tokio::fs::remove_dir_all(path).await; + }); + } McpError::internal(format!( "failed to spawn external MCP server {}: {e}", server.command @@ -176,6 +194,7 @@ impl ExternalMcpSession { stdin, stdout: BufReader::new(stdout), next_id: 1, + cleanup_dir: prepared.cleanup_dir, }) } @@ -264,6 +283,9 @@ impl ExternalMcpSession { async fn shutdown(&mut self) { let _ = self.child.kill().await; let _ = self.child.wait().await; + if let Some(path) = self.cleanup_dir.as_deref() { + let _ = tokio::fs::remove_dir_all(path).await; + } } } @@ -294,13 +316,13 @@ mod tests { #[test] fn no_jail_resolves_to_bare_command() { // jail == None ⇒ unchanged regardless of COLIBRI_JAIL_PRIV_MODE. - let s = ExternalMcpServer { - command: "mcp-fs".into(), - args: vec!["--root".into(), "/srv".into()], - env: BTreeMap::new(), - jail: None, - }; - let (program, argv) = s.resolved_command(); + let (program, argv) = jail_wrap( + "mcp-fs", + &["--root".to_string(), "/srv".to_string()], + None, + &PrivMode::Mdo, + DEFAULT_JAIL_HELPER, + ); assert_eq!(program, "mcp-fs"); assert_eq!(argv, vec!["--root".to_string(), "/srv".to_string()]); }