Skip to main content

isolate_integration/
sandbox.rs

1use anyhow::{Context, Result, bail};
2use std::collections::HashMap;
3use std::path::PathBuf;
4use tokio::fs;
5use tokio::process::Command;
6
7/// Directory binding rule
8#[derive(Debug, Clone)]
9pub struct DirectoryRule {
10    pub inside_path: PathBuf,
11    pub outside_path: Option<PathBuf>,
12    pub options: DirectoryOptions,
13}
14
15/// Directory binding options
16/// NOTE: Unless --no-default-dirs is specified, the default set of directory rules binds /bin, /dev (with devices allowed), /lib, /lib64 (if it exists), and /usr.
17/// It also binds the working directory to /box (read-write), mounts the proc filesystem at /proc, and creates a temporary directory /tmp.
18#[derive(Debug, Clone, Default)]
19pub struct DirectoryOptions {
20    pub read_write: bool,
21    pub allow_devices: bool,
22    pub no_exec: bool,
23    pub maybe: bool,
24    pub is_filesystem: bool,
25    pub is_tmp: bool,
26    pub no_recursive: bool,
27}
28
29/// Environment variable rule
30#[derive(Debug, Clone)]
31pub enum EnvRule {
32    Inherit(String),
33    Set(String, String),
34    FullEnv,
35}
36
37/// Resource limits for isolate sandbox
38/// All size-related items are in kilobytes (kB), time-related items are in seconds (s).
39/// NOTE: use cgroups-related options first to control memory precisely.
40#[derive(Debug, Clone)]
41pub struct ResourceLimits {
42    pub time_limit: Option<f64>,
43    pub wall_time_limit: Option<f64>,
44    pub extra_time: Option<f64>,
45    pub memory_limit: Option<u32>,
46    pub cg_memory_limit: Option<u32>,
47    pub stack_limit: Option<u32>,
48    pub open_files_limit: Option<u32>,
49    pub file_size_limit: Option<u32>,
50    pub core_limit: Option<u32>,
51    pub process_limit: Option<u32>,
52    pub quota: Option<(u32, u32)>,
53}
54
55/// Special options for isolate
56#[derive(Debug, Clone, Default)]
57pub struct SpecialOptions {
58    pub share_net: bool,
59    pub inherit_fds: bool,
60    pub tty_hack: bool,
61    pub special_files: bool,
62    pub use_cgroups: bool,
63    pub no_default_dirs: bool,
64    pub verbose: bool,
65    pub silent: bool,
66    pub wait: bool,
67    pub as_uid: Option<u32>,
68    pub as_gid: Option<u32>,
69}
70
71/// Execution result from isolate
72#[derive(Debug, Clone)]
73pub struct ExecutionResult {
74    pub exit_code: Option<i32>,
75    pub signal: Option<i32>,
76    pub time_used: f64,
77    pub wall_time_used: f64,
78    pub memory_used: u32,
79    pub cg_memory_used: Option<u32>,
80    pub killed: bool,
81    pub cg_oom_killed: bool,
82    pub status: String,
83    pub message: String,
84    pub stdout: String,
85    pub stderr: String,
86}
87
88/// Main isolate sandbox implementation
89pub struct IsolateSandbox {
90    pub box_id: u32,
91    pub isolate_bin: String,
92    pub directory_rules: Vec<DirectoryRule>,
93    pub env_rules: Vec<EnvRule>,
94    pub stdin_file: Option<String>,
95    pub stdout_file: Option<String>,
96    pub stderr_file: Option<String>,
97    pub stderr_to_stdout: bool,
98    pub chdir: Option<String>,
99    pub meta_file: Option<PathBuf>,
100    pub special_options: SpecialOptions,
101}
102
103impl ResourceLimits {
104    pub fn new() -> Self {
105        Self {
106            time_limit: None,
107            wall_time_limit: None,
108            extra_time: None,
109            memory_limit: None,
110            cg_memory_limit: None,
111            stack_limit: None,
112            open_files_limit: None,
113            file_size_limit: None,
114            core_limit: None,
115            process_limit: None,
116            quota: None,
117        }
118    }
119
120    pub fn with_time_limit(mut self, seconds: f64) -> Self {
121        self.time_limit = Some(seconds);
122        self
123    }
124
125    pub fn with_memory_limit(mut self, kilobytes: u32) -> Self {
126        self.memory_limit = Some(kilobytes);
127        self
128    }
129
130    pub fn with_wall_time_limit(mut self, seconds: f64) -> Self {
131        self.wall_time_limit = Some(seconds);
132        self
133    }
134
135    pub fn with_cg_memory_limit(mut self, kilobytes: u32) -> Self {
136        self.cg_memory_limit = Some(kilobytes);
137        self
138    }
139
140    pub fn with_process_limit(mut self, count: u32) -> Self {
141        self.process_limit = Some(count);
142        self
143    }
144}
145
146impl DirectoryRule {
147    pub fn bind(inside: impl Into<PathBuf>, outside: impl Into<PathBuf>) -> Self {
148        Self {
149            inside_path: inside.into(),
150            outside_path: Some(outside.into()),
151            options: DirectoryOptions::default(),
152        }
153    }
154
155    pub fn bind_same(path: impl Into<PathBuf>) -> Self {
156        let path = path.into();
157        Self {
158            inside_path: path.clone(),
159            outside_path: Some(path),
160            options: DirectoryOptions::default(),
161        }
162    }
163
164    pub fn tmp(inside: impl Into<PathBuf>) -> Self {
165        Self {
166            inside_path: inside.into(),
167            outside_path: None,
168            options: DirectoryOptions {
169                is_tmp: true,
170                read_write: true,
171                ..Default::default()
172            },
173        }
174    }
175
176    pub fn filesystem(name: impl Into<PathBuf>) -> Self {
177        Self {
178            inside_path: name.into(),
179            outside_path: None,
180            options: DirectoryOptions {
181                is_filesystem: true,
182                ..Default::default()
183            },
184        }
185    }
186
187    pub fn read_write(mut self) -> Self {
188        self.options.read_write = true;
189        self
190    }
191
192    pub fn allow_devices(mut self) -> Self {
193        self.options.allow_devices = true;
194        self
195    }
196
197    pub fn no_exec(mut self) -> Self {
198        self.options.no_exec = true;
199        self
200    }
201
202    pub fn maybe(mut self) -> Self {
203        self.options.maybe = true;
204        self
205    }
206
207    pub fn no_recursive(mut self) -> Self {
208        self.options.no_recursive = true;
209        self
210    }
211}
212
213impl IsolateSandbox {
214    pub fn new(box_id: u32) -> Self {
215        Self {
216            box_id,
217            isolate_bin: std::env::var("ISOLATE_BIN").unwrap_or_else(|_| "isolate".to_string()),
218            directory_rules: Vec::new(),
219            env_rules: vec![EnvRule::Set(
220                "LIBC_FATAL_STDERR_".to_string(), // send fatal errors to stderr by default
221                "1".to_string(),
222            )],
223            stdin_file: None,
224            stdout_file: None,
225            stderr_file: None,
226            stderr_to_stdout: false,
227            chdir: None,
228            meta_file: None,
229            special_options: Default::default(),
230        }
231    }
232
233    /// Initialize the sandbox
234    pub async fn init(&self, limits: &ResourceLimits) -> Result<()> {
235        let mut cmd = Command::new(&self.isolate_bin);
236
237        cmd.arg(format!("--box-id={}", self.box_id));
238        cmd.arg("--init");
239
240        if let Some((blocks, inodes)) = limits.quota {
241            cmd.arg(format!("--quota={},{}", blocks, inodes));
242        }
243
244        // Add special options
245        if self.special_options.use_cgroups {
246            cmd.arg("--cg");
247        }
248        if self.special_options.verbose {
249            cmd.arg("--verbose");
250        }
251        if self.special_options.silent {
252            cmd.arg("--silent");
253        }
254        if self.special_options.wait {
255            cmd.arg("--wait");
256        }
257        if let Some(uid) = self.special_options.as_uid {
258            cmd.arg(format!("--as-uid={}", uid));
259        }
260        if let Some(gid) = self.special_options.as_gid {
261            cmd.arg(format!("--as-gid={}", gid));
262        }
263
264        let output = cmd
265            .output()
266            .await
267            .context("Failed to execute isolate --init")?;
268
269        if !output.status.success() {
270            let stderr = String::from_utf8_lossy(&output.stderr);
271            bail!("isolate --init failed: {}", stderr);
272        }
273
274        Ok(())
275    }
276
277    /// Run a command in the sandbox
278    pub async fn run<I, S>(
279        &self,
280        program: &str,
281        args: I,
282        limits: &ResourceLimits,
283    ) -> Result<ExecutionResult>
284    where
285        I: IntoIterator<Item = S>,
286        S: AsRef<str>,
287    {
288        let mut cmd = Command::new(&self.isolate_bin);
289
290        cmd.arg(format!("--box-id={}", self.box_id));
291        cmd.arg("--run");
292
293        // Add resource limits
294        if let Some(time) = limits.time_limit {
295            cmd.arg(format!("--time={}", time));
296        }
297        if let Some(wall_time) = limits.wall_time_limit {
298            cmd.arg(format!("--wall-time={}", wall_time));
299        }
300        if let Some(extra_time) = limits.extra_time {
301            cmd.arg(format!("--extra-time={}", extra_time));
302        }
303        if let Some(memory) = limits.memory_limit {
304            cmd.arg(format!("--mem={}", memory));
305        }
306        if let Some(cg_memory) = limits.cg_memory_limit {
307            cmd.arg(format!("--cg-mem={}", cg_memory));
308        }
309        if let Some(stack) = limits.stack_limit {
310            cmd.arg(format!("--stack={}", stack));
311        }
312        if let Some(open_files) = limits.open_files_limit {
313            cmd.arg(format!("--open-files={}", open_files));
314        }
315        if let Some(file_size) = limits.file_size_limit {
316            cmd.arg(format!("--fsize={}", file_size));
317        }
318        if let Some(core) = limits.core_limit {
319            cmd.arg(format!("--core={}", core));
320        }
321        match limits.process_limit {
322            Some(processes) if processes > 0 => {
323                cmd.arg(format!("--processes={}", processes));
324            }
325            _ => {
326                cmd.arg("--processes");
327            }
328        }
329
330        // Add I/O redirection
331        if let Some(ref stdin) = self.stdin_file {
332            cmd.arg(format!("--stdin={}", stdin));
333        }
334        if let Some(ref stdout) = self.stdout_file {
335            cmd.arg(format!("--stdout={}", stdout));
336        }
337        if let Some(ref stderr) = self.stderr_file {
338            cmd.arg(format!("--stderr={}", stderr));
339        }
340        if self.stderr_to_stdout {
341            cmd.arg("--stderr-to-stdout");
342        }
343
344        // Add directory change
345        if let Some(ref chdir) = self.chdir {
346            cmd.arg(format!("--chdir={}", chdir));
347        }
348
349        // Add meta file
350        if let Some(ref meta) = self.meta_file {
351            cmd.arg(format!("--meta={}", meta.display()));
352        }
353
354        // Add directory rules
355        if self.special_options.no_default_dirs {
356            cmd.arg("--no-default-dirs");
357        }
358        for rule in &self.directory_rules {
359            let mut dir_arg = if rule.options.is_filesystem {
360                format!("{}:fs", rule.inside_path.display())
361            } else if rule.options.is_tmp {
362                format!("{}:tmp", rule.inside_path.display())
363            } else if let Some(ref outside) = rule.outside_path {
364                format!("{}={}", rule.inside_path.display(), outside.display())
365            } else {
366                rule.inside_path.display().to_string()
367            };
368
369            let mut options = Vec::new();
370            if rule.options.read_write {
371                options.push("rw");
372            }
373            if rule.options.allow_devices {
374                options.push("dev");
375            }
376            if rule.options.no_exec {
377                options.push("noexec");
378            }
379            if rule.options.maybe {
380                options.push("maybe");
381            }
382            if rule.options.no_recursive {
383                options.push("norec");
384            }
385
386            if !options.is_empty() {
387                dir_arg.push(':');
388                dir_arg.push_str(&options.join(","));
389            }
390
391            cmd.arg(format!("--dir={}", dir_arg));
392        }
393
394        // Add environment rules
395        for rule in &self.env_rules {
396            match rule {
397                EnvRule::Inherit(var) => {
398                    cmd.arg(format!("--env={}", var));
399                }
400                EnvRule::Set(var, value) => {
401                    cmd.arg(format!("--env={}={}", var, value));
402                }
403                EnvRule::FullEnv => {
404                    cmd.arg("--full-env");
405                }
406            }
407        }
408
409        // Add special options
410        if self.special_options.use_cgroups {
411            cmd.arg("--cg");
412        }
413        if self.special_options.share_net {
414            cmd.arg("--share-net");
415        }
416        if self.special_options.inherit_fds {
417            cmd.arg("--inherit-fds");
418        }
419        if self.special_options.tty_hack {
420            cmd.arg("--tty-hack");
421        }
422        if self.special_options.special_files {
423            cmd.arg("--special-files");
424        }
425        if self.special_options.verbose {
426            cmd.arg("--verbose");
427        }
428        if self.special_options.silent {
429            cmd.arg("--silent");
430        }
431
432        // Add the command to execute
433        cmd.arg("--").arg(program);
434        for arg in args {
435            cmd.arg(arg.as_ref());
436        }
437
438        // print the command in string for debugging to stderr
439        let command_to_string = |cmd: &Command| -> String {
440            let program = cmd.as_std().get_program().to_string_lossy();
441            let args: Vec<String> = cmd
442                .as_std()
443                .get_args()
444                .map(|arg| arg.to_string_lossy().to_string())
445                .collect();
446            format!("{} {}", program, args.join(" "))
447        };
448        eprintln!("Executing command: {}", command_to_string(&cmd));
449
450        let output = cmd
451            .output()
452            .await
453            .context("Failed to execute isolate --run")?;
454
455        // Read output files if specified
456        let stdout = if let Some(ref stdout_file) = self.stdout_file {
457            fs::read_to_string(stdout_file).await.unwrap_or_default()
458        } else {
459            String::from_utf8_lossy(&output.stdout).to_string()
460        };
461
462        let stderr = if let Some(ref stderr_file) = self.stderr_file {
463            fs::read_to_string(stderr_file).await.unwrap_or_default()
464        } else {
465            String::from_utf8_lossy(&output.stderr).to_string()
466        };
467
468        // Parse metadata if available
469        let metadata = if let Some(ref meta_file) = self.meta_file {
470            self.parse_metadata(meta_file).await.unwrap_or_default()
471        } else {
472            HashMap::new()
473        };
474
475        Ok(ExecutionResult {
476            exit_code: metadata.get("exitcode").and_then(|s| s.parse().ok()),
477            signal: metadata.get("exitsig").and_then(|s| s.parse().ok()),
478            time_used: metadata
479                .get("time")
480                .and_then(|s| s.parse().ok())
481                .unwrap_or(0.0),
482            wall_time_used: metadata
483                .get("time-wall")
484                .and_then(|s| s.parse().ok())
485                .unwrap_or(0.0),
486            memory_used: metadata
487                .get("max-rss")
488                .and_then(|s| s.parse().ok())
489                .unwrap_or(0),
490            cg_memory_used: metadata.get("cg-mem").and_then(|s| s.parse().ok()),
491            killed: metadata.get("killed").map(|s| s == "1").unwrap_or(false),
492            cg_oom_killed: metadata.get("cg-oom-killed").is_some(),
493            status: metadata.get("status").cloned().unwrap_or_default(),
494            message: metadata.get("message").cloned().unwrap_or_default(),
495            stdout,
496            stderr,
497        })
498    }
499
500    /// Cleanup the sandbox
501    pub async fn cleanup(&self) -> Result<()> {
502        let mut cmd = Command::new(&self.isolate_bin);
503
504        cmd.arg(format!("--box-id={}", self.box_id));
505        cmd.arg("--cleanup");
506
507        if self.special_options.use_cgroups {
508            cmd.arg("--cg");
509        }
510
511        let output = cmd
512            .output()
513            .await
514            .context("Failed to execute isolate --cleanup")?;
515
516        if !output.status.success() {
517            let stderr = String::from_utf8_lossy(&output.stderr);
518            eprintln!("isolate --cleanup warning: {}", stderr);
519        }
520
521        Ok(())
522    }
523
524    async fn parse_metadata(&self, meta_file: &PathBuf) -> Result<HashMap<String, String>> {
525        let content = fs::read_to_string(meta_file)
526            .await
527            .context("Failed to read metadata file")?;
528
529        let mut metadata = HashMap::new();
530        for line in content.lines() {
531            if let Some((key, value)) = line.split_once(':') {
532                metadata.insert(key.to_string(), value.to_string());
533            }
534        }
535
536        Ok(metadata)
537    }
538
539    /// The following are builder options.
540
541    pub fn with_directory_rule(mut self, rule: DirectoryRule) -> Self {
542        self.directory_rules.push(rule);
543        self
544    }
545
546    pub fn with_env_rule(mut self, rule: EnvRule) -> Self {
547        self.env_rules.push(rule);
548        self
549    }
550
551    pub fn with_stdin(mut self, file: impl Into<String>) -> Self {
552        self.stdin_file = Some(file.into());
553        self
554    }
555
556    pub fn with_stdout(mut self, file: impl Into<String>) -> Self {
557        self.stdout_file = Some(file.into());
558        self
559    }
560
561    pub fn with_stderr(mut self, file: impl Into<String>) -> Self {
562        self.stderr_file = Some(file.into());
563        self
564    }
565
566    pub fn with_stderr_to_stdout(mut self) -> Self {
567        self.stderr_to_stdout = true;
568        self
569    }
570
571    pub fn with_chdir(mut self, dir: impl Into<String>) -> Self {
572        self.chdir = Some(dir.into());
573        self
574    }
575
576    pub fn with_meta_file(mut self, file: impl Into<PathBuf>) -> Self {
577        self.meta_file = Some(file.into());
578        self
579    }
580
581    pub fn with_special_options(mut self, options: SpecialOptions) -> Self {
582        self.special_options = options;
583        self
584    }
585
586    pub fn use_cgroups(mut self) -> Self {
587        self.special_options.use_cgroups = true;
588        self
589    }
590
591    pub fn disable_cgroups(mut self) -> Self {
592        self.special_options.use_cgroups = false;
593        self
594    }
595
596    pub fn share_network(mut self) -> Self {
597        self.special_options.share_net = true;
598        self
599    }
600
601    pub fn no_default_dirs(mut self) -> Self {
602        self.special_options.no_default_dirs = true;
603        self
604    }
605
606    pub fn verbose(mut self) -> Self {
607        self.special_options.verbose = true;
608        self
609    }
610}