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