fixup

mullvad · Aug 21, 2024 · 71bff82 · 71bff82
1 parent b7dd3fe
commit 71bff82
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 26 deletions.
diff --git a/test/test-runner/src/main.rs b/test/test-runner/src/main.rs
@@ -388,8 +388,7 @@ impl Service for TestServer {
         cmd.stderr(Stdio::piped());
         cmd.kill_on_drop(true);
 
-        // TODO: do not hardcode
-        let mut child = util::as_unprivileged("mole", || cmd.spawn())
+        let mut child = util::as_unprivileged_user("mole", || cmd.spawn())
             .map_err(|error| {
                 log::error!("Failed to drop privileges: {error}");
                 test_rpc::Error::Syscall

diff --git a/test/test-runner/src/util.rs b/test/test-runner/src/util.rs
@@ -22,47 +22,69 @@ impl<F: FnOnce() + Send> OnDrop<F> {
     }
 }
 
-#[cfg(target_os = "windows")]
-pub fn as_unprivileged<T>(unpriv_user: &str, func: impl FnOnce() -> T) -> Result<T, nix::Error> {
-    // NOTE: no-op
-    let _ = unpriv_user;
-    Ok(func())
+#[derive(thiserror::Error, Debug)]
+#[error(transparent)]
+pub struct Error {
+    inner: InnerError,
 }
 
+#[cfg(unix)]
 #[derive(thiserror::Error, Debug)]
-pub enum Error {
+enum InnerError {
     #[error("Failed to get the specified user")]
     GetUser(#[source] nix::Error),
     #[error("The specified user was not found")]
     MissingUser,
+    #[error("Failed to set uid")]
+    SetUid(#[source] nix::Error),
+    #[error("Failed to set gid")]
+    SetGid(#[source] nix::Error),
+}
+
+#[cfg(target_os = "windows")]
+#[derive(thiserror::Error, Debug)]
+enum InnerError {}
+
+impl From<InnerError> for Error {
+    fn from(inner: InnerError) -> Self {
+        Self { inner }
+    }
+}
+
+#[cfg(target_os = "windows")]
+pub fn as_unprivileged_user<T>(
+    unpriv_user: &str,
+    func: impl FnOnce() -> T,
+) -> Result<T, nix::Error> {
+    // NOTE: no-op
+    let _ = unpriv_user;
+    Ok(func())
 }
 
 #[cfg(unix)]
-pub fn as_unprivileged<T>(unpriv_user: &str, func: impl FnOnce() -> T) -> Result<T, Error> {
+pub fn as_unprivileged_user<T>(unpriv_user: &str, func: impl FnOnce() -> T) -> Result<T, Error> {
     let original_uid = nix::unistd::getuid();
     let original_gid = nix::unistd::getgid();
 
     let user = nix::unistd::User::from_name(unpriv_user)
-        .map_err(Error::GetUser)?
-        .ok_or(Error::MissingUser)?;
+        .map_err(InnerError::GetUser)?
+        .ok_or(InnerError::MissingUser)?;
     let uid = user.uid;
     let gid = user.gid;
 
-    if let Err(error) = nix::unistd::setegid(gid) {
-        log::error!("Failed to set gid: {error}");
-    }
-    if let Err(error) = nix::unistd::seteuid(uid) {
-        log::error!("Failed to set uid: {error}");
-    }
-
-    let func_result = func();
+    nix::unistd::setegid(gid).map_err(InnerError::SetGid)?;
+    OnDrop::new(|| {
+        if let Err(error) = nix::unistd::setegid(original_gid) {
+            log::error!("Failed to restore gid: {error}");
+        }
+    });
 
-    if let Err(error) = nix::unistd::seteuid(original_uid) {
-        log::error!("Failed to restore uid: {error}");
-    }
-    if let Err(error) = nix::unistd::setegid(original_gid) {
-        log::error!("Failed to restore gid: {error}");
-    }
+    nix::unistd::seteuid(uid).map_err(InnerError::SetUid)?;
+    OnDrop::new(|| {
+        if let Err(error) = nix::unistd::seteuid(original_uid) {
+            log::error!("Failed to restore uid: {error}");
+        }
+    });
 
-    Ok(func_result)
+    Ok(func())
 }