diff --git a/shell/Cargo.toml b/shell/Cargo.toml new file mode 100644 index 000000000..772dd1eb3 --- /dev/null +++ b/shell/Cargo.toml @@ -0,0 +1,27 @@ +[workspace] + +[package] +name = "iii-shell" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "iii-shell" +path = "src/main.rs" + +[dependencies] +iii-sdk = "=0.11.3" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "signal", "process", "time", "io-util"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +serde_yaml = "0.9" +anyhow = "1" +thiserror = "2" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } +clap = { version = "4", features = ["derive"] } +uuid = { version = "1", features = ["v4"] } +once_cell = "1" +regex = "1" +shell-words = "1" diff --git a/shell/README.md b/shell/README.md new file mode 100644 index 000000000..5ed18e601 --- /dev/null +++ b/shell/README.md @@ -0,0 +1,69 @@ +# iii-shell + +Unix shell execution worker for iii agents. Mike's priority-1 fundamental — every agent worker that needs to touch the OS (run a build, read a file via `cat`, list a directory, call a CLI) goes through this worker so there is a single place to enforce allowlists, timeouts, and output caps. + +## Functions + +| id | shape | +|----|-------| +| `shell::exec` | run to completion, return `{exit_code, stdout, stderr, duration_ms, timed_out, stdout_truncated, stderr_truncated}` | +| `shell::exec_bg` | spawn in background, return `{job_id, argv}` | +| `shell::kill` | kill a running `job_id` | +| `shell::status` | return `{job: JobRecord}` for a `job_id` | +| `shell::list` | return all jobs + counts | + +## HTTP triggers + +``` +POST /api/shell/exec → shell::exec +POST /api/shell/exec_bg → shell::exec_bg +POST /api/shell/kill → shell::kill +POST /api/shell/status → shell::status +GET /api/shell/list → shell::list +``` + +## Safety + +- `allowlist` — if non-empty, command (basename) must be present. Empty list = open. +- `denylist_patterns` — regex patterns tested against the full joined argv. Example: `rm\s+-rf\s+/`, `:()\s*\{\s*:\|` (fork bomb), `mkfs`, `shutdown`. +- `max_timeout_ms` — hard cap; per-call `timeout_ms` is clamped. +- `max_output_bytes` — stdout/stderr are truncated at this size, flagged via `*_truncated`. +- `inherit_env: false` by default. Only variables in `allowed_env` are forwarded. +- `working_dir` — pins cwd. +- `max_concurrent_jobs` — rejects new `exec_bg` requests past the cap. +- `job_retention_secs` — old finished jobs are pruned on every `shell::list` call. + +## Example + +```bash +curl -X POST localhost:3111/api/shell/exec -d '{ + "command": "ls", + "args": ["-la", "/tmp"], + "timeout_ms": 5000 +}' +# → {"exit_code": 0, "stdout": "total …", "stderr": "", "duration_ms": 12, ...} + +curl -X POST localhost:3111/api/shell/exec_bg -d '{ + "command": "cargo", + "args": ["build", "--release"] +}' +# → {"job_id": "job-abc…", "argv": ["cargo", "build", "--release"]} + +curl -X POST localhost:3111/api/shell/status -d '{"job_id": "job-abc…"}' +``` + +## Run locally + +```bash +cargo run --release -- --config ./config.yaml --url ws://127.0.0.1:49134 +``` + +## What this is NOT + +- Not a PTY. Interactive shells, TUIs, password prompts all break. +- Not a remote executor. Runs on the worker's host only. +- Not a sandbox. For isolation use `sandbox-docker`/`sandbox-firecracker` and call through `shell` only for trusted commands. + +## Deferred + +- `shell::exec_stream` — live stdout/stderr via iii Streams (for long-running commands). Next iteration. diff --git a/shell/build.rs b/shell/build.rs new file mode 100644 index 000000000..33143a5e0 --- /dev/null +++ b/shell/build.rs @@ -0,0 +1,6 @@ +fn main() { + println!( + "cargo:rustc-env=TARGET={}", + std::env::var("TARGET").unwrap_or_default() + ); +} diff --git a/shell/config.yaml b/shell/config.yaml new file mode 100644 index 000000000..6fb154b67 --- /dev/null +++ b/shell/config.yaml @@ -0,0 +1,60 @@ +max_timeout_ms: 30000 +default_timeout_ms: 10000 +max_output_bytes: 1048576 +working_dir: null +inherit_env: false +allowed_env: + - PATH + - HOME + - LANG + - LC_ALL + - TERM +# Default allowlist is intentionally read-only. Tools that can shell out +# (git hooks, curl -o/file://, find -exec, awk system(), sed e/-i, cargo +# build.rs, node -e, python3 -c, npm run) are left out on purpose — add them +# per deployment after you've decided on the threat model. This worker is +# NOT a sandbox. +allowlist: + - ls + - cat + - pwd + - echo + - grep + - wc + - head + - tail + - sort + - uniq + - cut + - date + - whoami + - hostname + - which + - jq + - uname + - df + - du + - ps + - env + - printenv + - basename + - dirname +denylist_patterns: + - "rm\\s+-rf\\s+/" + - ":\\(\\)\\s*\\{\\s*:\\|" + - "mkfs" + - "dd\\s+if=" + - "shutdown" + - "reboot" + - "/etc/passwd" + - "/etc/shadow" + # Sub-execution / write escapes for commonly-added tools + - "\\bfind\\b[^|;&]*-exec(dir)?\\b" + - "\\bawk\\b[^|;&]*system\\s*\\(" + - "\\bsed\\b[^|;&]*(-i\\b|\\be\\b)" + - "\\bcurl\\b[^|;&]*(file://|-o\\s|--output-dir\\b|-F\\s+@)" + - "\\bgit\\b[^|;&]*(--upload-pack|--receive-pack|core\\.pager|core\\.hooksPath|GIT_SSH_COMMAND)" + - "\\b(node|python3?)\\b[^|;&]*\\s-(e|c)\\b" + - "\\bnpm\\b[^|;&]*\\brun\\b" +max_concurrent_jobs: 16 +job_retention_secs: 3600 diff --git a/shell/src/config.rs b/shell/src/config.rs new file mode 100644 index 000000000..024244b05 --- /dev/null +++ b/shell/src/config.rs @@ -0,0 +1,206 @@ +use anyhow::{Context, Result}; +use regex::Regex; +use serde::{Deserialize, Serialize}; +use std::fs; +use std::path::PathBuf; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ShellConfig { + #[serde(default = "default_max_timeout_ms")] + pub max_timeout_ms: u64, + + #[serde(default = "default_default_timeout_ms")] + pub default_timeout_ms: u64, + + #[serde(default = "default_max_output_bytes")] + pub max_output_bytes: usize, + + #[serde(default)] + pub working_dir: Option, + + #[serde(default)] + pub inherit_env: bool, + + #[serde(default = "default_allowed_env")] + pub allowed_env: Vec, + + #[serde(default)] + pub allowlist: Vec, + + #[serde(default)] + pub denylist_patterns: Vec, + + #[serde(default = "default_max_concurrent_jobs")] + pub max_concurrent_jobs: usize, + + #[serde(default = "default_job_retention_secs")] + pub job_retention_secs: u64, + + #[serde(default, skip)] + pub compiled_denylist: Vec, +} + +fn default_max_timeout_ms() -> u64 { + 30_000 +} +fn default_default_timeout_ms() -> u64 { + 10_000 +} +fn default_max_output_bytes() -> usize { + 1_048_576 +} +fn default_allowed_env() -> Vec { + vec!["PATH", "HOME", "LANG", "LC_ALL", "TERM"] + .into_iter() + .map(String::from) + .collect() +} +fn default_max_concurrent_jobs() -> usize { + 16 +} +fn default_job_retention_secs() -> u64 { + 3600 +} + +impl Default for ShellConfig { + fn default() -> Self { + Self { + max_timeout_ms: default_max_timeout_ms(), + default_timeout_ms: default_default_timeout_ms(), + max_output_bytes: default_max_output_bytes(), + working_dir: None, + inherit_env: false, + allowed_env: default_allowed_env(), + allowlist: Vec::new(), + denylist_patterns: Vec::new(), + max_concurrent_jobs: default_max_concurrent_jobs(), + job_retention_secs: default_job_retention_secs(), + compiled_denylist: Vec::new(), + } + } +} + +pub fn load_config(path: &str) -> Result { + let content = fs::read_to_string(path).with_context(|| format!("read {}", path))?; + let mut cfg: ShellConfig = + serde_yaml::from_str(&content).with_context(|| format!("parse {}", path))?; + cfg.compile_denylist()?; + Ok(cfg) +} + +impl ShellConfig { + pub fn compile_denylist(&mut self) -> Result<()> { + self.compiled_denylist = self + .denylist_patterns + .iter() + .map(|p| Regex::new(p).with_context(|| format!("bad denylist pattern: {}", p))) + .collect::>>()?; + Ok(()) + } + + pub fn is_command_allowed(&self, argv: &[String]) -> Result<(), String> { + let cmd = argv + .first() + .ok_or_else(|| "empty command".to_string())? + .clone(); + + if !self.allowlist.is_empty() { + let base = std::path::Path::new(&cmd) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or(&cmd); + if !self.allowlist.iter().any(|a| a == base || a == &cmd) { + return Err(format!("command '{}' not in allowlist", base)); + } + } + + let joined = argv.join(" "); + for re in &self.compiled_denylist { + if re.is_match(&joined) { + return Err(format!("command matches denylist: {}", re.as_str())); + } + } + Ok(()) + } + + pub fn resolve_timeout(&self, requested: Option) -> u64 { + let t = requested.unwrap_or(self.default_timeout_ms); + t.min(self.max_timeout_ms) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cfg_with(allow: Vec<&str>, deny: Vec<&str>) -> ShellConfig { + let mut c = ShellConfig { + allowlist: allow.into_iter().map(String::from).collect(), + denylist_patterns: deny.into_iter().map(String::from).collect(), + ..Default::default() + }; + c.compile_denylist().unwrap(); + c + } + + #[test] + fn test_defaults() { + let c = ShellConfig::default(); + assert_eq!(c.max_timeout_ms, 30_000); + assert_eq!(c.default_timeout_ms, 10_000); + assert!(!c.inherit_env); + assert_eq!(c.max_concurrent_jobs, 16); + } + + #[test] + fn test_allowlist_permits() { + let c = cfg_with(vec!["ls", "cat"], vec![]); + assert!(c.is_command_allowed(&["ls".into(), "-la".into()]).is_ok()); + } + + #[test] + fn test_allowlist_rejects() { + let c = cfg_with(vec!["ls"], vec![]); + let err = c + .is_command_allowed(&["nmap".into()]) + .expect_err("must reject"); + assert!(err.contains("not in allowlist")); + } + + #[test] + fn test_allowlist_empty_means_open() { + let c = cfg_with(vec![], vec![]); + assert!(c.is_command_allowed(&["anything".into()]).is_ok()); + } + + #[test] + fn test_allowlist_basename_match() { + let c = cfg_with(vec!["ls"], vec![]); + assert!(c + .is_command_allowed(&["/usr/bin/ls".into(), "-la".into()]) + .is_ok()); + } + + #[test] + fn test_denylist_blocks() { + let c = cfg_with(vec![], vec![r"rm\s+-rf\s+/"]); + let err = c + .is_command_allowed(&["rm".into(), "-rf".into(), "/".into()]) + .expect_err("must reject"); + assert!(err.contains("denylist")); + } + + #[test] + fn test_empty_argv_rejected() { + let c = ShellConfig::default(); + assert!(c.is_command_allowed(&[]).is_err()); + } + + #[test] + fn test_resolve_timeout_caps_at_max() { + let c = ShellConfig::default(); + assert_eq!(c.resolve_timeout(Some(60_000)), 30_000); + assert_eq!(c.resolve_timeout(Some(5_000)), 5_000); + assert_eq!(c.resolve_timeout(None), 10_000); + } +} diff --git a/shell/src/exec.rs b/shell/src/exec.rs new file mode 100644 index 000000000..52fbb8232 --- /dev/null +++ b/shell/src/exec.rs @@ -0,0 +1,190 @@ +use crate::config::ShellConfig; +use anyhow::Result; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::process::Command; + +pub struct ExecOutcome { + pub stdout: String, + pub stderr: String, + pub exit_code: Option, + pub duration_ms: u64, + pub timed_out: bool, + pub stdout_truncated: bool, + pub stderr_truncated: bool, +} + +pub fn parse_argv(command: &str, args: Option<&Vec>) -> Result, String> { + if let Some(args) = args { + let mut v = vec![command.to_string()]; + v.extend(args.iter().cloned()); + Ok(v) + } else { + shell_words::split(command).map_err(|e| format!("parse command: {}", e)) + } +} + +pub fn build_command(argv: &[String], cfg: &ShellConfig) -> Result { + let program = argv.first().ok_or_else(|| "empty command".to_string())?; + let mut cmd = Command::new(program); + if argv.len() > 1 { + cmd.args(&argv[1..]); + } + if !cfg.inherit_env { + cmd.env_clear(); + for k in &cfg.allowed_env { + if let Ok(v) = std::env::var(k) { + cmd.env(k, v); + } + } + } + if let Some(dir) = &cfg.working_dir { + cmd.current_dir(dir); + } + cmd.stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + Ok(cmd) +} + +pub async fn run_to_completion( + argv: &[String], + cfg: &ShellConfig, + timeout_ms: u64, +) -> Result { + let started = std::time::Instant::now(); + let mut cmd = build_command(argv, cfg)?; + let mut child = cmd.spawn().map_err(|e| format!("spawn: {}", e))?; + + let mut stdout_reader = child.stdout.take().ok_or("no stdout pipe")?; + let mut stderr_reader = child.stderr.take().ok_or("no stderr pipe")?; + + let limit = cfg.max_output_bytes; + let timeout = Duration::from_millis(timeout_ms); + + let stdout_task = tokio::spawn(async move { read_bounded(&mut stdout_reader, limit).await }); + let stderr_task = tokio::spawn(async move { read_bounded(&mut stderr_reader, limit).await }); + + let wait_res = tokio::time::timeout(timeout, child.wait()).await; + + let (exit_code, timed_out) = match wait_res { + Ok(Ok(status)) => (status.code(), false), + Ok(Err(e)) => return Err(format!("wait: {}", e)), + Err(_) => { + let _ = child.start_kill(); + let _ = child.wait().await; + (None, true) + } + }; + + let (stdout_bytes, stdout_truncated) = stdout_task.await.map_err(|e| format!("stdout: {}", e))?; + let (stderr_bytes, stderr_truncated) = stderr_task.await.map_err(|e| format!("stderr: {}", e))?; + + Ok(ExecOutcome { + stdout: String::from_utf8_lossy(&stdout_bytes).into_owned(), + stderr: String::from_utf8_lossy(&stderr_bytes).into_owned(), + exit_code, + duration_ms: started.elapsed().as_millis() as u64, + timed_out, + stdout_truncated, + stderr_truncated, + }) +} + +async fn read_bounded(reader: &mut R, limit: usize) -> (Vec, bool) { + let mut buf = Vec::with_capacity(limit.min(8192)); + let mut chunk = [0u8; 8192]; + let mut truncated = false; + loop { + match reader.read(&mut chunk).await { + Ok(0) => break, + Ok(n) => { + if buf.len() + n > limit { + let take = limit.saturating_sub(buf.len()); + buf.extend_from_slice(&chunk[..take]); + truncated = true; + break; + } + buf.extend_from_slice(&chunk[..n]); + } + Err(_) => break, + } + } + (buf, truncated) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_cfg() -> ShellConfig { + let mut c = ShellConfig { + inherit_env: true, + max_output_bytes: 4096, + ..Default::default() + }; + c.compile_denylist().unwrap(); + c + } + + #[test] + fn test_parse_argv_with_args_field() { + let got = parse_argv("echo", Some(&vec!["hello".into(), "world".into()])).unwrap(); + assert_eq!(got, vec!["echo", "hello", "world"]); + } + + #[test] + fn test_parse_argv_from_shell_words() { + let got = parse_argv(r#"echo "hello world""#, None).unwrap(); + assert_eq!(got, vec!["echo", "hello world"]); + } + + #[test] + fn test_parse_argv_bad_quoting() { + assert!(parse_argv(r#"echo "unterminated"#, None).is_err()); + } + + #[tokio::test] + async fn test_run_echo() { + let cfg = test_cfg(); + let out = run_to_completion(&["echo".into(), "hi".into()], &cfg, 5000) + .await + .unwrap(); + assert_eq!(out.exit_code, Some(0)); + assert_eq!(out.stdout.trim(), "hi"); + assert!(!out.timed_out); + } + + #[tokio::test] + async fn test_run_nonexistent_command() { + let cfg = test_cfg(); + let err = run_to_completion(&["_nope_no_exist_".into()], &cfg, 1000).await; + assert!(err.is_err()); + } + + #[tokio::test] + async fn test_timeout_kills() { + let cfg = test_cfg(); + let out = run_to_completion(&["sleep".into(), "5".into()], &cfg, 200) + .await + .unwrap(); + assert!(out.timed_out); + assert_eq!(out.exit_code, None); + } + + #[tokio::test] + async fn test_output_truncation() { + let mut cfg = test_cfg(); + cfg.max_output_bytes = 16; + let out = run_to_completion( + &["sh".into(), "-c".into(), "printf 'x%.0s' $(seq 1 100)".into()], + &cfg, + 3000, + ) + .await + .unwrap(); + assert!(out.stdout_truncated); + assert_eq!(out.stdout.len(), 16); + } +} diff --git a/shell/src/functions/exec.rs b/shell/src/functions/exec.rs new file mode 100644 index 000000000..8662b164e --- /dev/null +++ b/shell/src/functions/exec.rs @@ -0,0 +1,57 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::IIIError; +use serde_json::{json, Value}; + +use crate::config::ShellConfig; +use crate::exec::{parse_argv, run_to_completion}; + +pub fn build_handler( + config: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let cfg = config.clone(); + Box::pin(async move { handle(cfg, payload).await }) + } +} + +async fn handle(cfg: Arc, payload: Value) -> Result { + let command = payload + .get("command") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'command'".to_string()))?; + let args: Option> = payload + .get("args") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_str().map(String::from)).collect()); + let timeout_ms = payload + .get("timeout_ms") + .and_then(|v| v.as_u64()); + + let argv = parse_argv(command, args.as_ref()) + .map_err(|e| IIIError::Handler(format!("argv: {}", e)))?; + + cfg.is_command_allowed(&argv) + .map_err(|e| IIIError::Handler(e))?; + + let timeout = cfg.resolve_timeout(timeout_ms); + + let out = run_to_completion(&argv, &cfg, timeout) + .await + .map_err(|e| IIIError::Handler(format!("exec: {}", e)))?; + + Ok(json!({ + "exit_code": out.exit_code, + "stdout": out.stdout, + "stderr": out.stderr, + "duration_ms": out.duration_ms, + "timed_out": out.timed_out, + "stdout_truncated": out.stdout_truncated, + "stderr_truncated": out.stderr_truncated, + })) +} diff --git a/shell/src/functions/exec_bg.rs b/shell/src/functions/exec_bg.rs new file mode 100644 index 000000000..5ea4c616b --- /dev/null +++ b/shell/src/functions/exec_bg.rs @@ -0,0 +1,174 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::IIIError; +use serde_json::{json, Value}; +use uuid::Uuid; + +use crate::config::ShellConfig; +use crate::exec::{build_command, parse_argv}; +use crate::jobs::{self, JobHandle, JobRecord, JobStatus}; +use tokio::io::AsyncReadExt; + +pub fn build_handler( + config: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let cfg = config.clone(); + Box::pin(async move { handle(cfg, payload).await }) + } +} + +async fn handle(cfg: Arc, payload: Value) -> Result { + let command = payload + .get("command") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'command'".to_string()))?; + let args: Option> = payload + .get("args") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_str().map(String::from)).collect()); + + let argv = parse_argv(command, args.as_ref()) + .map_err(|e| IIIError::Handler(format!("argv: {}", e)))?; + + cfg.is_command_allowed(&argv) + .map_err(IIIError::Handler)?; + + let running = jobs::running_count().await; + if running >= cfg.max_concurrent_jobs { + return Err(IIIError::Handler(format!( + "max concurrent jobs ({}) reached", + cfg.max_concurrent_jobs + ))); + } + + let mut cmd = build_command(&argv, &cfg).map_err(IIIError::Handler)?; + let mut child = cmd + .spawn() + .map_err(|e| IIIError::Handler(format!("spawn: {}", e)))?; + + let stdout_pipe = child.stdout.take(); + let stderr_pipe = child.stderr.take(); + + let id = format!("job-{}", Uuid::new_v4()); + let record = JobRecord { + id: id.clone(), + argv: argv.clone(), + started_at_ms: jobs::now_ms(), + finished_at_ms: None, + status: JobStatus::Running, + exit_code: None, + stdout: String::new(), + stderr: String::new(), + stdout_truncated: false, + stderr_truncated: false, + }; + jobs::insert(JobHandle { + record, + child: Some(child), + }) + .await; + + let id_clone = id.clone(); + let limit = cfg.max_output_bytes; + tokio::spawn(async move { + let handle = match jobs::get(&id_clone).await { + Some(h) => h, + None => return, + }; + + // Drain stdout and stderr concurrently. Sequential reads deadlock + // when the child fills one pipe's buffer (~64 KiB on Linux) before + // closing the other — matches the pattern used by run_to_completion. + let stdout_task = stdout_pipe.map(|mut out| { + tokio::spawn(async move { + let mut buf = Vec::new(); + let mut trunc = false; + read_bounded(&mut out, limit, &mut buf, &mut trunc).await; + (buf, trunc) + }) + }); + let stderr_task = stderr_pipe.map(|mut err| { + tokio::spawn(async move { + let mut buf = Vec::new(); + let mut trunc = false; + read_bounded(&mut err, limit, &mut buf, &mut trunc).await; + (buf, trunc) + }) + }); + + let (stdout_buf, stdout_trunc) = match stdout_task { + Some(t) => t.await.unwrap_or_else(|_| (Vec::new(), false)), + None => (Vec::new(), false), + }; + let (stderr_buf, stderr_trunc) = match stderr_task { + Some(t) => t.await.unwrap_or_else(|_| (Vec::new(), false)), + None => (Vec::new(), false), + }; + + { + let mut h = handle.lock().await; + if let Some(mut ch) = h.child.take() { + drop(h); + let wait_res = ch.wait().await; + let mut h2 = handle.lock().await; + match wait_res { + Ok(s) => { + h2.record.exit_code = s.code(); + if h2.record.status == JobStatus::Running { + h2.record.status = if s.success() { + JobStatus::Finished + } else { + JobStatus::Failed + }; + } + } + Err(_) => { + h2.record.status = JobStatus::Failed; + } + } + } + } + + let mut h = handle.lock().await; + h.record.stdout = String::from_utf8_lossy(&stdout_buf).into_owned(); + h.record.stderr = String::from_utf8_lossy(&stderr_buf).into_owned(); + h.record.stdout_truncated = stdout_trunc; + h.record.stderr_truncated = stderr_trunc; + h.record.finished_at_ms = Some(jobs::now_ms()); + }); + + Ok(json!({ + "job_id": id, + "argv": argv, + })) +} + +async fn read_bounded( + reader: &mut R, + limit: usize, + buf: &mut Vec, + truncated: &mut bool, +) { + let mut chunk = [0u8; 8192]; + loop { + match reader.read(&mut chunk).await { + Ok(0) => break, + Ok(n) => { + if buf.len() + n > limit { + let take = limit.saturating_sub(buf.len()); + buf.extend_from_slice(&chunk[..take]); + *truncated = true; + break; + } + buf.extend_from_slice(&chunk[..n]); + } + Err(_) => break, + } + } +} diff --git a/shell/src/functions/kill.rs b/shell/src/functions/kill.rs new file mode 100644 index 000000000..1179fd3c5 --- /dev/null +++ b/shell/src/functions/kill.rs @@ -0,0 +1,54 @@ +use std::future::Future; +use std::pin::Pin; + +use iii_sdk::IIIError; +use serde_json::{json, Value}; + +use crate::jobs::{self, JobStatus}; + +pub fn build_handler() +-> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| Box::pin(async move { handle(payload).await }) +} + +async fn handle(payload: Value) -> Result { + let job_id = payload + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'job_id'".to_string()))?; + + let handle = jobs::get(job_id) + .await + .ok_or_else(|| IIIError::Handler(format!("no such job: {}", job_id)))?; + + let mut h = handle.lock().await; + if h.record.status != JobStatus::Running { + return Ok(json!({ + "job_id": job_id, + "killed": false, + "status": h.record.status, + "reason": "not running", + })); + } + let Some(child) = h.child.as_mut() else { + return Ok(json!({ + "job_id": job_id, + "killed": false, + "status": h.record.status, + "reason": "missing child handle", + })); + }; + child + .start_kill() + .map_err(|e| IIIError::Handler(format!("failed to kill job {}: {}", job_id, e)))?; + h.record.status = JobStatus::Killed; + h.record.finished_at_ms = Some(jobs::now_ms()); + Ok(json!({ + "job_id": job_id, + "killed": true, + "status": h.record.status, + })) +} diff --git a/shell/src/functions/list.rs b/shell/src/functions/list.rs new file mode 100644 index 000000000..a81371f40 --- /dev/null +++ b/shell/src/functions/list.rs @@ -0,0 +1,28 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::IIIError; +use serde_json::{json, Value}; + +use crate::config::ShellConfig; +use crate::jobs; + +pub fn build_handler( + config: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let cfg = config.clone(); + Box::pin(async move { + jobs::remove_old(cfg.job_retention_secs).await; + let all = jobs::list_all().await; + Ok(json!({ + "jobs": all, + "count": all.len(), + })) + }) + } +} diff --git a/shell/src/functions/mod.rs b/shell/src/functions/mod.rs new file mode 100644 index 000000000..4a3517048 --- /dev/null +++ b/shell/src/functions/mod.rs @@ -0,0 +1,5 @@ +pub mod exec; +pub mod exec_bg; +pub mod kill; +pub mod list; +pub mod status; diff --git a/shell/src/functions/status.rs b/shell/src/functions/status.rs new file mode 100644 index 000000000..7976dbc9f --- /dev/null +++ b/shell/src/functions/status.rs @@ -0,0 +1,30 @@ +use std::future::Future; +use std::pin::Pin; + +use iii_sdk::IIIError; +use serde_json::{json, Value}; + +use crate::jobs; + +pub fn build_handler() +-> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| Box::pin(async move { handle(payload).await }) +} + +async fn handle(payload: Value) -> Result { + let job_id = payload + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'job_id'".to_string()))?; + + let handle = jobs::get(job_id) + .await + .ok_or_else(|| IIIError::Handler(format!("no such job: {}", job_id)))?; + let h = handle.lock().await; + Ok(json!({ + "job": h.record, + })) +} diff --git a/shell/src/jobs.rs b/shell/src/jobs.rs new file mode 100644 index 000000000..62eca7849 --- /dev/null +++ b/shell/src/jobs.rs @@ -0,0 +1,144 @@ +use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; +use tokio::process::Child; +use tokio::sync::Mutex; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum JobStatus { + Running, + Finished, + Killed, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobRecord { + pub id: String, + pub argv: Vec, + pub started_at_ms: u64, + pub finished_at_ms: Option, + pub status: JobStatus, + pub exit_code: Option, + pub stdout: String, + pub stderr: String, + pub stdout_truncated: bool, + pub stderr_truncated: bool, +} + +pub struct JobHandle { + pub record: JobRecord, + pub child: Option, +} + +pub struct Jobs { + pub map: Mutex>>>, +} + +impl Jobs { + fn new() -> Self { + Self { + map: Mutex::new(HashMap::new()), + } + } +} + +pub static JOBS: Lazy = Lazy::new(Jobs::new); + +pub fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +pub async fn insert(handle: JobHandle) -> String { + let id = handle.record.id.clone(); + let boxed = Arc::new(Mutex::new(handle)); + JOBS.map.lock().await.insert(id.clone(), boxed); + id +} + +pub async fn get(id: &str) -> Option>> { + JOBS.map.lock().await.get(id).cloned() +} + +// Snapshot the map before awaiting per-job locks. Holding the map guard +// across `handle.lock().await` head-of-line-blocks every other job +// operation (insert, get) for the duration of the iteration. +async fn snapshot() -> Vec<(String, Arc>)> { + let guard = JOBS.map.lock().await; + guard.iter().map(|(k, v)| (k.clone(), v.clone())).collect() +} + +pub async fn remove_old(retention_secs: u64) { + let now = now_ms(); + let threshold_ms = retention_secs.saturating_mul(1000); + let handles = snapshot().await; + let mut to_remove: Vec = Vec::new(); + for (id, handle) in handles { + let h = handle.lock().await; + if let Some(fin) = h.record.finished_at_ms { + if now.saturating_sub(fin) > threshold_ms { + to_remove.push(id); + } + } + } + if !to_remove.is_empty() { + let mut guard = JOBS.map.lock().await; + for id in to_remove { + guard.remove(&id); + } + } +} + +pub async fn list_all() -> Vec { + let handles = snapshot().await; + let mut out = Vec::with_capacity(handles.len()); + for (_, handle) in handles { + out.push(handle.lock().await.record.clone()); + } + out +} + +pub async fn running_count() -> usize { + let handles = snapshot().await; + let mut n = 0; + for (_, handle) in handles { + if handle.lock().await.record.status == JobStatus::Running { + n += 1; + } + } + n +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_insert_and_get() { + let rec = JobRecord { + id: "test-1".to_string(), + argv: vec!["echo".to_string()], + started_at_ms: now_ms(), + finished_at_ms: None, + status: JobStatus::Running, + exit_code: None, + stdout: String::new(), + stderr: String::new(), + stdout_truncated: false, + stderr_truncated: false, + }; + insert(JobHandle { + record: rec.clone(), + child: None, + }) + .await; + let got = get("test-1").await.expect("job exists"); + assert_eq!(got.lock().await.record.id, "test-1"); + } +} diff --git a/shell/src/main.rs b/shell/src/main.rs new file mode 100644 index 000000000..2969c2cec --- /dev/null +++ b/shell/src/main.rs @@ -0,0 +1,198 @@ +use anyhow::Result; +use clap::Parser; +use iii_sdk::{ + register_worker, InitOptions, OtelConfig, RegisterFunctionMessage, RegisterTriggerInput, +}; +use serde_json::json; +use std::sync::Arc; + +mod config; +mod exec; +mod functions; +mod jobs; +mod manifest; + +#[derive(Parser, Debug)] +#[command(name = "iii-shell", about = "Unix shell execution worker for iii agents")] +struct Cli { + #[arg(long, default_value = "./config.yaml")] + config: String, + + #[arg(long, default_value = "ws://127.0.0.1:49134")] + url: String, + + #[arg(long)] + manifest: bool, +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + + let cli = Cli::parse(); + + if cli.manifest { + let m = manifest::build_manifest(); + println!("{}", serde_json::to_string_pretty(&m).unwrap()); + return Ok(()); + } + + let shell_config = match config::load_config(&cli.config) { + Ok(c) => { + tracing::info!( + allowlist_size = c.allowlist.len(), + denylist_size = c.denylist_patterns.len(), + max_timeout_ms = c.max_timeout_ms, + max_concurrent = c.max_concurrent_jobs, + "loaded config from {}", + cli.config + ); + c + } + Err(e) => { + tracing::warn!(error = %e, path = %cli.config, "failed to load config, using defaults"); + let mut c = config::ShellConfig::default(); + c.compile_denylist()?; + c + } + }; + let shared = Arc::new(shell_config); + + tracing::info!(url = %cli.url, "connecting to III engine"); + let iii = register_worker( + &cli.url, + InitOptions { + otel: Some(OtelConfig::default()), + ..Default::default() + }, + ); + + let _exec_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "shell::exec".to_string(), + description: Some("Execute a command and return full output".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": { + "command": { "type": "string", "description": "Program name or full command line if 'args' omitted" }, + "args": { "type": "array", "items": { "type": "string" } }, + "timeout_ms": { "type": "integer" } + }, + "required": ["command"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "exit_code": { "type": ["integer", "null"] }, + "stdout": { "type": "string" }, + "stderr": { "type": "string" }, + "duration_ms": { "type": "integer" }, + "timed_out": { "type": "boolean" }, + "stdout_truncated": { "type": "boolean" }, + "stderr_truncated": { "type": "boolean" } + } + })), + metadata: None, + invocation: None, + }, + functions::exec::build_handler(shared.clone()), + ); + + let _exec_bg_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "shell::exec_bg".to_string(), + description: Some("Spawn a command in background, return job_id".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": { + "command": { "type": "string" }, + "args": { "type": "array", "items": { "type": "string" } } + }, + "required": ["command"] + })), + response_format: Some(json!({ + "type": "object", + "properties": { + "job_id": { "type": "string" }, + "argv": { "type": "array" } + } + })), + metadata: None, + invocation: None, + }, + functions::exec_bg::build_handler(shared.clone()), + ); + + let _kill_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "shell::kill".to_string(), + description: Some("Kill a running background job".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": { "job_id": { "type": "string" } }, + "required": ["job_id"] + })), + response_format: None, + metadata: None, + invocation: None, + }, + functions::kill::build_handler(), + ); + + let _status_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "shell::status".to_string(), + description: Some("Get status of a background job".to_string()), + request_format: Some(json!({ + "type": "object", + "properties": { "job_id": { "type": "string" } }, + "required": ["job_id"] + })), + response_format: None, + metadata: None, + invocation: None, + }, + functions::status::build_handler(), + ); + + let _list_fn = iii.register_function_with( + RegisterFunctionMessage { + id: "shell::list".to_string(), + description: Some("List all background jobs".to_string()), + request_format: Some(json!({ "type": "object", "properties": {} })), + response_format: None, + metadata: None, + invocation: None, + }, + functions::list::build_handler(shared.clone()), + ); + + for (fn_id, path, method) in [ + ("shell::exec", "shell/exec", "POST"), + ("shell::exec_bg", "shell/exec_bg", "POST"), + ("shell::kill", "shell/kill", "POST"), + ("shell::status", "shell/status", "POST"), + ("shell::list", "shell/list", "GET"), + ] { + if let Err(e) = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: fn_id.to_string(), + config: json!({ "api_path": path, "http_method": method }), + metadata: None, + }) { + tracing::warn!(error = %e, "failed to register http trigger for {}", fn_id); + } + } + + tracing::info!("iii-shell registered 5 functions and 5 HTTP triggers, ready"); + + tokio::signal::ctrl_c().await?; + tracing::info!("iii-shell shutting down"); + iii.shutdown_async().await; + Ok(()) +} diff --git a/shell/src/manifest.rs b/shell/src/manifest.rs new file mode 100644 index 000000000..9368e1794 --- /dev/null +++ b/shell/src/manifest.rs @@ -0,0 +1,55 @@ +use serde_json::{json, Value}; + +pub fn build_manifest() -> Value { + json!({ + "name": "iii-shell", + "version": env!("CARGO_PKG_VERSION"), + "description": "Unix shell execution worker for iii agents", + "functions": [ + { + "id": "shell::exec", + "description": "Execute a command synchronously and return stdout/stderr (capped at max_output_bytes; truncation flagged per stream)", + }, + { + "id": "shell::exec_bg", + "description": "Spawn a command in the background and return job_id", + }, + { + "id": "shell::kill", + "description": "Kill a running background job", + }, + { + "id": "shell::status", + "description": "Get status of a background job", + }, + { + "id": "shell::list", + "description": "List all background jobs (running + recently completed)", + }, + ], + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_manifest_has_required_fields() { + let m = build_manifest(); + assert!(m.get("name").is_some()); + assert!(m.get("version").is_some()); + assert!(m.get("functions").is_some()); + let fns = m.get("functions").unwrap().as_array().unwrap(); + assert_eq!(fns.len(), 5); + } + + #[test] + fn test_manifest_json_output() { + let m = build_manifest(); + let s = serde_json::to_string(&m).unwrap(); + assert!(s.contains("shell::exec")); + assert!(s.contains("shell::exec_bg")); + assert!(s.contains("shell::kill")); + } +}