feat: split syscall filter

Author: yjl9903

src/context/builder.rs

@@ -1,5 +1,5 @@
 use crate::context::{CatBoxCompileContext, CatBoxContext, CatBoxJudgeContext, CatBoxRunContext};
-use crate::syscall::SyscallFilter;
+use crate::syscall::{RestrictedSyscall, SyscallFilter};
 use crate::utils::mount::MountPoint;
 use crate::utils::{into_c_string, parse_env, GidType, MemoryLimitType, TimeLimitType, UidType};
 use crate::{CatBox, CatBoxError, CatBoxOption};
@@ -22,6 +22,7 @@ pub struct CatBoxBuilder {
   memory_limit: Option<MemoryLimitType>,
   uid: Option<UidType>,
   gid: Option<GidType>,
+  cwd: Option<PathBuf>,
 }
 
 /// Build CatBox running option
@@ -42,6 +43,7 @@ impl CatBoxBuilder {
       memory_limit: None,
       uid: None,
       gid: None,
+      cwd: None,
     }
   }
 
@@ -93,6 +95,10 @@ impl CatBoxBuilder {
     if let Some(gid) = self.gid {
       option.gid = Gid::from(gid);
     }
+    // Set default cwd
+    if let Some(cwd) = &self.cwd {
+      option.cwd = cwd.clone();
+    }
     // Set default env
     for env_pair in self.env.iter() {
       option.env.push(env_pair.clone());
@@ -125,8 +131,8 @@ impl CatBoxBuilder {
   }
 
   /// Set default force mode
-  pub fn set_default_force(mut self, flag: Option<bool>) -> Self {
-    self.force = flag;
+  pub fn set_default_force(mut self, flag: bool) -> Self {
+    self.force = Some(flag);
     self
   }
 
@@ -152,6 +158,12 @@ impl CatBoxBuilder {
     self
   }
 
+  /// Set default cwd
+  pub fn set_default_cwd(mut self, path: Option<PathBuf>) -> Self {
+    self.cwd = path;
+    self
+  }
+
   /// Parse default env list
   pub fn parse_env_list(mut self, list: Vec<String>) -> Result<Self, CatBoxError> {
     for env_var in list {
@@ -273,13 +285,29 @@ impl CatBoxOptionBuilder {
     }
   }
 
-  /// Set ptrace syscall filter
-  pub fn set_ptrace(mut self, flag: bool) -> Self {
-    if flag {
-      self.option.ptrace = Some(SyscallFilter::default());
-    } else {
-      self.option.ptrace = None;
+  /// Parse ptrace syscall filter
+  pub fn parse_ptrace_presets(mut self, presets: Option<Vec<String>>) -> Result<Self, CatBoxError> {
+    if let Some(presets) = presets {
+      self.option.ptrace = SyscallFilter::parse_presets(presets)?;
     }
+    Ok(self)
+  }
+
+  /// Set ptrace feature
+  pub fn ptrace(mut self, preset: RestrictedSyscall) -> Self {
+    let mut filter = self
+      .option
+      .ptrace
+      .get_or_insert(SyscallFilter::new())
+      .to_owned();
+    filter.enable(preset);
+    self.option.ptrace = Some(filter);
+    self
+  }
+
+  /// Disable ptrace
+  pub fn disable_ptrace(mut self) -> Self {
+    self.option.ptrace = None;
     self
   }
 
@@ -291,6 +319,14 @@ impl CatBoxOptionBuilder {
     self
   }
 
+  /// Set work directory in chroot or not
+  pub fn set_cwd(mut self, path: Option<PathBuf>) -> Self {
+    if let Some(path) = path {
+      self.option.cwd = path;
+    }
+    self
+  }
+
   /// Set work directory in chroot
   pub fn cwd<P: Into<PathBuf>>(mut self, path: P) -> Self {
     self.option.cwd = path.into();

src/lib.rs

@@ -36,6 +36,7 @@
 pub use catbox::run;
 pub use context::{CatBox, CatBoxBuilder, CatBoxOption, CatBoxOptionBuilder, CatBoxResult};
 pub use error::CatBoxError;
+pub use syscall::{RestrictedSyscall, SyscallFilter};
 
 mod catbox;
 mod cgroup;

src/main.rs

@@ -11,7 +11,7 @@ use crate::catbox::run;
 use crate::context::{CatBox, CatBoxBuilder, CatBoxOption};
 use crate::error::{CatBoxError, CatBoxExit};
 // use crate::preset::make_compile_params;
-use crate::utils::{default_format, MemoryLimitType, TimeLimitType};
+use crate::utils::{default_format, GidType, MemoryLimitType, TimeLimitType, UidType};
 
 mod catbox;
 mod cgroup;
@@ -30,26 +30,29 @@ struct Cli {
   #[arg(long, requires = "report", help = "Output JSON format report")]
   json: bool,
 
-  #[arg(short, long, help = "Time limit (unit: ms)")]
+  #[arg(short, long, help = "Time limit (unit: ms) [default: 1000]")]
   time: Option<TimeLimitType>,
 
-  #[arg(short, long, help = "Memory limit (unit: KB)")]
+  #[arg(short, long, help = "Memory limit (unit: KB) [default: 262144]")]
   memory: Option<MemoryLimitType>,
 
   #[arg(long, value_name = "KEY=VALUE", help = "Pass environment variables")]
   env: Vec<String>,
 
+  #[arg(long, help = "Current working directory [default: ./]")]
+  cwd: Option<PathBuf>,
+
   #[arg(long, help = "Child process uid")]
-  uid: Option<u32>,
+  uid: Option<UidType>,
 
   #[arg(long, help = "Child process gid")]
-  gid: Option<u32>,
+  gid: Option<GidType>,
 
-  #[arg(long, help = "Run in current user")]
+  #[arg(long, help = "Run in current user [default: false]")]
   user: bool,
 
-  #[arg(short, long, help = "Force security control")]
-  force: Option<bool>,
+  #[arg(short, long, help = "Force security control [default: false]")]
+  force: bool,
 
   #[structopt(subcommand)]
   command: Commands,
@@ -83,11 +86,15 @@ enum Commands {
     #[arg(long, help = "The number of processes")]
     process: Option<u64>,
 
+    #[arg(
+      long,
+      value_name = "PRESET",
+      help = "Enable ptrace presets [support: none|net|process|all]"
+    )]
+    ptrace: Option<Vec<String>>,
+
     #[arg(long, help = "Disable chroot")]
     no_chroot: bool,
-
-    #[arg(long, help = "Disable ptrace")]
-    no_ptrace: bool,
   },
 
   #[command(about = "Compile user code")]
@@ -139,6 +146,7 @@ impl Cli {
     .set_current_user(self.user)
     .set_default_uid(self.uid)
     .set_default_gid(self.gid)
+    .set_default_cwd(self.cwd)
     .parse_env_list(self.env)?;
 
     let catbox = match self.command {
@@ -151,16 +159,16 @@ impl Cli {
         read,
         write,
         process,
+        ptrace,
         no_chroot,
-        no_ptrace,
       } => builder
         .command(program, arguments)
         .set_process(process)
         .set_stdin(stdin)
         .set_stdout(stdout)
         .set_stderr(stderr)
-        .set_ptrace(!no_ptrace)
         .set_chroot(!no_chroot)
+        .parse_ptrace_presets(ptrace)?
         .parse_mount_read(read)?
         .parse_mount_write(write)?
         .done(),

src/syscall.rs

@@ -3,60 +3,121 @@ use std::collections::HashMap;
 use std::ffi::{c_long, c_ulonglong};
 use std::fmt::{Debug, Formatter};
 
+use crate::CatBoxError;
 use nix::libc::{
-  user_regs_struct, SYS_accept, SYS_accept4, SYS_bind, SYS_clone, SYS_connect, SYS_execve,
-  SYS_execveat, SYS_fork, SYS_getpeername, SYS_getsockname, SYS_getsockopt, SYS_listen,
+  user_regs_struct, SYS_accept, SYS_accept4, SYS_bind, SYS_clone, SYS_clone3, SYS_connect,
+  SYS_execve, SYS_execveat, SYS_fork, SYS_getpeername, SYS_getsockname, SYS_getsockopt, SYS_listen,
   SYS_setsockopt, SYS_shutdown, SYS_socket, SYS_socketpair, SYS_vfork,
 };
 use nix::unistd::Pid;
 
 type SyscallId = c_ulonglong;
 
-/// 禁止系统调用
-/// 允许有限次系统调用
+/// Syscall permission
 #[derive(Clone)]
 pub enum SyscallPerm {
+  /// Forbid all
   Forbid,
+  /// Use a filter function to check whether it is ok
   FilterFn(fn(pid: &Pid, regs: &user_regs_struct) -> bool),
+  /// Allow a few times
   Allow(i32),
 }
 
-/// 系统调用过滤器
-/// 黑名单过滤，若不在映射内，则允许；否则，禁止或者允许有限次
+/// Syscall filter
+/// It is a black list filter, and it supports forbid syscall or allow a few times
 #[derive(Debug, Clone)]
 pub struct SyscallFilter {
   map: HashMap<SyscallId, SyscallPerm>,
 }
 
+/// Syscall filter preset category
+#[derive(Debug, Copy, Clone)]
+pub enum RestrictedSyscall {
+  Net,
+  Process,
+  Thread,
+}
+
 impl SyscallFilter {
-  pub fn default() -> Self {
-    let mut filter = SyscallFilter {
+  /// Create an empty syscall filter
+  pub fn new() -> Self {
+    let filter = SyscallFilter {
       map: HashMap::new(),
     };
-    // 禁用网络
     filter
-      .forbid(SYS_socket)
-      .forbid(SYS_socketpair)
-      .forbid(SYS_setsockopt)
-      .forbid(SYS_getsockopt)
-      .forbid(SYS_getsockname)
-      .forbid(SYS_getpeername)
-      .forbid(SYS_bind)
-      .forbid(SYS_listen)
-      .forbid(SYS_accept)
-      .forbid(SYS_accept4)
-      .forbid(SYS_connect)
-      .forbid(SYS_shutdown);
-    // 禁用进程相关
+  }
+
+  /// Create a default syscall filter with all the presets open
+  pub fn default() -> Self {
+    let mut filter = Self::new();
     filter
-      .allow(SYS_execve, 1)
-      .allow(SYS_execveat, 1)
-      .forbid(SYS_fork)
-      .forbid(SYS_vfork)
-      .forbid(SYS_clone);
+      .enable(RestrictedSyscall::Net)
+      .enable(RestrictedSyscall::Process);
     filter
   }
 
+  /// Enable preset
+  pub fn enable(self: &mut Self, feature: RestrictedSyscall) -> &mut Self {
+    match feature {
+      RestrictedSyscall::Net => {
+        self
+          .forbid(SYS_socket)
+          .forbid(SYS_socketpair)
+          .forbid(SYS_setsockopt)
+          .forbid(SYS_getsockopt)
+          .forbid(SYS_getsockname)
+          .forbid(SYS_getpeername)
+          .forbid(SYS_bind)
+          .forbid(SYS_listen)
+          .forbid(SYS_accept)
+          .forbid(SYS_accept4)
+          .forbid(SYS_connect)
+          .forbid(SYS_shutdown);
+      }
+      RestrictedSyscall::Process => {
+        self
+          .allow(SYS_execve, 1)
+          .allow(SYS_execveat, 1)
+          .forbid(SYS_fork)
+          .forbid(SYS_vfork)
+          .forbid(SYS_clone)
+          .forbid(SYS_clone3);
+      }
+      RestrictedSyscall::Thread => {}
+    };
+    self
+  }
+
+  /// Try parsing presets string
+  pub fn parse_presets(presets: Vec<String>) -> Result<Option<Self>, CatBoxError> {
+    let mut filter = Self::new();
+    let presets = presets
+      .into_iter()
+      .flat_map(|str| str.split(" ").map(str::to_owned).collect::<Vec<_>>())
+      .map(|p| p.trim().to_ascii_lowercase())
+      .filter(|p| p.len() > 0)
+      .collect::<Vec<String>>();
+    for preset in presets {
+      match preset.as_str() {
+        "none" => return Ok(None),
+        "net" | "network" => {
+          filter.enable(RestrictedSyscall::Net);
+        }
+        "process" => {
+          filter.enable(RestrictedSyscall::Process);
+        }
+        "all" => {
+          filter
+            .enable(RestrictedSyscall::Net)
+            .enable(RestrictedSyscall::Process);
+        }
+        _ => return Err(CatBoxError::cli("Parse ptrace syscall filter string fails")),
+      };
+    }
+    Ok(Some(filter))
+  }
+
   pub fn forbid(self: &mut Self, id: c_long) -> &mut Self {
     self.map.insert(id as SyscallId, SyscallPerm::forbid());
     self

tests/aplusb.rs

@@ -1,4 +1,4 @@
-use catj::{run, CatBoxBuilder, CatBoxResult};
+use catj::{run, CatBoxBuilder, CatBoxResult, RestrictedSyscall};
 use log::info;
 use nix::sys::signal::Signal;
 use std::env::current_dir;
@@ -35,7 +35,7 @@ fn compile_cpp(dir: &PathBuf, file: &String) -> String {
     .stdout("/dev/null")
     // .stderr("/dev/null")
     .current_user()
-    .set_ptrace(false)
+    .disable_ptrace()
     .process(10)
     .chroot()
     .cwd(current_dir().unwrap())