Move towards ensuring safe code cannot cause UB even within internal …

…libraries

Signed-off-by: Alex Saveau <[email protected]>

src/core/files.rs

@@ -53,7 +53,7 @@ pub fn create_files_and_dirs(
 #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace"))]
 fn create_dirs(num_dirs: usize, dir: &mut FastPathBuf) -> Result<(), io::Error> {
     for i in 0..num_dirs {
-        with_dir_name(i, |s| dir.push(s));
+        let mut dir = with_dir_name(i, |s| dir.push(s));
 
         create_dir_all(&dir)
             .attach_printable_lazy(|| format!("Failed to create directory {dir:?}"))?;
@@ -78,20 +78,20 @@ fn create_files(
 
     let mut start_file = 0;
     if num_files > 0 {
-        with_file_name(offset, |s| file.push(s));
+        let mut guard = with_file_name(offset, |s| file.push(s));
 
-        match contents.create_file(file, 0, true, &mut state) {
+        match contents.create_file(&mut guard, 0, true, &mut state) {
             Ok(bytes) => {
                 bytes_written += bytes;
                 start_file += 1;
-                file.pop();
+                guard.pop();
             }
             Err(e) => {
                 if e.kind() == NotFound {
                     #[cfg(feature = "tracing")]
-                    tracing::event!(tracing::Level::TRACE, file = ?file, "Parent directory not created in time");
+                    tracing::event!(tracing::Level::TRACE, file = ?guard, "Parent directory not created in time");
 
-                    file.pop();
+                    guard.pop();
                     create_dir_all(&file)
                         .attach_printable_lazy(|| format!("Failed to create directory {file:?}"))?;
                 } else {
@@ -102,10 +102,15 @@ fn create_files(
         }
     }
     for i in start_file..num_files {
-        with_file_name(i + offset, |s| file.push(s));
+        let mut file = with_file_name(i + offset, |s| file.push(s));
 
         bytes_written += contents
-            .create_file(file, i.try_into().unwrap_or(usize::MAX), false, &mut state)
+            .create_file(
+                &mut file,
+                i.try_into().unwrap_or(usize::MAX),
+                false,
+                &mut state,
+            )
             .attach_printable_lazy(|| format!("Failed to create file {file:?}"))?;
 
         file.pop();

src/core/scheduler.rs

@@ -184,7 +184,9 @@ pub async fn run(
             scheduler.stack.push(directory);
             with_dir_name(0, |s| scheduler.target_dir.push(s));
         } else if !is_completing {
-            with_dir_name(next_stack_dir, |s| scheduler.target_dir.set_file_name(s));
+            with_dir_name(next_stack_dir, |s| unsafe {
+                scheduler.target_dir.set_file_name(s);
+            });
         }
     }
     #[cfg(feature = "tracing")]
@@ -474,10 +476,12 @@ fn handle_directory_completion(
         ref child_dir_counts,
     }) = stack.last()
     {
-        target_dir.pop();
+        unsafe {
+            target_dir.pop();
+        }
 
         if !child_dir_counts.is_empty() {
-            with_dir_name(total_dirs - child_dir_counts.len(), |s| {
+            with_dir_name(total_dirs - child_dir_counts.len(), |s| unsafe {
                 target_dir.set_file_name(s);
             });
         }

src/utils/fast_path.rs

@@ -1,25 +1,27 @@
 use std::{
     ffi::OsStr,
     fmt,
-    ops::Deref,
+    ops::{Deref, DerefMut},
     os::unix::ffi::{OsStrExt, OsStringExt},
     path::{Path, PathBuf, MAIN_SEPARATOR},
 };
 
 /// A specialized [`PathBuf`][std::path::PathBuf] implementation that takes
 /// advantage of a few assumptions. Specifically, it *only* supports adding
 /// single-level directories (e.g. "foo", "foo/bar" is not allowed) and updating
-/// the current file name. Any other usage is UB.
+/// the current file name.
 pub struct FastPathBuf {
     inner: Vec<u8>,
     last_len: usize,
 }
 
 impl FastPathBuf {
+    #[must_use]
     pub fn new() -> Self {
         Self::with_capacity(0)
     }
 
+    #[must_use]
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
             inner: Vec::with_capacity(capacity),
@@ -36,22 +38,12 @@ impl FastPathBuf {
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace"))]
-    pub fn push(&mut self, name: &str) {
-        let Self {
-            ref mut inner,
-            ref mut last_len,
-        } = *self;
-
-        *last_len = inner.len();
-
-        // Reserve an extra byte for the eventually appended NUL terminator
-        inner.reserve(1 + name.len() + 1);
-        inner.push(MAIN_SEPARATOR as u8);
-        inner.extend_from_slice(name.as_bytes());
+    pub fn push(&mut self, name: &str) -> PopGuard {
+        PopGuard::push(self, name)
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace"))]
-    pub fn pop(&mut self) {
+    pub unsafe fn pop(&mut self) {
         let Self {
             ref mut inner,
             last_len,
@@ -69,12 +61,15 @@ impl FastPathBuf {
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace"))]
-    pub fn set_file_name(&mut self, name: &str) {
-        self.pop();
+    pub unsafe fn set_file_name(&mut self, name: &str) {
+        unsafe {
+            self.pop();
+        }
         self.push(name);
     }
 
     #[cfg(all(unix, not(miri)))]
+    #[must_use]
     pub fn to_cstr_mut(&mut self) -> unix::CStrFastPathBufGuard {
         unix::CStrFastPathBufGuard::new(self)
     }
@@ -94,6 +89,12 @@ fn bytes_as_path(bytes: &[u8]) -> &Path {
     OsStr::from_bytes(bytes).as_ref()
 }
 
+impl Default for FastPathBuf {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl Deref for FastPathBuf {
     type Target = Path;
 
@@ -143,6 +144,56 @@ impl Clone for FastPathBuf {
     }
 }
 
+pub struct PopGuard<'a>(&'a mut FastPathBuf);
+
+impl<'a> PopGuard<'a> {
+    fn push(path: &'a mut FastPathBuf, name: &str) -> Self {
+        let FastPathBuf {
+            ref mut inner,
+            ref mut last_len,
+        } = *path;
+
+        *last_len = inner.len();
+
+        // Reserve an extra byte for the eventually appended NUL terminator
+        inner.reserve(1 + name.len() + 1);
+        inner.push(MAIN_SEPARATOR as u8);
+        inner.extend_from_slice(name.as_bytes());
+
+        Self(path)
+    }
+
+    pub fn pop(&mut self) {
+        unsafe { self.0.pop() }
+    }
+}
+
+impl Deref for PopGuard<'_> {
+    type Target = FastPathBuf;
+
+    fn deref(&self) -> &Self::Target {
+        self.0
+    }
+}
+
+impl DerefMut for PopGuard<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0
+    }
+}
+
+impl AsRef<Path> for PopGuard<'_> {
+    fn as_ref(&self) -> &Path {
+        self.0
+    }
+}
+
+impl fmt::Debug for PopGuard<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(&**self, f)
+    }
+}
+
 #[cfg(all(unix, not(miri)))]
 mod unix {
     use std::{ffi::CStr, ops::Deref};
@@ -154,6 +205,7 @@ mod unix {
     }
 
     impl<'a> CStrFastPathBufGuard<'a> {
+        #[must_use]
         pub fn new(buf: &mut FastPathBuf) -> CStrFastPathBufGuard {
             let FastPathBuf {
                 ref mut inner,

src/utils/file_names.rs

@@ -14,7 +14,7 @@ struct FileNameCache;
 /// since this cache is so small, we construct it at compile time and ship it
 /// with the binary.
 impl FileNameCache {
-    fn with_file_name<T, F: FnOnce(&str) -> T>(i: u16, f: F) -> T {
+    unsafe fn with_file_name<T, F: FnOnce(&str) -> T>(i: u16, f: F) -> T {
         static CACHE: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/file_name_cache.bin"));
 
         debug_assert!(i < Self::max_cache_size());
@@ -43,7 +43,7 @@ impl FileNameCache {
 #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(f)))]
 pub fn with_file_name<T>(i: u64, f: impl FnOnce(&str) -> T) -> T {
     if i < FileNameCache::max_cache_size().into() {
-        FileNameCache::with_file_name(i.try_into().unwrap(), f)
+        unsafe { FileNameCache::with_file_name(i.try_into().unwrap(), f) }
     } else {
         f(itoa::Buffer::new().format(i))
     }
@@ -53,7 +53,10 @@ pub fn with_file_name<T>(i: u64, f: impl FnOnce(&str) -> T) -> T {
 pub fn with_dir_name<T>(i: usize, f: impl FnOnce(&str) -> T) -> T {
     const SUFFIX: &str = ".dir";
     with_file_name(i.try_into().unwrap(), |s| {
+        #[allow(clippy::assertions_on_constants)]
+        const { assert!(usize::BITS <= 128, "Unsupported usize width.") }
         let mut buf = [MaybeUninit::<u8>::uninit(); 39 + SUFFIX.len()]; // 39 to support u128
+
         unsafe {
             let buf_ptr = buf.as_mut_ptr().cast::<u8>();
             ptr::copy_nonoverlapping(s.as_ptr(), buf_ptr, s.len());