use panic::catch_unwind to handle panics

1 files changed, 7 insertions, 21 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 6f6407a..9130582 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,6 +7,7 @@ use pam::module::{PamHandle, PamHooks};
 use std::ffi::CStr;
 use std::fmt::{Display, Write as _};
 use std::io::{ErrorKind, Write};
+use std::panic;
 use std::path::Path;
 use std::process;
 
@@ -30,30 +31,15 @@ impl From<std::io::Error> for SessionError {
     }
 }
 
-trait MaxDisplayLength: Display {
-    const MAX_DISPLAY_LENGTH: usize;
-}
-impl MaxDisplayLength for u32 {
-    const MAX_DISPLAY_LENGTH: usize = u32::MAX.ilog10() as usize + 1;
-}
-
-fn safe_to_string<T: MaxDisplayLength>(v: T) -> Result<String, SessionError> {
-    let mut buf = String::new();
-    buf.try_reserve_exact(T::MAX_DISPLAY_LENGTH)
-        .or(Err(SessionError))?;
-    write!(buf, "{v}").unwrap();
-    Ok(buf)
-}
-
-fn open_session(h: &mut PamHandle, mountpoint: &str) -> Result<(), SessionError> {
+fn open_session(h: &PamHandle, mountpoint: &str) -> Result<(), SessionError> {
     let user = h.get_user(None).or(Err(SessionError))?;
     let user = users::get_user_by_name(&user).ok_or(SessionError)?;
-    let uid = safe_to_string(user.uid())?;
+    let uid = user.uid().to_string();
     let d = Dir::open_ambient_dir(mountpoint, cap_std::ambient_authority())?;
     let d = create_and_open_dir(&d, "user")?;
     let d = create_and_open_dir(&d, &uid)?;
     let d = create_and_open_dir(&d, "leaf")?;
-    let pid = safe_to_string(process::id())?;
+    let pid = process::id().to_string();
     let mut options = OpenOptions::new();
     options.write(true);
     let mut procs = d.open_with("cgroup.procs", &options)?;
@@ -72,9 +58,9 @@ impl PamHooks for PAMUserCG {
         _args: Vec<&CStr>,
         _flags: PamFlag,
     ) -> PamResultCode {
-        match open_session(h, CG_MOUNT) {
-            Ok(()) => PamResultCode::PAM_SUCCESS,
-            Err(SessionError) => PamResultCode::PAM_SESSION_ERR,
+        match panic::catch_unwind(|| open_session(h, CG_MOUNT)) {
+            Ok(Ok(())) => PamResultCode::PAM_SUCCESS,
+            _ => PamResultCode::PAM_SESSION_ERR,
         }
     }
 }