fixup

mullvad · Aug 21, 2024 · ed0dde2 · ed0dde2
1 parent b7dd3fe
commit ed0dde2
Showing 1 changed file with 42 additions and 23 deletions.
diff --git a/test/test-runner/src/util.rs b/test/test-runner/src/util.rs
@@ -22,19 +22,40 @@ 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<T>(unpriv_user: &str, func: impl FnOnce() -> T) -> Result<T, nix::Error> {
+    // NOTE: no-op
+    let _ = unpriv_user;
+    Ok(func())
 }
 
 #[cfg(unix)]
@@ -43,26 +64,24 @@ pub fn as_unprivileged<T>(unpriv_user: &str, func: impl FnOnce() -> T) -> Result
     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())
 }