Implement a non-panicking to_string

1 files changed, 20 insertions, 4 deletions

diff --git a/src/lib.rs b/src/lib.rs
index cf93a9c..bd7c949 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ use cap_std::fs::{Dir, OpenOptions};
 use pam::constants::{PamFlag, PamResultCode};
 use pam::module::{PamHandle, PamHooks};
 use std::ffi::CStr;
+use std::fmt::{Display, Write as _};
 use std::io::{ErrorKind, Write};
 use std::path::Path;
 use std::process;
@@ -31,15 +32,30 @@ 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) -> Result<(), SessionError> {
     let user = h.get_user(None).or(Err(SessionError))?;
     let user = users::get_user_by_name(&user).ok_or(SessionError)?;
-    let aa = cap_std::ambient_authority();
-    let d = Dir::open_ambient_dir(CG_MOUNT, aa)?;
+    let uid = safe_to_string(user.uid())?;
+    let d = Dir::open_ambient_dir(CG_MOUNT, cap_std::ambient_authority())?;
     let d = create_and_open_dir(&d, "user")?;
-    let d = create_and_open_dir(&d, &user.uid().to_string())?;
+    let d = create_and_open_dir(&d, &uid)?;
     let d = create_and_open_dir(&d, "leaf")?;
-    let pid = process::id().to_string();
+    let pid = safe_to_string(process::id())?;
     let mut options = OpenOptions::new();
     options.write(true);
     let mut procs = d.open_with("cgroup.procs", &options)?;