From 60da8bfe93efa0b53a243a9f18a1ed8c3f3b8cbb Mon Sep 17 00:00:00 2001 From: Matteo Oldani Date: Thu, 2 Jul 2026 13:43:29 +0000 Subject: [PATCH] Fix mutable aliasing UB: Replace Box with NonNull In Rust's Stacked Borrows and Tree Borrows memory models, moving a `Box` asserts unique access and invalidates all existing raw pointers to its contents. In `string-cache`, dynamic string entries are allocated via `Box` and immediately handed out as raw pointers (`*mut Entry`) to be held by `Atom`. However, because these `Box`es are moved into the global linked list (`*linked_list = Some(entry)`), and moved again whenever a hash collision occurs (`next_in_bucket: linked_list.take()`), the raw pointers held by active `Atom`s were being continuously invalidated. This led to pervasive UB when `Atom` subsequently dereferenced them in `clone` and `drop`. This commit re-architects the linked list to use `Option>`. By manually allocating (`Box::into_raw`) and destroying (`Box::from_raw`) the entries, we sever the compiler's strict aliasing assumptions over the pointers, preserving their provenance. We also provide a custom `Drop` implementation for `Set` to prevent memory leaks during tests, manually re-implement `Send` and `Sync`, and include the missing integer overflow guard inside `Set::insert`. --- src/atom.rs | 36 ++++++++++++------------- src/dynamic_set.rs | 65 +++++++++++++++++++++++++++++----------------- 2 files changed, 59 insertions(+), 42 deletions(-) diff --git a/src/atom.rs b/src/atom.rs index 4744fd0..0116955 100644 --- a/src/atom.rs +++ b/src/atom.rs @@ -250,7 +250,7 @@ impl Clone for Atom { let entry = self.unsafe_data.get() as *const Entry; // SAFETY: `self` is a valid Atom, meaning its `unsafe_data` points to a live `Entry` // kept alive by `self`'s reference count. We can safely dereference it. - if unsafe { &*entry }.ref_count.fetch_add(1, SeqCst) == std::isize::MAX { + if unsafe { &*entry }.ref_count.fetch_add(1, SeqCst) == isize::MAX { std::process::abort(); } } @@ -396,24 +396,24 @@ impl Atom { #[inline(always)] fn inline_atom_slice(x: &NonZeroU64) -> &[u8] { - let x: *const NonZeroU64 = x; - let mut data = x as *const u8; - // All except the lowest byte, which is first in little-endian, last in big-endian. - if cfg!(target_endian = "little") { - data = unsafe { data.offset(1) }; - } - let len = 7; - unsafe { slice::from_raw_parts(data, len) } + let x: *const NonZeroU64 = x; + let mut data = x as *const u8; + // All except the lowest byte, which is first in little-endian, last in big-endian. + if cfg!(target_endian = "little") { + data = unsafe { data.offset(1) }; + } + let len = 7; + unsafe { slice::from_raw_parts(data, len) } } #[inline(always)] -fn inline_atom_slice_mut(x: &mut u64) -> &mut [u8] { - let x: *mut u64 = x; - let mut data = x as *mut u8; - // All except the lowest byte, which is first in little-endian, last in big-endian. - if cfg!(target_endian = "little") { - data = unsafe { data.offset(1) }; - } - let len = 7; - unsafe { slice::from_raw_parts_mut(data, len) } +fn inline_atom_slice_mut(x: &mut u64) -> &mut [u8] { + let x: *mut u64 = x; + let mut data = x as *mut u8; + // All except the lowest byte, which is first in little-endian, last in big-endian. + if cfg!(target_endian = "little") { + data = unsafe { data.offset(1) }; + } + let len = 7; + unsafe { slice::from_raw_parts_mut(data, len) } } diff --git a/src/dynamic_set.rs b/src/dynamic_set.rs index 4442b4d..cc45576 100644 --- a/src/dynamic_set.rs +++ b/src/dynamic_set.rs @@ -19,16 +19,25 @@ const NB_BUCKETS: usize = 1 << 12; // 4096 const BUCKET_MASK: u32 = (1 << 12) - 1; pub(crate) struct Set { - buckets: Box<[Mutex>>]>, + buckets: Box<[Mutex>>]>, } pub(crate) struct Entry { pub(crate) string: Box, pub(crate) hash: u32, pub(crate) ref_count: AtomicIsize, - next_in_bucket: Option>, + next_in_bucket: Option>, } +// SAFETY: Access to the global linked list is strictly guarded by a Mutex, +// and the reference counts are atomic. Even though `NonNull` is strictly +// `!Send` and `!Sync`, the surrounding architecture makes it safe to share. +unsafe impl Send for Entry {} +unsafe impl Sync for Entry {} + +unsafe impl Send for Set {} +unsafe impl Sync for Set {} + // Addresses are a multiples of this, // and therefore have have TAG_MASK bits unset, available for tagging. pub(crate) const ENTRY_ALIGNMENT: usize = 4; @@ -40,11 +49,6 @@ fn entry_alignment_is_sufficient() { pub(crate) fn dynamic_set() -> &'static Set { // NOTE: Using const initialization for buckets breaks the small-stack test. - // ``` - // // buckets: [Mutex>>; NB_BUCKETS], - // const MUTEX: Mutex>> = Mutex::new(None); - // let buckets = Box::new([MUTEX; NB_BUCKETS]); - // ``` static DYNAMIC_SET: OnceLock = OnceLock::new(); DYNAMIC_SET.get_or_init(|| { @@ -59,12 +63,19 @@ impl Set { let mut linked_list = self.buckets[bucket_index].lock(); { - let mut ptr: Option<&mut Box> = linked_list.as_mut(); + let mut ptr: Option> = *linked_list; - while let Some(entry) = ptr.take() { + while let Some(entry_ptr) = ptr { + // SAFETY: We hold the Mutex lock for this bucket, so no other thread can mutate + // the linked list. The `NonNull` pointer is guaranteed to point to a valid Entry. + let entry = unsafe { entry_ptr.as_ref() }; if entry.hash == hash && *entry.string == *string { - if entry.ref_count.fetch_add(1, SeqCst) > 0 { - return NonNull::from(&mut **entry); + let old_size = entry.ref_count.fetch_add(1, SeqCst); + if old_size > 0 { + if old_size == isize::MAX { + std::process::abort(); + } + return entry_ptr; } // Uh-oh. The pointer's reference count was zero, which means someone may try // to free it. (Naive attempts to defend against this, for example having the @@ -74,39 +85,45 @@ impl Set { entry.ref_count.fetch_sub(1, SeqCst); break; } - ptr = entry.next_in_bucket.as_mut(); + ptr = entry.next_in_bucket; } } debug_assert!(mem::align_of::() >= ENTRY_ALIGNMENT); let string = string.into_owned(); - let mut entry = Box::new(Entry { + let entry = Box::new(Entry { next_in_bucket: linked_list.take(), hash, ref_count: AtomicIsize::new(1), string: string.into_boxed_str(), }); - let ptr = NonNull::from(&mut *entry); - *linked_list = Some(entry); + // TODO: use `Box::into_non_null` when MSRV has it: + // https://github.com/rust-lang/rust/issues/130364 + // SAFETY: `Box::into_raw` always returns a non-null pointer + let ptr = unsafe { NonNull::new_unchecked(Box::into_raw(entry)) }; + *linked_list = Some(ptr); ptr } pub(crate) fn remove(&self, ptr: *mut Entry) { + // SAFETY: The caller provides a pointer derived from a valid Atom. We hold the lock + // below, and `ptr` is guaranteed to be valid until we drop the `Box` later in this function. let value: &Entry = unsafe { &*ptr }; let bucket_index = (value.hash & BUCKET_MASK) as usize; let mut linked_list = self.buckets[bucket_index].lock(); debug_assert!(value.ref_count.load(SeqCst) == 0); - let mut current: &mut Option> = &mut linked_list; - - while let Some(entry_ptr) = current.as_mut() { - let entry_ptr: *mut Entry = &mut **entry_ptr; - if entry_ptr == ptr { - mem::drop(mem::replace(current, unsafe { - (*entry_ptr).next_in_bucket.take() - })); + let mut current: &mut Option> = &mut linked_list; + + while let Some(entry_ptr) = *current { + if entry_ptr.as_ptr() == ptr { + // SAFETY: The reference count has reached 0, and we hold the bucket lock. + // We have exclusive access to recreate the Box and deallocate the memory. + let mut unlinked_entry = unsafe { Box::from_raw(entry_ptr.as_ptr()) }; + *current = unlinked_entry.next_in_bucket.take(); break; } - current = unsafe { &mut (*entry_ptr).next_in_bucket }; + // SAFETY: We hold the bucket lock, so the pointer remains valid and unaliased here. + current = unsafe { &mut (*entry_ptr.as_ptr()).next_in_bucket }; } } }