clawft_kernel/workload_runtime/
supervise.rs

1//! cog-runner-style process supervision (ADR-100 section 3): a cleared
2//! environment, rlimits applied in the child before exec, an unprivileged
3//! user, a wall-clock cap, capped output capture, and terminate-then-kill.
4
5use std::path::PathBuf;
6use std::sync::{Arc, Mutex};
7use std::time::{Duration, Instant};
8
9use tokio::io::AsyncReadExt;
10use tokio::process::{Child, Command};
11use tokio::task::JoinHandle;
12
13use super::evidence::{Capture, RunEvidence};
14use super::types::RuntimeError;
15
16/// Grace before SIGKILL when a run exceeds its wall-clock cap.
17pub const TIMEOUT_GRACE: Duration = Duration::from_secs(2);
18
19/// Resource limits applied with `setrlimit` in the child.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct ProcLimits {
22    /// Address-space cap (`RLIMIT_AS`) in bytes. Enforced on Linux; macOS
23    /// does not implement it, so there it is best effort.
24    pub mem_bytes: u64,
25    /// CPU seconds (`RLIMIT_CPU`), for bounded runs.
26    pub cpu_secs: Option<u64>,
27    /// Open files (`RLIMIT_NOFILE`).
28    pub nofile: u64,
29    /// Largest file the process may write (`RLIMIT_FSIZE`).
30    pub fsize_bytes: u64,
31}
32
33impl ProcLimits {
34    /// Limits for a cog with `ram_mb` memory.
35    pub fn for_ram_mb(ram_mb: u32) -> Self {
36        Self {
37            mem_bytes: u64::from(ram_mb) * 1024 * 1024,
38            cpu_secs: None,
39            nofile: 256,
40            fsize_bytes: 64 * 1024 * 1024,
41        }
42    }
43}
44
45/// What to launch.
46#[derive(Debug, Clone)]
47pub struct LaunchSpec {
48    /// Executable.
49    pub program: PathBuf,
50    /// Arguments.
51    pub args: Vec<String>,
52    /// Full environment (the parent's is cleared).
53    pub env: Vec<(String, String)>,
54    /// Working directory.
55    pub cwd: PathBuf,
56    /// rlimits.
57    pub limits: ProcLimits,
58    /// Bytes captured per stream.
59    pub output_limit: usize,
60    /// `(uid, gid)` to drop to; required when the supervisor runs as root.
61    pub run_as: Option<(u32, u32)>,
62}
63
64/// A running supervised process.
65pub struct Supervised {
66    child: Child,
67    out: Arc<Mutex<Capture>>,
68    err: Arc<Mutex<Capture>>,
69    readers: Vec<JoinHandle<()>>,
70    started: Instant,
71    /// Process group (the leader's pid) until the group has been dealt with;
72    /// kept so the group can be signalled after the leader is reaped.
73    pgid: Option<libc::pid_t>,
74    args: Vec<String>,
75    exited: Option<std::process::ExitStatus>,
76}
77
78/// True when this process runs as root.
79pub fn running_as_root() -> bool {
80    // SAFETY: geteuid has no preconditions and cannot fail.
81    unsafe { libc::geteuid() == 0 }
82}
83
84fn set_limit(resource: libc::c_int, value: u64) -> std::io::Result<()> {
85    let lim = libc::rlimit {
86        rlim_cur: value as libc::rlim_t,
87        rlim_max: value as libc::rlim_t,
88    };
89    // SAFETY: plain syscall on a valid, initialized struct.
90    if unsafe { libc::setrlimit(resource as _, &lim) } == 0 {
91        Ok(())
92    } else {
93        Err(std::io::Error::last_os_error())
94    }
95}
96
97fn apply_in_child(limits: ProcLimits, run_as: Option<(u32, u32)>) -> std::io::Result<()> {
98    let as_result = set_limit(libc::RLIMIT_AS as libc::c_int, limits.mem_bytes);
99    if cfg!(target_os = "linux") {
100        as_result?;
101    }
102    if let Some(cpu) = limits.cpu_secs {
103        set_limit(libc::RLIMIT_CPU as libc::c_int, cpu)?;
104    }
105    set_limit(libc::RLIMIT_NOFILE as libc::c_int, limits.nofile)?;
106    set_limit(libc::RLIMIT_FSIZE as libc::c_int, limits.fsize_bytes)?;
107    set_limit(libc::RLIMIT_CORE as libc::c_int, 0)?;
108    if let Some((uid, gid)) = run_as {
109        // SAFETY: async-signal-safe syscalls between fork and exec.
110        unsafe {
111            if libc::setgroups(0, std::ptr::null()) != 0
112                || libc::setgid(gid as libc::gid_t) != 0
113                || libc::setuid(uid as libc::uid_t) != 0
114            {
115                return Err(std::io::Error::last_os_error());
116            }
117        }
118    }
119    Ok(())
120}
121
122fn pump<R: tokio::io::AsyncRead + Unpin + Send + 'static>(
123    mut r: R,
124    sink: Arc<Mutex<Capture>>,
125) -> JoinHandle<()> {
126    tokio::spawn(async move {
127        let mut buf = [0u8; 8192];
128        loop {
129            match r.read(&mut buf).await {
130                Ok(0) | Err(_) => break,
131                Ok(n) => {
132                    if let Ok(mut c) = sink.lock() {
133                        c.push(&buf[..n]);
134                    }
135                }
136            }
137        }
138    })
139}
140
141impl Supervised {
142    /// Launch under limits. Refuses to run as root without `run_as`.
143    pub fn spawn(spec: LaunchSpec) -> Result<Self, RuntimeError> {
144        if running_as_root() && spec.run_as.is_none() {
145            return Err(RuntimeError::AdmissionRefused(
146                "native workloads run unprivileged: configure run_as when the host runs as root"
147                    .into(),
148            ));
149        }
150        if let Some((0, _)) = spec.run_as {
151            return Err(RuntimeError::InvalidConfig(
152                "run_as uid 0 is not unprivileged".into(),
153            ));
154        }
155        let mut cmd = Command::new(&spec.program);
156        cmd.args(&spec.args)
157            .env_clear()
158            .env("PATH", "/usr/local/bin:/usr/bin:/bin")
159            // After the default so a caller that must reach a toolchain
160            // outside the default PATH (model servers) can say so.
161            .envs(spec.env.iter().cloned())
162            .current_dir(&spec.cwd)
163            .stdin(std::process::Stdio::null())
164            .stdout(std::process::Stdio::piped())
165            .stderr(std::process::Stdio::piped())
166            .kill_on_drop(true)
167            .process_group(0);
168        let (limits, run_as) = (spec.limits, spec.run_as);
169        // SAFETY: the closure only makes async-signal-safe syscalls.
170        unsafe {
171            cmd.pre_exec(move || apply_in_child(limits, run_as));
172        }
173        let mut child = cmd
174            .spawn()
175            .map_err(|e| RuntimeError::Backend(format!("spawn {}: {e}", spec.program.display())))?;
176        let pgid = child.id().map(|p| p as libc::pid_t);
177        let out = Arc::new(Mutex::new(Capture::new(spec.output_limit)));
178        let err = Arc::new(Mutex::new(Capture::new(spec.output_limit)));
179        let mut readers = Vec::new();
180        if let Some(o) = child.stdout.take() {
181            readers.push(pump(o, out.clone()));
182        }
183        if let Some(e) = child.stderr.take() {
184            readers.push(pump(e, err.clone()));
185        }
186        Ok(Self {
187            child,
188            out,
189            err,
190            readers,
191            started: Instant::now(),
192            pgid,
193            args: spec.args,
194            exited: None,
195        })
196    }
197
198    /// OS process id, while running.
199    pub fn pid(&self) -> Option<u32> {
200        self.child.id()
201    }
202
203    /// Non-blocking exit check.
204    pub fn try_exit(&mut self) -> Option<std::process::ExitStatus> {
205        if self.exited.is_none() {
206            self.exited = self.child.try_wait().ok().flatten();
207        }
208        self.exited
209    }
210
211    fn signal_group(&self, sig: libc::c_int) {
212        if let Some(pgid) = self.pgid {
213            // SAFETY: signalling the process group this supervisor created.
214            unsafe {
215                libc::kill(-pgid, sig);
216            }
217        }
218    }
219
220    /// Wait up to `max`; on timeout terminate, then kill after
221    /// [`TIMEOUT_GRACE`], and mark the evidence `killed_for_timeout`.
222    pub async fn wait_for(mut self, max: Duration) -> RunEvidence {
223        let timed_out = tokio::time::timeout(max, self.child.wait()).await.is_err();
224        if timed_out {
225            self.terminate_inner(TIMEOUT_GRACE).await;
226        }
227        let mut ev = self.finish().await;
228        ev.killed_for_timeout = timed_out;
229        ev
230    }
231
232    /// SIGTERM the group, SIGKILL after `grace`, and collect evidence.
233    pub async fn terminate(mut self, grace: Duration) -> RunEvidence {
234        if self.try_exit().is_none() {
235            self.terminate_inner(grace).await;
236        }
237        self.finish().await
238    }
239
240    async fn terminate_inner(&mut self, grace: Duration) {
241        self.signal_group(libc::SIGTERM);
242        if tokio::time::timeout(grace, self.child.wait())
243            .await
244            .is_err()
245        {
246            self.signal_group(libc::SIGKILL);
247            let _ = self.child.kill().await;
248        }
249    }
250
251    async fn finish(mut self) -> RunEvidence {
252        let status = match self.exited {
253            Some(s) => Some(s),
254            None => self.child.wait().await.ok(),
255        };
256        // The leader is gone: nothing it forked may outlive it.
257        self.signal_group(libc::SIGKILL);
258        self.pgid = None;
259        // Readers finish at EOF; a grandchild holding the pipe open is
260        // bounded by this wait.
261        for r in self.readers.drain(..) {
262            let _ = tokio::time::timeout(Duration::from_secs(2), r).await;
263        }
264        use std::os::unix::process::ExitStatusExt;
265        let ev = RunEvidence {
266            args: self.args.clone(),
267            exit_code: status.and_then(|s| s.code()),
268            signal: status.and_then(|s| s.signal()),
269            elapsed_ms: self.started.elapsed().as_millis() as u64,
270            ..RunEvidence::default()
271        };
272        let out = self
273            .out
274            .lock()
275            .map(|c| c.clone())
276            .unwrap_or_else(|_| Capture::new(0));
277        let err = self
278            .err
279            .lock()
280            .map(|c| c.clone())
281            .unwrap_or_else(|_| Capture::new(0));
282        ev.with_output(&out, &err)
283    }
284}
285
286impl Drop for Supervised {
287    /// A supervisor dropped while its process still runs takes the whole
288    /// group with it (`kill_on_drop` alone only reaches the leader).
289    fn drop(&mut self) {
290        self.signal_group(libc::SIGKILL);
291    }
292}
clawft_kernel/workload_runtime/supervise.rs

clawft_kernel/workload_runtime/
supervise.rs
