feat(rsjudge-runner): ✨ add require_caps function, remove repetition …

…from user module

crates/rsjudge-runner/examples/exploit.rs

@@ -1,8 +1,7 @@
 use std::{path::PathBuf, process::Command};
 
-use anyhow::anyhow;
-use caps::{has_cap, read, CapSet, Capability};
-use rsjudge_runner::{user::builder, RunAs};
+use caps::{read, CapSet, Capability};
+use rsjudge_runner::{require_caps, user::builder, RunAs};
 use rsjudge_utils::command::check_output;
 
 /// An attempt to exploit the runner by running a binary with a setuid call.
@@ -22,19 +21,11 @@ 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,
-    ]
-    .into_iter()
-    .try_for_each(|cap| {
-        if has_cap(None, CapSet::Effective, cap)? {
-            Ok(())
-        } else {
-            Err(anyhow!("Missing capability: {:?}", cap))
-        }
-    })?;
+    ])?;
 
     // Get the path to the examples.
     // This crate is located at crates/rsjudge-runner,

crates/rsjudge-runner/src/caps_check.rs

@@ -0,0 +1,29 @@
+use std::convert::identity;
+
+use caps::{errors::CapsError, has_cap, Capability};
+
+use crate::error::{Error, Result};
+
+pub fn require_caps<I>(caps: I) -> Result<()>
+where
+    I: IntoIterator<Item = Capability>,
+{
+    let missing_caps = caps
+        .into_iter()
+        .map(|cap| match has_cap(None, caps::CapSet::Effective, cap) {
+            Ok(has_cap) => Ok((!has_cap).then_some(cap)),
+            Err(e) => Err(e),
+        })
+        .collect::<Result<Vec<_>, CapsError>>()?
+        .into_iter()
+        .filter_map(identity)
+        .collect::<Vec<_>>();
+
+    if missing_caps.is_empty() {
+        Ok(())
+    } else {
+        Err(Error::CapsRequired {
+            caps: missing_caps.into_boxed_slice(),
+        })
+    }
+}

crates/rsjudge-runner/src/error.rs

@@ -1,17 +1,23 @@
-use caps::Capability;
+use caps::{errors::CapsError, Capability};
 use thiserror::Error;
 
 #[derive(Debug, Error)]
 pub enum Error {
-    #[error("{cap} is required to run as another user.")]
-    CapsRequired { cap: Capability },
-    #[error("User '{name}' not found")]
-    UserNotFound { name: String },
+    /// Capabilities required but not set.
+    #[error("{caps:?} required but not set.")]
+    CapsRequired { caps: Box<[Capability]> },
 
-    #[error("I/O error: {0}")]
+    /// The requested user is not found.
+    #[error("User '{username}' not found")]
+    UserNotFound { username: &'static str },
+
+    /// A wrapper for `std::io::Error`.
+    #[error(transparent)]
     Io(#[from] std::io::Error),
-    #[error("{0}")]
-    CapsError(#[from] caps::errors::CapsError),
+
+    /// A wrapper for `caps::errors::CapsError`.
+    #[error(transparent)]
+    CapsError(#[from] CapsError),
 }
 
-pub type Result<T> = std::result::Result<T, Error>;
+pub type Result<T, E = Error> = std::result::Result<T, E>;

crates/rsjudge-runner/src/lib.rs

@@ -2,13 +2,17 @@
 
 use std::{os::unix::process::CommandExt as _, process::Command};
 
-use caps::{has_cap, CapSet, Capability};
+use caps::Capability;
 use nix::unistd::{setgroups, Gid};
 use uzers::User;
 
-use crate::error::{Error, Result};
+pub use crate::{
+    caps_check::require_caps,
+    error::{Error, Result},
+};
 
-pub mod error;
+mod caps_check;
+mod error;
 pub mod user;
 pub trait RunAs {
     type Error;
@@ -18,17 +22,11 @@ pub trait RunAs {
 impl RunAs for Command {
     type Error = Error;
     fn run_as(&mut self, user: &User) -> Result<&mut Self> {
-        if !has_cap(None, CapSet::Effective, Capability::CAP_SETUID)? {
-            Err(Error::CapsRequired {
-                cap: Capability::CAP_SETUID,
-            })?;
-        }
-
-        has_cap(None, CapSet::Effective, Capability::CAP_SETGID).map_err(|_| {
-            Error::CapsRequired {
-                cap: Capability::CAP_SETGID,
-            }
-        })?;
+        require_caps([
+            Capability::CAP_SETUID,
+            Capability::CAP_SETGID,
+            Capability::CAP_DAC_READ_SEARCH,
+        ])?;
 
         let uid = user.uid();
         let gid = user.primary_group_id();

crates/rsjudge-runner/src/user.rs

@@ -1,36 +1,34 @@
+//! Functions to get user instances.
+//!
+//! All functions return a reference to a static instance of [`uzers::User`] if succeeded.
+
 use std::sync::OnceLock;
 
 use uzers::{get_user_by_name, User};
 
 use crate::error::{Error, Result};
 
-pub static SUPERVISOR: OnceLock<Option<User>> = OnceLock::new();
-pub static BUILDER: OnceLock<Option<User>> = OnceLock::new();
-pub static RUNNER: OnceLock<Option<User>> = OnceLock::new();
-
-pub fn supervisor<'a>() -> Result<&'a User> {
-    SUPERVISOR
-        .get_or_init(|| get_user_by_name("rsjudge-supervisor"))
-        .as_ref()
-        .ok_or_else(|| Error::UserNotFound {
-            name: "rsjudge-supervisor".to_string(),
-        })
-}
-
-pub fn builder<'a>() -> Result<&'a User> {
-    BUILDER
-        .get_or_init(|| get_user_by_name("rsjudge-builder"))
-        .as_ref()
-        .ok_or_else(|| Error::UserNotFound {
-            name: "rsjudge-builder".to_string(),
-        })
+/// Generate functions to get user instances.
+macro_rules! users {
+    ($($vis:vis fn $id:ident() => $name:literal);* $(;)?) => {
+        $(
+            #[doc = concat!("Get an instance of user `", $name, "`.")]
+            ///
+            /// # Errors
+            /// Returns an error if the user is not found.
+            pub fn $id() -> Result<&'static User> {
+                static INNER: OnceLock<Option<User>> = OnceLock::new();
+                INNER
+                    .get_or_init(|| get_user_by_name($name))
+                    .as_ref()
+                    .ok_or_else(|| Error::UserNotFound { username: $name })
+            }
+        )*
+    };
 }
 
-pub fn runner<'a>() -> Result<&'a User> {
-    RUNNER
-        .get_or_init(|| get_user_by_name("rsjudge-runner"))
-        .as_ref()
-        .ok_or_else(|| Error::UserNotFound {
-            name: "rsjudge-runner".to_string(),
-        })
+users! {
+    pub fn supervisor() => "rsjudge-supervisor";
+    pub fn builder() => "rsjudge-builder";
+    pub fn runner() => "rsjudge-runner";
 }