fix(rsjudge-runner): 🍻 remove problematic caps_check

crates/rsjudge-grpc/Cargo.toml

@@ -9,7 +9,7 @@ license.workspace = true
 rust-version.workspace = true
 
 [package.metadata.cargo-machete]
-ignored = ["prost"]
+ignored = ["prost", "prost-types"]
 
 [dependencies]
 anyhow = "1.0.82"

crates/rsjudge-runner/examples/demo.rs

@@ -3,9 +3,10 @@
 use std::process::Command;
 
 use anyhow::anyhow;
+use caps::Capability;
 use rsjudge_runner::{
     user::{builder, runner},
-    RunAs,
+    CapHandle, RunAs,
 };
 use uzers::{get_current_uid, get_user_by_uid};
 fn main() -> anyhow::Result<()> {
@@ -14,6 +15,9 @@ fn main() -> anyhow::Result<()> {
         .output()?;
     println!("{}", String::from_utf8_lossy(&self_output.stdout));
 
+    CapHandle::new(Capability::CAP_SETUID)?;
+    CapHandle::new(Capability::CAP_SETGID)?;
+
     let builder_output = Command::new("id").run_as(builder()?)?.output()?;
     println!("{}", String::from_utf8_lossy(&builder_output.stdout));
 

crates/rsjudge-runner/examples/exploit.rs

@@ -2,8 +2,8 @@
 
 use std::{path::PathBuf, process::Command};
 
-use caps::{read, CapSet, Capability};
-use rsjudge_runner::{require_caps, user::builder, RunAs};
+use caps::{read, CapSet};
+use rsjudge_runner::{user::builder, RunAs};
 use rsjudge_utils::command::check_output;
 
 /// An attempt to exploit the runner by running a binary with a setuid call.
@@ -23,12 +23,6 @@ fn main() -> anyhow::Result<()> {
     dbg!(read(None, CapSet::Inheritable).unwrap());
     dbg!(read(None, CapSet::Permitted).unwrap());
 
-    require_caps([
-        Capability::CAP_SETUID,
-        Capability::CAP_SETGID,
-        Capability::CAP_DAC_READ_SEARCH,
-    ])?;
-
     // Get the path to the examples.
     // This crate is located at crates/rsjudge-runner,
     // so we need to go up two levels to find the workspace root.

crates/rsjudge-runner/src/cap_handle.rs

@@ -0,0 +1,49 @@
+use std::{cell::RefCell, collections::HashMap, rc::Rc};
+
+use caps::{drop as drop_cap, has_cap, raise as raise_cap, Capability};
+
+use crate::Result;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct CapHandle {
+    cap: Capability,
+    ref_count: Rc<()>,
+}
+
+impl CapHandle {
+    thread_local! {
+        /// Local capability reference count.
+        static LOCAL_CAPS: RefCell<HashMap<Capability,Rc<()>>> = RefCell::new(HashMap::new());
+    }
+
+    pub fn new(cap: Capability) -> Result<Self> {
+        let ref_count = Self::LOCAL_CAPS
+            .with_borrow_mut(|local_caps| local_caps.entry(cap).or_default().clone());
+        try_raise_cap(cap)?;
+        Ok(Self { cap, ref_count })
+    }
+}
+
+impl Drop for CapHandle {
+    fn drop(&mut self) {
+        if Rc::strong_count(&self.ref_count) == 1 {
+            // Last reference.
+            let _ = drop_cap(None, caps::CapSet::Effective, self.cap);
+            Self::LOCAL_CAPS.with_borrow_mut(|local_caps| {
+                local_caps.remove(&self.cap);
+            });
+        }
+    }
+}
+
+fn try_raise_cap(cap: Capability) -> Result<bool> {
+    if has_cap(None, caps::CapSet::Effective, cap)? {
+        // Already has cap.
+        Ok(true)
+    } else if has_cap(None, caps::CapSet::Permitted, cap)? {
+        raise_cap(None, caps::CapSet::Effective, cap)?;
+        Ok(true)
+    } else {
+        Ok(false)
+    }
+}

crates/rsjudge-runner/src/lib.rs

@@ -2,23 +2,28 @@
 
 #![cfg_attr(not(test), warn(clippy::print_stdout, clippy::print_stderr))]
 
-use std::{os::unix::process::CommandExt as _, process::Command};
+use std::{
+    io::{self, ErrorKind},
+    os::unix::process::CommandExt as _,
+    process::Command,
+};
 
 use caps::Capability;
 use nix::unistd::{setgroups, Gid};
 use rsjudge_utils::log_if_error;
 use uzers::User;
 
 pub use crate::{
-    caps_check::require_caps,
+    cap_handle::CapHandle,
     error::{Error, Result},
 };
 
-mod caps_check;
 mod error;
 
 #[macro_use]
 mod user_macro;
+
+mod cap_handle;
 pub mod user;
 
 pub trait RunAs {
@@ -29,12 +34,6 @@ pub trait RunAs {
 impl RunAs for Command {
     type Error = Error;
     fn run_as(&mut self, user: &User) -> Result<&mut Self> {
-        log_if_error!(require_caps([
-            Capability::CAP_SETUID,
-            Capability::CAP_SETGID,
-            Capability::CAP_DAC_READ_SEARCH,
-        ]))?;
-
         let uid = user.uid();
         let gid = user.primary_group_id();
 
@@ -53,6 +52,8 @@ impl RunAs for Command {
                 .map(|g| Gid::from_raw(g.gid()))
                 .collect();
             let set_groups = move || {
+                CapHandle::new(Capability::CAP_SETGID)
+                    .map_err(|e| io::Error::new(ErrorKind::PermissionDenied, e.to_string()))?;
                 log_if_error!(setgroups(&groups))?;
                 Ok(())
             };