Harden Large allocation layout invariants

src/internal.rs

@@ -74,14 +74,13 @@ pub(crate) struct Large<T>(pub(crate) NonNull<Allocated<T>>);
 
 impl<T> Large<T> {
     pub(crate) fn new_layout(cap: usize) -> Layout {
-        let additional_size = size_of::<T>().checked_mul(cap).expect("new capacity is too large");
-        let size = size_of::<Allocated<T>>()
-            .checked_add(additional_size)
-            .expect("new capacity is too large");
-        let align = align_of::<Allocated<T>>();
-        // SAFETY: size of Allocated<T> must be a multiple of align of Allocated<T>,
-        // which must be a multiple of align of T due to the `data` field.
-        unsafe { Layout::from_size_align_unchecked(size, align) }
+        let header = Layout::new::<Allocated<T>>();
+        let data = Layout::array::<T>(cap).expect("new capacity is too large");
+        let (layout, offset) = header.extend(data).expect("new capacity is too large");
+        // `Layout::new::<Allocated<T>>()` already includes padding for `T`,
+        // so the trailing data starts immediately after the header.
+        debug_assert_eq!(offset, size_of::<Allocated<T>>());
+        layout.pad_to_align()
     }
 
     pub(crate) fn as_allocated(&self) -> (&Allocated<T>, *const T) {
@@ -262,6 +261,9 @@ impl<T> Allocated<T> {
     ///
     /// The data behind the header are allowed to be uninitialized.
     pub(crate) unsafe fn data_start(this: NonNull<Self>) -> *mut T {
-        unsafe { (&raw mut (*this.as_ptr()).data_start).cast() }
+        let base = this.as_ptr().cast::<u8>();
+        // SAFETY: base points to a valid allocation whose data starts immediately after the header.
+        // `new_layout` uses the same header size (see the debug assertion there).
+        unsafe { base.add(size_of::<Self>()).cast::<T>() }
     }
 }

src/tests.rs

@@ -1,5 +1,6 @@
 use alloc::string::{String, ToString};
 use alloc::vec::Vec;
+use alloc::alloc::dealloc;
 use core::cell::Cell;
 use core::panic::AssertUnwindSafe;
 use core::{mem, ops};
@@ -14,6 +15,49 @@ fn assert_size() {
     assert_eq!(size_of::<WordVec<i8, 7>>(), 8);
 }
 
+#[test]
+fn test_large_layout_padding_and_data_start() {
+    use crate::internal::{Allocated, Large};
+
+    #[allow(dead_code)]
+    #[repr(align(32))]
+    struct Align32(u8);
+
+    fn assert_layout<T>() {
+        use core::alloc::Layout;
+        use core::ptr::NonNull;
+
+        let layout = Large::<T>::new_layout(1);
+        assert_eq!(layout.size() % layout.align(), 0);
+
+        let large = Large::<T>::new_empty(1);
+        let layout = large.current_layout();
+        assert_eq!(layout, Large::<T>::new_layout(1));
+        struct DeallocGuard<T> {
+            ptr: NonNull<Allocated<T>>,
+            layout: Layout,
+        }
+        impl<T> Drop for DeallocGuard<T> {
+            fn drop(&mut self) {
+                unsafe {
+                    dealloc(self.ptr.as_ptr().cast(), self.layout);
+                }
+            }
+        }
+        let _guard = DeallocGuard { ptr: large.0, layout };
+
+        let (allocated, data_start) = large.as_allocated();
+        let base = allocated as *const Allocated<T> as *const u8;
+        let offset = unsafe { data_start.cast::<u8>().offset_from(base) as usize };
+        assert_eq!(offset, size_of::<Allocated<T>>());
+    }
+
+    assert_layout::<u8>();
+    assert_layout::<u64>();
+    assert_layout::<u128>();
+    assert_layout::<Align32>();
+}
+
 struct AssertDrop<'a> {
     string:  String,
     counter: &'a Cell<usize>,