From aa439cff19c0e6060cfe1c9df404ff364016844d Mon Sep 17 00:00:00 2001
From: Andreas Molzer <andreas.molzer@gmx.de>
Date: Sun, 10 Nov 2024 23:42:38 +0100
Subject: [PATCH] Verify alignment for AlignedVec

Turning the pointer of an aligned chunk into a reference to the overlay
element type requires us to verify the element type's alignment
invariant on the address. This check need not happen dynamically if we
can provide a stronger static assertion.

The use of `AlignedVec` is over-aligning the storage with regards to the
elements. This guarantees that all aligned block addresses are at least
as aligned as required for elements. With this commit we validate that
only vectors with appropriate element types can be instantiated. This
verifies the constructor instead of the type parameters due to
simplicity. Turning the violation into a compile error may be beneficial
during development but should not be strictly necessary. Passing through
a constructor is a dominator code path of the creation of a reference
from its storage (at least right now).
---
 src/align.rs | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/align.rs b/src/align.rs
index 01bc7f68e..265d11a5f 100644
--- a/src/align.rs
+++ b/src/align.rs
@@ -146,8 +146,14 @@ impl<T: Copy, C: AlignedByteChunk> AlignedVec<T, C> {
         assert!(mem::size_of::<C>() == mem::align_of::<C>());
     }
 
+    const fn check_inner_type_is_aligned() {
+        assert!(mem::align_of::<T>() <= mem::align_of::<C>());
+    }
+
     pub const fn new() -> Self {
         Self::check_byte_chunk_type_is_aligned();
+        Self::check_inner_type_is_aligned();
+
         Self {
             inner: Vec::new(),
             len: 0,
@@ -172,6 +178,9 @@ impl<T: Copy, C: AlignedByteChunk> AlignedVec<T, C> {
     pub fn as_slice(&self) -> &[T] {
         // SAFETY: The first `len` elements have been
         // initialized to `T`s in `Self::resize_with`.
+        // SAFETY: The pointer is sufficiently aligned,
+        // as the chunks are always over-aligned with
+        // respect to `T`.
         unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
     }
 
@@ -179,6 +188,9 @@ impl<T: Copy, C: AlignedByteChunk> AlignedVec<T, C> {
     pub fn as_mut_slice(&mut self) -> &mut [T] {
         // SAFETY: The first `len` elements have been
         // initialized to `T`s in `Self::resize_with`.
+        // SAFETY: The pointer is sufficiently aligned,
+        // as the chunks are always over-aligned with
+        // respect to `T`.
         unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
     }
 
@@ -208,6 +220,9 @@ impl<T: Copy, C: AlignedByteChunk> AlignedVec<T, C> {
         for offset in old_len..new_len {
             // SAFETY: We've allocated enough space to write
             // up to `new_len` elements into the buffer.
+            // SAFETY: The pointer is sufficiently aligned,
+            // as the chunks are always over-aligned with
+            // respect to `T`.
             unsafe { self.as_mut_ptr().add(offset).write(value) };
         }
 
@@ -288,3 +303,13 @@ fn align_vec_fails() {
     // work.
     assert_eq!(v.as_slice()[isize::MAX as usize], 0);
 }
+
+#[test]
+#[should_panic]
+fn under_aligned_storage_fails() {
+    let mut v = AlignedVec::<u128, Align1<[u8; 1]>>::new();
+
+    // Would be UB: unaligned write of type u128 to pointer aligned to 1 (or 8 in most allocators,
+    // but statically we only ensure an alignment of 1).
+    v.resize(1, 0u128);
+}