alice.pds.demo.boltless.dev
This repository has no description
1use anyhow::{Context, Result};
2use nix::sys::signal::{Signal, kill};
3use nix::unistd::{Gid, Pid, Uid, User, getgrouplist, setgid, setgroups, setuid};
4use std::ffi::{CString, OsStr, OsString};
5use std::io;
6use std::os::unix::process::ExitStatusExt;
7use std::path::PathBuf;
8use std::process::Stdio;
9use std::time::Duration;
10use tokio::io::{AsyncRead, AsyncReadExt};
11use tokio::process::{Child, Command};
12use tokio::sync::mpsc::{self, Receiver, Sender};
13use tokio::task::JoinHandle;
14use tracing::warn;
15
16#[derive(Clone, Debug)]
17pub struct Spec {
18 pub program: OsString,
19 pub args: Vec<OsString>,
20 pub env: Vec<(OsString, OsString)>,
21 pub cwd: Option<PathBuf>,
22 pub timeout: Option<Duration>,
23 pub uid: Option<u32>,
24 pub gid: Option<u32>,
25}
26
27impl Spec {
28 pub fn new(program: impl Into<OsString>) -> Self {
29 Self {
30 program: program.into(),
31 args: Vec::new(),
32 env: Vec::new(),
33 cwd: None,
34 timeout: None,
35 uid: None,
36 gid: None,
37 }
38 }
39
40 pub fn arg(mut self, arg: impl Into<OsString>) -> Self {
41 self.args.push(arg.into());
42 self
43 }
44
45 pub fn args<I, S>(mut self, args: I) -> Self
46 where
47 I: IntoIterator<Item = S>,
48 S: Into<OsString>,
49 {
50 self.args.extend(args.into_iter().map(Into::into));
51 self
52 }
53
54 pub fn envs<I, K, V>(mut self, env: I) -> Self
55 where
56 I: IntoIterator<Item = (K, V)>,
57 K: Into<OsString>,
58 V: Into<OsString>,
59 {
60 self.env.extend(
61 env.into_iter()
62 .map(|(key, value)| (key.into(), value.into())),
63 );
64 self
65 }
66
67 pub fn cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
68 self.cwd = Some(cwd.into());
69 self
70 }
71
72 pub fn timeout(mut self, timeout: Duration) -> Self {
73 self.timeout = Some(timeout);
74 self
75 }
76
77 pub fn run_as(mut self, uid: u32, gid: u32) -> Self {
78 self.uid = Some(uid);
79 self.gid = Some(gid);
80 self
81 }
82}
83
84#[derive(Clone, Debug)]
85pub struct ExitResult {
86 pub exit_code: i32,
87 pub error: Option<String>,
88 pub timed_out: bool,
89}
90
91#[derive(Clone, Debug)]
92pub struct CaptureOutput {
93 pub exit: ExitResult,
94 pub stdout: Vec<u8>,
95 pub stderr: Vec<u8>,
96}
97
98impl CaptureOutput {
99 pub fn success(&self) -> bool {
100 self.exit.exit_code == 0 && self.exit.error.is_none()
101 }
102
103 pub fn combined_lossy(&self) -> String {
104 let mut data = self.stdout.clone();
105 data.extend_from_slice(&self.stderr);
106 String::from_utf8_lossy(&data).trim().to_owned()
107 }
108}
109
110#[derive(Clone, Copy, Debug)]
111pub enum OutKind {
112 Stdout,
113 Stderr,
114}
115
116#[derive(Clone, Debug)]
117pub struct OutData {
118 pub data: Vec<u8>,
119 pub kind: OutKind,
120}
121
122pub struct StreamingCommand {
123 events: Receiver<OutData>,
124 exit: JoinHandle<Result<ExitResult>>,
125}
126
127impl StreamingCommand {
128 pub fn into_parts(self) -> (Receiver<OutData>, JoinHandle<Result<ExitResult>>) {
129 (self.events, self.exit)
130 }
131}
132
133pub async fn run_capture(spec: Spec) -> Result<CaptureOutput> {
134 let running = spawn_streaming(spec)?;
135 let mut stdout = Vec::new();
136 let mut stderr = Vec::new();
137 let (mut events, exit_task) = running.into_parts();
138
139 while let Some(event) = events.recv().await {
140 match event.kind {
141 OutKind::Stdout => stdout.extend_from_slice(&event.data),
142 OutKind::Stderr => stderr.extend_from_slice(&event.data),
143 }
144 }
145
146 let exit = exit_task
147 .await
148 .unwrap_or_else(|error| Err(anyhow::anyhow!("command supervisor failed: {error}")))?;
149 Ok(CaptureOutput {
150 exit,
151 stdout,
152 stderr,
153 })
154}
155
156fn read_oom_kill_count() -> u32 {
157 let content = match std::fs::read_to_string("/proc/vmstat") {
158 Ok(c) => c,
159 Err(_) => return 0,
160 };
161 for line in content.lines() {
162 let mut parts = line.split_whitespace();
163 if parts.next() == Some("oom_kill") {
164 if let Some(count_str) = parts.next() {
165 if let Ok(count) = count_str.parse::<u32>() {
166 return count;
167 }
168 }
169 }
170 }
171 0
172}
173
174pub fn spawn_streaming(mut spec: Spec) -> Result<StreamingCommand> {
175 let oom_kill_before = read_oom_kill_count();
176 let mut child = spawn(&mut spec)?;
177 let stdout = child.stdout.take().context("stdout pipe missing")?;
178 let stderr = child.stderr.take().context("stderr pipe missing")?;
179
180 let (events_tx, events_rx) = mpsc::channel(64);
181 let stdout_thread = spawn_reader(stdout, events_tx.clone(), OutKind::Stdout);
182 let stderr_thread = spawn_reader(stderr, events_tx.clone(), OutKind::Stderr);
183 drop(events_tx);
184
185 let exit = tokio::spawn(async move {
186 let exit = wait_child(&mut child, spec.timeout, oom_kill_before).await;
187
188 // ensure all output is observed before exiting
189 // this assumes children dont daemonize and hold onto the stdout/err
190 stdout_thread.await.context("stdout reader task failed")?;
191 stderr_thread.await.context("stderr reader task failed")?;
192
193 Ok(exit)
194 });
195
196 Ok(StreamingCommand {
197 events: events_rx,
198 exit,
199 })
200}
201
202fn spawn(spec: &mut Spec) -> Result<Child> {
203 let mut cmd = Command::new(&spec.program);
204 cmd.args(&spec.args)
205 .envs(spec.env.iter().map(|(key, value)| (key, value)))
206 .stdout(Stdio::piped())
207 .stderr(Stdio::piped());
208
209 if let Some(cwd) = &spec.cwd {
210 cmd.current_dir(cwd);
211 }
212
213 // don't use rust's .uid() / .gid() methods here because they clear
214 // supplemantary groups, which means for example adding a user to "docker"
215 // group won't actually let it access the sock.
216 // https://github.com/rust-lang/rust/issues/90747
217 if let (Some(uid), Some(gid)) = (spec.uid, spec.gid) {
218 let username = User::from_uid(Uid::from_raw(uid))
219 .ok()
220 .flatten()
221 .map(|u| u.name)
222 .with_context(|| format!("lookup passwd entry for uid {uid}"))?;
223 let cname = CString::new(username)
224 .with_context(|| format!("username for uid {uid} contained a null byte"))?;
225 // resolve groups beforehand so we don't have to read /etc/group in the pre_exec
226 let groups =
227 getgrouplist(&cname, Gid::from_raw(gid)).context("resolve supplementary groups")?;
228 // SAFETY: pre_exec runs between fork and execve in the child.
229 // we only call async-signal-safe syscalls and we don't touch any
230 // shared state, no allocator, no mutexes, no globals.
231 unsafe {
232 cmd.pre_exec(move || {
233 setgroups(&groups).map_err(io::Error::from)?;
234 setgid(Gid::from_raw(gid)).map_err(io::Error::from)?;
235 setuid(Uid::from_raw(uid)).map_err(io::Error::from)?;
236 Ok(())
237 });
238 }
239 }
240
241 // allow us to kill this whole process tree on deadline
242 cmd.process_group(0);
243
244 cmd.spawn()
245 .with_context(|| format!("spawn {}", display_os(&spec.program)))
246}
247
248async fn wait_child(
249 child: &mut Child,
250 timeout: Option<Duration>,
251 oom_kill_before: u32,
252) -> ExitResult {
253 let wait = child.wait();
254 let status = match timeout {
255 Some(timeout) => match tokio::time::timeout(timeout, wait).await {
256 Ok(status) => status,
257 Err(_) => {
258 if let Some(pid) = child.id()
259 && let Err(error) = kill(Pid::from_raw(-(pid as i32)), Signal::SIGKILL)
260 {
261 warn!(pid, %error, "failed to kill process group");
262 }
263 let _ = child.wait().await;
264 return ExitResult {
265 exit_code: 124,
266 error: Some("command timed out".to_owned()),
267 timed_out: true,
268 };
269 }
270 },
271 None => wait.await,
272 };
273
274 match status {
275 Ok(status) => {
276 let code = status.code();
277 let signal = status.signal();
278 let exit_code = code.or_else(|| signal.map(|sig| 128 + sig)).unwrap_or(1);
279
280 let mut error = None;
281 if signal == Some(9) {
282 let oom_kill_after = read_oom_kill_count();
283 if oom_kill_after > oom_kill_before {
284 error = Some("guest process killed by guest kernel OOM".to_owned());
285 }
286 }
287
288 ExitResult {
289 exit_code,
290 error,
291 timed_out: false,
292 }
293 }
294 Err(error) => ExitResult {
295 exit_code: 1,
296 error: Some(error.to_string()),
297 timed_out: false,
298 },
299 }
300}
301
302fn spawn_reader(
303 mut reader: impl AsyncRead + Unpin + Send + 'static,
304 events: Sender<OutData>,
305 kind: OutKind,
306) -> JoinHandle<()> {
307 tokio::spawn(async move {
308 let mut buf = [0_u8; 32 * 1024];
309 loop {
310 match reader.read(&mut buf).await {
311 Ok(0) => return,
312 Ok(n) => {
313 let event = OutData {
314 data: buf[..n].to_vec(),
315 kind,
316 };
317 if events.send(event).await.is_err() {
318 return;
319 }
320 }
321 Err(error) => {
322 warn!(%error, "failed to read command stream");
323 return;
324 }
325 }
326 }
327 })
328}
329
330fn display_os(value: &OsStr) -> String {
331 value.to_string_lossy().into_owned()
332}