1use anyhow::{Context, Result, bail};
2use std::collections::HashMap;
3use std::path::PathBuf;
4use tokio::fs;
5use tokio::process::Command;
6
7#[derive(Debug, Clone)]
9pub struct DirectoryRule {
10 pub inside_path: PathBuf,
11 pub outside_path: Option<PathBuf>,
12 pub options: DirectoryOptions,
13}
14
15#[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#[derive(Debug, Clone)]
31pub enum EnvRule {
32 Inherit(String),
33 Set(String, String),
34 FullEnv,
35}
36
37#[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#[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#[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
88pub 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(), "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 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 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 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 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 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 if let Some(ref chdir) = self.chdir {
346 cmd.arg(format!("--chdir={}", chdir));
347 }
348
349 if let Some(ref meta) = self.meta_file {
351 cmd.arg(format!("--meta={}", meta.display()));
352 }
353
354 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 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 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 cmd.arg("--").arg(program);
434 for arg in args {
435 cmd.arg(arg.as_ref());
436 }
437
438 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 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 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 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 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}