[v7] Generic `Allocator` support for Rust

[PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Danilo Krummrich 2 months ago

Add a kernel specific `Allocator` trait, that in contrast to the one in
Rust's core library doesn't require unstable features and supports GFP
flags.

Subsequent patches add the following trait implementors: `Kmalloc`,
`Vmalloc` and `KVmalloc`.

Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Danilo Krummrich <dakr@kernel.org>
---
 rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 112 insertions(+)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 1966bd407017..6c21bd2edad9 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -11,6 +11,7 @@
 /// Indicates an allocation error.
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct AllocError;
+use core::{alloc::Layout, ptr::NonNull};
 
 /// Flags to be used when allocating memory.
 ///
@@ -86,3 +87,114 @@ pub mod flags {
     /// small allocations.
     pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
 }
+
+/// The kernel's [`Allocator`] trait.
+///
+/// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffers described
+/// via [`Layout`].
+///
+/// [`Allocator`] is designed to be implemented as a ZST; [`Allocator`] functions do not operate on
+/// an object instance.
+///
+/// In order to be able to support `#[derive(SmartPointer)]` later on, we need to avoid a design
+/// that requires an `Allocator` to be instantiated, hence its functions must not contain any kind
+/// of `self` parameter.
+///
+/// # Safety
+///
+/// - A memory allocation returned from an allocator must remain valid until it is explicitly freed.
+///
+/// - Any pointer to a valid memory allocation must be valid to be passed to any other [`Allocator`]
+///   function of the same type.
+///
+/// - Implementers must ensure that all trait functions abide by the guarantees documented in the
+///   `# Guarantees` sections.
+//
+// Note that `Allocator::{realloc,free}` don't have an `old_layout` argument (like stdlib's
+// corresponding `Allocator` trait functions have), since the implemented (kernel) allocators
+// neither need nor honor such an argument. Thus, it would be misleading to make this API require it
+// anyways.
+//
+// More generally, this trait isn't intended for implementers to encode a lot of semantics, but
+// rather provide a thin generalization layer for the kernel's allocators.
+//
+// Depending on future requirements, the requirements for this trait may change as well and
+// implementing allocators that need to encode more semantics may become desirable.
+pub unsafe trait Allocator {
+    /// Allocate memory based on `layout` and `flags`.
+    ///
+    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
+    /// constraints (i.e. minimum size and alignment as specified by `layout`).
+    ///
+    /// This function is equivalent to `realloc` when called with `None`.
+    ///
+    /// # Guarantees
+    ///
+    /// When the return value is `Ok(ptr)`, then `ptr` is
+    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
+    ///   [`Allocator::free`] or [`Allocator::realloc`],
+    /// - aligned to `layout.align()`,
+    ///
+    /// Additionally, `Flags` are honored as documented in
+    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
+    fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
+        // SAFETY: Passing `None` to `realloc` is valid by it's safety requirements and asks for a
+        // new memory allocation.
+        unsafe { Self::realloc(None, layout, flags) }
+    }
+
+    /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
+    ///
+    /// If the requested size is zero, `realloc` behaves equivalent to `free`.
+    ///
+    /// If the requested size is larger than the size of the existing allocation, a successful call
+    /// to `realloc` guarantees that the new or grown buffer has at least `Layout::size` bytes, but
+    /// may also be larger.
+    ///
+    /// If the requested size is smaller than the size of the existing allocation, `realloc` may or
+    /// may not shrink the buffer; this is implementation specific to the allocator.
+    ///
+    /// On allocation failure, the existing buffer, if any, remains valid.
+    ///
+    /// The buffer is represented as `NonNull<[u8]>`.
+    ///
+    /// # Safety
+    ///
+    /// If `ptr == Some(p)`, then `p` must point to an existing and valid memory allocation created
+    /// by this allocator. The alignment encoded in `layout` must be smaller than or equal to the
+    /// alignment requested in the previous `alloc` or `realloc` call of the same allocation.
+    ///
+    /// Additionally, `ptr` is allowed to be `None`; in this case a new memory allocation is
+    /// created.
+    ///
+    /// # Guarantees
+    ///
+    /// This function has the same guarantees as [`Allocator::alloc`]. When `ptr == Some(p)`, then
+    /// it additionally guarantees that:
+    /// - the contents of the memory pointed to by `p` are preserved up to the lesser of the new
+    ///   and old size,
+    ///   and old size, i.e.
+    ///   `ret_ptr[0..min(layout.size(), old_size)] == p[0..min(layout.size(), old_size)]`, where
+    ///   `old_size` is the size of the allocation that `p` points at.
+    /// - when the return value is `Err(AllocError)`, then `p` is still valid.
+    unsafe fn realloc(
+        ptr: Option<NonNull<u8>>,
+        layout: Layout,
+        flags: Flags,
+    ) -> Result<NonNull<[u8]>, AllocError>;
+
+    /// Free an existing memory allocation.
+    ///
+    /// # Safety
+    ///
+    /// `ptr` must point to an existing and valid memory allocation created by this `Allocator` and
+    /// must not be a dangling pointer.
+    ///
+    /// The memory allocation at `ptr` must never again be read from or written to.
+    unsafe fn free(ptr: NonNull<u8>) {
+        // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
+        // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
+        // smaller than or equal to the alignment previously used with this allocation.
+        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), Flags(0)) };
+    }
+}
-- 
2.46.0

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Gary Guo 1 month, 4 weeks ago

On Thu, 12 Sep 2024 00:52:37 +0200
Danilo Krummrich <dakr@kernel.org> wrote:

> Add a kernel specific `Allocator` trait, that in contrast to the one in
> Rust's core library doesn't require unstable features and supports GFP
> flags.
> 
> Subsequent patches add the following trait implementors: `Kmalloc`,
> `Vmalloc` and `KVmalloc`.

Hi Danilo,

I think the current design is unsound regarding ZST.

Let's say that `Allocator::alloc` gets called with a ZST type with
alignment of 4096. Your implementation will call into `krelloc` with
new_size of 0, which gets turned into of `kfree` of null pointer, which
is no-op. Everything is fine so far. Krealloc returns `ZERO_SIZE_PTR`,
and then implementation of `<Kmalloc as Allocator>::realloc` throws it
away and returns `NonNull::dangling`.

Since `NonNull::dangling` is called with T=u8, this means the pointer
returns is 1, and it's invalid for ZSTs with larger alignments.

And this is unfixable even if the realloc implementation is changed.
Let's say the realloc now returns a dangling pointer that is suitable
aligned. Now let's see what happens when the `Allocator::free` is
called. `kfree` would be trying to free a Rust-side ZST pointer, but it
has no way to know that it's ZST!

I can see 3 ways of fixing this:
1. Reject ZSTs that have larger alignment than 16 and fix the realloc
implementation to return suitable aligned ZST pointer. I don't
particularly like the idea of allocating ZST can fail though.
2. Say ZST must be handled by the caller, and make alloc function
unsafe. This means that we essentially revert to the `GlobalAlloc`
design of Rust, and all callers have to check for ZST.
3. Accept the `old_layout` and use it to check whether the allocation
is ZST allocation.

My personal preference is 3.

Best,
Gary

> 
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> ---
>  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 112 insertions(+)
> 
> diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> index 1966bd407017..6c21bd2edad9 100644
> --- a/rust/kernel/alloc.rs
> +++ b/rust/kernel/alloc.rs
> @@ -11,6 +11,7 @@
>  /// Indicates an allocation error.
>  #[derive(Copy, Clone, PartialEq, Eq, Debug)]
>  pub struct AllocError;
> +use core::{alloc::Layout, ptr::NonNull};
>  
>  /// Flags to be used when allocating memory.
>  ///
> @@ -86,3 +87,114 @@ pub mod flags {
>      /// small allocations.
>      pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
>  }
> +
> +/// The kernel's [`Allocator`] trait.
> +///
> +/// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffers described
> +/// via [`Layout`].
> +///
> +/// [`Allocator`] is designed to be implemented as a ZST; [`Allocator`] functions do not operate on
> +/// an object instance.

I think whether the Allocator is ZST or not doesn't matter anymore
since we say that functions do not operate on an instance.

> +///
> +/// In order to be able to support `#[derive(SmartPointer)]` later on, we need to avoid a design
> +/// that requires an `Allocator` to be instantiated, hence its functions must not contain any kind
> +/// of `self` parameter.


nit: this is the reason for `Allocator` to not have instances, so
should be merged together with the preceding sentence into one
paragraph.

> +///
> +/// # Safety
> +///
> +/// - A memory allocation returned from an allocator must remain valid until it is explicitly freed.
> +///
> +/// - Any pointer to a valid memory allocation must be valid to be passed to any other [`Allocator`]
> +///   function of the same type.
> +///
> +/// - Implementers must ensure that all trait functions abide by the guarantees documented in the
> +///   `# Guarantees` sections.
> +//
> +// Note that `Allocator::{realloc,free}` don't have an `old_layout` argument (like stdlib's
> +// corresponding `Allocator` trait functions have), since the implemented (kernel) allocators
> +// neither need nor honor such an argument. Thus, it would be misleading to make this API require it
> +// anyways.

I would drop the "honor" part, and drop the sentence saying it's
misleading to require it. The documentation should say why we don't
have the `old_layout` argument (because we don't need it for now), and
shouldn't be trying to dissuade it from being added if it ends up being
useful in the future.

> +//
> +// More generally, this trait isn't intended for implementers to encode a lot of semantics, but
> +// rather provide a thin generalization layer for the kernel's allocators.
> +//
> +// Depending on future requirements, the requirements for this trait may change as well and
> +// implementing allocators that need to encode more semantics may become desirable.

Not sure what's the purpose of these two paragraphs. They sound
contradictory to each other.

> +pub unsafe trait Allocator {
> +    /// Allocate memory based on `layout` and `flags`.
> +    ///
> +    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
> +    /// constraints (i.e. minimum size and alignment as specified by `layout`).
> +    ///
> +    /// This function is equivalent to `realloc` when called with `None`.
> +    ///
> +    /// # Guarantees
> +    ///
> +    /// When the return value is `Ok(ptr)`, then `ptr` is
> +    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
> +    ///   [`Allocator::free`] or [`Allocator::realloc`],
> +    /// - aligned to `layout.align()`,
> +    ///
> +    /// Additionally, `Flags` are honored as documented in
> +    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
> +    fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
> +        // SAFETY: Passing `None` to `realloc` is valid by it's safety requirements and asks for a
> +        // new memory allocation.
> +        unsafe { Self::realloc(None, layout, flags) }
> +    }
> +
> +    /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
> +    ///
> +    /// If the requested size is zero, `realloc` behaves equivalent to `free`.
> +    ///
> +    /// If the requested size is larger than the size of the existing allocation, a successful call
> +    /// to `realloc` guarantees that the new or grown buffer has at least `Layout::size` bytes, but
> +    /// may also be larger.
> +    ///
> +    /// If the requested size is smaller than the size of the existing allocation, `realloc` may or
> +    /// may not shrink the buffer; this is implementation specific to the allocator.
> +    ///
> +    /// On allocation failure, the existing buffer, if any, remains valid.
> +    ///
> +    /// The buffer is represented as `NonNull<[u8]>`.
> +    ///
> +    /// # Safety
> +    ///
> +    /// If `ptr == Some(p)`, then `p` must point to an existing and valid memory allocation created

nit: maybe

	If `ptr` is `Some(p)`?

because `ptr` carries the provenance not only addresses.

> +    /// by this allocator. The alignment encoded in `layout` must be smaller than or equal to the
> +    /// alignment requested in the previous `alloc` or `realloc` call of the same allocation.
> +    ///
> +    /// Additionally, `ptr` is allowed to be `None`; in this case a new memory allocation is
> +    /// created.
> +    ///
> +    /// # Guarantees
> +    ///
> +    /// This function has the same guarantees as [`Allocator::alloc`]. When `ptr == Some(p)`, then
> +    /// it additionally guarantees that:
> +    /// - the contents of the memory pointed to by `p` are preserved up to the lesser of the new
> +    ///   and old size,
> +    ///   and old size, i.e.
> +    ///   `ret_ptr[0..min(layout.size(), old_size)] == p[0..min(layout.size(), old_size)]`, where
> +    ///   `old_size` is the size of the allocation that `p` points at.
> +    /// - when the return value is `Err(AllocError)`, then `p` is still valid.
> +    unsafe fn realloc(
> +        ptr: Option<NonNull<u8>>,
> +        layout: Layout,
> +        flags: Flags,
> +    ) -> Result<NonNull<[u8]>, AllocError>;
> +
> +    /// Free an existing memory allocation.
> +    ///
> +    /// # Safety
> +    ///
> +    /// `ptr` must point to an existing and valid memory allocation created by this `Allocator` and
> +    /// must not be a dangling pointer.
> +    ///
> +    /// The memory allocation at `ptr` must never again be read from or written to.
> +    unsafe fn free(ptr: NonNull<u8>) {
> +        // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
> +        // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
> +        // smaller than or equal to the alignment previously used with this allocation.
> +        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), Flags(0)) };
> +    }
> +}

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Danilo Krummrich 1 month, 4 weeks ago

Hi Gary,

thanks for taking a look.

On Sun, Sep 15, 2024 at 04:28:13PM +0100, Gary Guo wrote:
> On Thu, 12 Sep 2024 00:52:37 +0200
> Danilo Krummrich <dakr@kernel.org> wrote:
> 
> > Add a kernel specific `Allocator` trait, that in contrast to the one in
> > Rust's core library doesn't require unstable features and supports GFP
> > flags.
> > 
> > Subsequent patches add the following trait implementors: `Kmalloc`,
> > `Vmalloc` and `KVmalloc`.
> 
> Hi Danilo,
> 
> I think the current design is unsound regarding ZST.
> 
> Let's say that `Allocator::alloc` gets called with a ZST type with
> alignment of 4096. Your implementation will call into `krelloc` with
> new_size of 0, which gets turned into of `kfree` of null pointer, which
> is no-op. Everything is fine so far. Krealloc returns `ZERO_SIZE_PTR`,
> and then implementation of `<Kmalloc as Allocator>::realloc` throws it
> away and returns `NonNull::dangling`.
> 
> Since `NonNull::dangling` is called with T=u8, this means the pointer
> returns is 1, and it's invalid for ZSTs with larger alignments.

Right, this interface is not meant to handle "allocations" for ZSTs.

But you're right, since `alloc` is a safe function, we should return a properly
aligned pointer.

> 
> And this is unfixable even if the realloc implementation is changed.
> Let's say the realloc now returns a dangling pointer that is suitable
> aligned. Now let's see what happens when the `Allocator::free` is
> called. `kfree` would be trying to free a Rust-side ZST pointer, but it
> has no way to know that it's ZST!

Right, that's why it's not valid to call `free` with dangling pointers.

From the safety comment of `free`:

"`ptr` must point to an existing and valid memory allocation created by this
`Allocator` and must not be a dangling pointer."

We still need the same in `realloc` though.

> 
> I can see 3 ways of fixing this:
> 1. Reject ZSTs that have larger alignment than 16 and fix the realloc
> implementation to return suitable aligned ZST pointer. I don't
> particularly like the idea of allocating ZST can fail though.
> 2. Say ZST must be handled by the caller, and make alloc function
> unsafe. This means that we essentially revert to the `GlobalAlloc`
> design of Rust, and all callers have to check for ZST.
> 3. Accept the `old_layout` and use it to check whether the allocation
> is ZST allocation.
> 
> My personal preference is 3.

There is also 4.

Let `alloc` and `realloc` return a properly aligned dangling pointer for
`size == 0` and don't accept dangling pointers in `realloc` and `free`.

And 5.

Reject the combination of `None` and `size == 0` entirely, as earlier proposed
by Benno.

I'm fine with both, 4. and 5. with a slight preference for 4.

I'd also go along with 1., as a mix of 4. and 5.

I really don't like making `alloc` unsafe, and I really don't want to have
`old_layout` in `free`. Please let's not discuss this again. :-)

> 
> Best,
> Gary
> 
> > 
> > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> > ---
> >  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
> >  1 file changed, 112 insertions(+)
> > 
> > diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> > index 1966bd407017..6c21bd2edad9 100644
> > --- a/rust/kernel/alloc.rs
> > +++ b/rust/kernel/alloc.rs
> > @@ -11,6 +11,7 @@
> >  /// Indicates an allocation error.
> >  #[derive(Copy, Clone, PartialEq, Eq, Debug)]
> >  pub struct AllocError;
> > +use core::{alloc::Layout, ptr::NonNull};
> >  
> >  /// Flags to be used when allocating memory.
> >  ///
> > @@ -86,3 +87,114 @@ pub mod flags {
> >      /// small allocations.
> >      pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
> >  }
> > +
> > +/// The kernel's [`Allocator`] trait.
> > +///
> > +/// An implementation of [`Allocator`] can allocate, re-allocate and free memory buffers described
> > +/// via [`Layout`].
> > +///
> > +/// [`Allocator`] is designed to be implemented as a ZST; [`Allocator`] functions do not operate on
> > +/// an object instance.
> 
> I think whether the Allocator is ZST or not doesn't matter anymore
> since we say that functions do not operate on an instance.

IMO, It's still valid to say that it's designed to be implemented as ZST,
especially since it's never instantiated.

> 
> > +///
> > +/// In order to be able to support `#[derive(SmartPointer)]` later on, we need to avoid a design
> > +/// that requires an `Allocator` to be instantiated, hence its functions must not contain any kind
> > +/// of `self` parameter.
> 
> 
> nit: this is the reason for `Allocator` to not have instances, so
> should be merged together with the preceding sentence into one
> paragraph.
> 
> > +///
> > +/// # Safety
> > +///
> > +/// - A memory allocation returned from an allocator must remain valid until it is explicitly freed.
> > +///
> > +/// - Any pointer to a valid memory allocation must be valid to be passed to any other [`Allocator`]
> > +///   function of the same type.
> > +///
> > +/// - Implementers must ensure that all trait functions abide by the guarantees documented in the
> > +///   `# Guarantees` sections.
> > +//
> > +// Note that `Allocator::{realloc,free}` don't have an `old_layout` argument (like stdlib's
> > +// corresponding `Allocator` trait functions have), since the implemented (kernel) allocators
> > +// neither need nor honor such an argument. Thus, it would be misleading to make this API require it
> > +// anyways.
> 
> I would drop the "honor" part, and drop the sentence saying it's
> misleading to require it. The documentation should say why we don't
> have the `old_layout` argument (because we don't need it for now), and
> shouldn't be trying to dissuade it from being added if it ends up being
> useful in the future.

But the honor part is the crucial one for me. Just because it's not needed it
could still be optional (and honored). But that's not the case. And because of
that it would indeed be misleading to have the argument.

That's really just my reason, I'm not trying to dissuade anyone from adding it
if it's actually needed.

If there's a good reason to add it, and it's honored by (other) implementations,
it wouldn't be generally misleading anymore.

> 
> > +//
> > +// More generally, this trait isn't intended for implementers to encode a lot of semantics, but
> > +// rather provide a thin generalization layer for the kernel's allocators.
> > +//
> > +// Depending on future requirements, the requirements for this trait may change as well and
> > +// implementing allocators that need to encode more semantics may become desirable.
> 
> Not sure what's the purpose of these two paragraphs. They sound
> contradictory to each other.

I tried to articulate what it is currently intended for, but, depending on
future requirements, this may change.

I was asked by Benno to somehow add that we're open for changes...

> 
> > +pub unsafe trait Allocator {
> > +    /// Allocate memory based on `layout` and `flags`.
> > +    ///
> > +    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
> > +    /// constraints (i.e. minimum size and alignment as specified by `layout`).
> > +    ///
> > +    /// This function is equivalent to `realloc` when called with `None`.
> > +    ///
> > +    /// # Guarantees
> > +    ///
> > +    /// When the return value is `Ok(ptr)`, then `ptr` is
> > +    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
> > +    ///   [`Allocator::free`] or [`Allocator::realloc`],
> > +    /// - aligned to `layout.align()`,
> > +    ///
> > +    /// Additionally, `Flags` are honored as documented in
> > +    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
> > +    fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
> > +        // SAFETY: Passing `None` to `realloc` is valid by it's safety requirements and asks for a
> > +        // new memory allocation.
> > +        unsafe { Self::realloc(None, layout, flags) }
> > +    }
> > +
> > +    /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
> > +    ///
> > +    /// If the requested size is zero, `realloc` behaves equivalent to `free`.
> > +    ///
> > +    /// If the requested size is larger than the size of the existing allocation, a successful call
> > +    /// to `realloc` guarantees that the new or grown buffer has at least `Layout::size` bytes, but
> > +    /// may also be larger.
> > +    ///
> > +    /// If the requested size is smaller than the size of the existing allocation, `realloc` may or
> > +    /// may not shrink the buffer; this is implementation specific to the allocator.
> > +    ///
> > +    /// On allocation failure, the existing buffer, if any, remains valid.
> > +    ///
> > +    /// The buffer is represented as `NonNull<[u8]>`.
> > +    ///
> > +    /// # Safety
> > +    ///
> > +    /// If `ptr == Some(p)`, then `p` must point to an existing and valid memory allocation created
> 
> nit: maybe
> 
> 	If `ptr` is `Some(p)`?
> 
> because `ptr` carries the provenance not only addresses.

I don't mind either versions, this one was proposed by Benno.

> 
> > +    /// by this allocator. The alignment encoded in `layout` must be smaller than or equal to the
> > +    /// alignment requested in the previous `alloc` or `realloc` call of the same allocation.
> > +    ///
> > +    /// Additionally, `ptr` is allowed to be `None`; in this case a new memory allocation is
> > +    /// created.
> > +    ///
> > +    /// # Guarantees
> > +    ///
> > +    /// This function has the same guarantees as [`Allocator::alloc`]. When `ptr == Some(p)`, then
> > +    /// it additionally guarantees that:
> > +    /// - the contents of the memory pointed to by `p` are preserved up to the lesser of the new
> > +    ///   and old size,
> > +    ///   and old size, i.e.
> > +    ///   `ret_ptr[0..min(layout.size(), old_size)] == p[0..min(layout.size(), old_size)]`, where
> > +    ///   `old_size` is the size of the allocation that `p` points at.
> > +    /// - when the return value is `Err(AllocError)`, then `p` is still valid.
> > +    unsafe fn realloc(
> > +        ptr: Option<NonNull<u8>>,
> > +        layout: Layout,
> > +        flags: Flags,
> > +    ) -> Result<NonNull<[u8]>, AllocError>;
> > +
> > +    /// Free an existing memory allocation.
> > +    ///
> > +    /// # Safety
> > +    ///
> > +    /// `ptr` must point to an existing and valid memory allocation created by this `Allocator` and
> > +    /// must not be a dangling pointer.
> > +    ///
> > +    /// The memory allocation at `ptr` must never again be read from or written to.
> > +    unsafe fn free(ptr: NonNull<u8>) {
> > +        // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
> > +        // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
> > +        // smaller than or equal to the alignment previously used with this allocation.
> > +        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), Flags(0)) };
> > +    }
> > +}
>

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Gary Guo 1 month, 4 weeks ago

On Sun, 15 Sep 2024 19:02:40 +0200
Danilo Krummrich <dakr@kernel.org> wrote:

> Hi Gary,
> 
> thanks for taking a look.
> 
> On Sun, Sep 15, 2024 at 04:28:13PM +0100, Gary Guo wrote:
> > On Thu, 12 Sep 2024 00:52:37 +0200
> > Danilo Krummrich <dakr@kernel.org> wrote:
> >   
> > > Add a kernel specific `Allocator` trait, that in contrast to the one in
> > > Rust's core library doesn't require unstable features and supports GFP
> > > flags.
> > > 
> > > Subsequent patches add the following trait implementors: `Kmalloc`,
> > > `Vmalloc` and `KVmalloc`.  
> > 
> > Hi Danilo,
> > 
> > I think the current design is unsound regarding ZST.
> > 
> > Let's say that `Allocator::alloc` gets called with a ZST type with
> > alignment of 4096. Your implementation will call into `krelloc` with
> > new_size of 0, which gets turned into of `kfree` of null pointer, which
> > is no-op. Everything is fine so far. Krealloc returns `ZERO_SIZE_PTR`,
> > and then implementation of `<Kmalloc as Allocator>::realloc` throws it
> > away and returns `NonNull::dangling`.
> > 
> > Since `NonNull::dangling` is called with T=u8, this means the pointer
> > returns is 1, and it's invalid for ZSTs with larger alignments.  
> 
> Right, this interface is not meant to handle "allocations" for ZSTs.
> 
> But you're right, since `alloc` is a safe function, we should return a properly
> aligned pointer.
> 
> > 
> > And this is unfixable even if the realloc implementation is changed.
> > Let's say the realloc now returns a dangling pointer that is suitable
> > aligned. Now let's see what happens when the `Allocator::free` is
> > called. `kfree` would be trying to free a Rust-side ZST pointer, but it
> > has no way to know that it's ZST!  
> 
> Right, that's why it's not valid to call `free` with dangling pointers.
> 
> From the safety comment of `free`:
> 
> "`ptr` must point to an existing and valid memory allocation created by this
> `Allocator` and must not be a dangling pointer."
> 
> We still need the same in `realloc` though.

I don't agree with this reading. If you allocate something with `alloc`
and it doesn't return an error then you should be able to feed it to
`free`. Whether the allocator does actual allocation when size is zero
or return a dangling pointer shouldn't matter to the caller.

The fact you `Kmalloc` returns a dangling pointer for ZST is an
implementation detail and the caller shouldn't care (and it also
couldn't check whether it's a dangling pointer). Nothing in your
`alloc` doc mention about dangling pointer return for zero-sized alloc
at all.

> 
> > 
> > I can see 3 ways of fixing this:
> > 1. Reject ZSTs that have larger alignment than 16 and fix the realloc
> > implementation to return suitable aligned ZST pointer. I don't
> > particularly like the idea of allocating ZST can fail though.
> > 2. Say ZST must be handled by the caller, and make alloc function
> > unsafe. This means that we essentially revert to the `GlobalAlloc`
> > design of Rust, and all callers have to check for ZST.
> > 3. Accept the `old_layout` and use it to check whether the allocation
> > is ZST allocation.
> > 
> > My personal preference is 3.  
> 
> There is also 4.
> 
> Let `alloc` and `realloc` return a properly aligned dangling pointer for
> `size == 0` and don't accept dangling pointers in `realloc` and `free`.

I'll consider the API design to be bad if I can't pass allocated pointer to
free. If caller needs to handle ZST specially then we might as well
just ban it completely.

> And 5.
> 
> Reject the combination of `None` and `size == 0` entirely, as earlier proposed
> by Benno.
> 
> I'm fine with both, 4. and 5. with a slight preference for 4.
> 
> I'd also go along with 1., as a mix of 4. and 5.
> 
> I really don't like making `alloc` unsafe, and I really don't want to have
> `old_layout` in `free`. Please let's not discuss this again. :-)

I don't buy it.

Your argument for having `old_layout` is so that the caller doesn't
need to care about the size. But as demonstrated the caller *does* need
to care about whether the size is zero.

Our previous discussion doesn't cover the particular case of ZST and
you said that it reason arise that we need this extra parameter, then
it could be added. It feels to me that sane behaviour when it comes
to ZST allocation is a very good reason.

> 
> > 
> > Best,
> > Gary
> >   
> > > 
> > > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > > Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> > > ---
> > >  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
> > >  1 file changed, 112 insertions(+)

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Danilo Krummrich 1 month, 4 weeks ago

On Sun, Sep 15, 2024 at 08:22:42PM +0100, Gary Guo wrote:
> On Sun, 15 Sep 2024 19:02:40 +0200
> Danilo Krummrich <dakr@kernel.org> wrote:
> > There is also 4.
> > 
> > Let `alloc` and `realloc` return a properly aligned dangling pointer for
> > `size == 0` and don't accept dangling pointers in `realloc` and `free`.
> 
> I'll consider the API design to be bad if I can't pass allocated pointer to
> free. If caller needs to handle ZST specially then we might as well
> just ban it completely.

Fine for me -- I don't see a need to support ZSTs with `Allocator`.

I think its main purpose is to give us an interface to actually allocate memory.
We probably don't need it to be a "dangling pointer generator".

> 
> > And 5.
> > 
> > Reject the combination of `None` and `size == 0` entirely, as earlier proposed
> > by Benno.
> > 
> > I'm fine with both, 4. and 5. with a slight preference for 4.
> > 
> > I'd also go along with 1., as a mix of 4. and 5.
> > 
> > I really don't like making `alloc` unsafe, and I really don't want to have
> > `old_layout` in `free`. Please let's not discuss this again. :-)
> 
> I don't buy it.
> 
> Your argument for having `old_layout` is so that the caller doesn't
> need to care about the size. But as demonstrated the caller *does* need
> to care about whether the size is zero.
> 
> Our previous discussion doesn't cover the particular case of ZST and
> you said that it reason arise that we need this extra parameter, then
> it could be added. It feels to me that sane behaviour when it comes
> to ZST allocation is a very good reason.

I don't see why we should "optimize" the API for creating dangling pointers and
be able to pass them to `free` (which does not serve any practical purpose).

I don't want to add arguments that are meaningless for the actual backing
allocators (such as Kmalloc, Vmalloc, etc.), only to be able to generate and
"free" pointers for ZSTs with arbitrary alignment.

Do we even have use cases for ZSTs with other alignments?

> 
> > 
> > > 
> > > Best,
> > > Gary
> > >   
> > > > 
> > > > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > > > Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> > > > ---
> > > >  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
> > > >  1 file changed, 112 insertions(+)
>

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Gary Guo 1 month, 4 weeks ago

On Sun, 15 Sep 2024 20:22:42 +0100
Gary Guo <gary@garyguo.net> wrote:

> On Sun, 15 Sep 2024 19:02:40 +0200
> Danilo Krummrich <dakr@kernel.org> wrote:
> 
> > Hi Gary,
> > 
> > thanks for taking a look.
> > 
> > On Sun, Sep 15, 2024 at 04:28:13PM +0100, Gary Guo wrote:  
> > > On Thu, 12 Sep 2024 00:52:37 +0200
> > > Danilo Krummrich <dakr@kernel.org> wrote:
> > >     
> > > > Add a kernel specific `Allocator` trait, that in contrast to the one in
> > > > Rust's core library doesn't require unstable features and supports GFP
> > > > flags.
> > > > 
> > > > Subsequent patches add the following trait implementors: `Kmalloc`,
> > > > `Vmalloc` and `KVmalloc`.    
> > > 
> > > Hi Danilo,
> > > 
> > > I think the current design is unsound regarding ZST.
> > > 
> > > Let's say that `Allocator::alloc` gets called with a ZST type with
> > > alignment of 4096. Your implementation will call into `krelloc` with
> > > new_size of 0, which gets turned into of `kfree` of null pointer, which
> > > is no-op. Everything is fine so far. Krealloc returns `ZERO_SIZE_PTR`,
> > > and then implementation of `<Kmalloc as Allocator>::realloc` throws it
> > > away and returns `NonNull::dangling`.
> > > 
> > > Since `NonNull::dangling` is called with T=u8, this means the pointer
> > > returns is 1, and it's invalid for ZSTs with larger alignments.    
> > 
> > Right, this interface is not meant to handle "allocations" for ZSTs.
> > 
> > But you're right, since `alloc` is a safe function, we should return a properly
> > aligned pointer.
> >   
> > > 
> > > And this is unfixable even if the realloc implementation is changed.
> > > Let's say the realloc now returns a dangling pointer that is suitable
> > > aligned. Now let's see what happens when the `Allocator::free` is
> > > called. `kfree` would be trying to free a Rust-side ZST pointer, but it
> > > has no way to know that it's ZST!    
> > 
> > Right, that's why it's not valid to call `free` with dangling pointers.
> > 
> > From the safety comment of `free`:
> > 
> > "`ptr` must point to an existing and valid memory allocation created by this
> > `Allocator` and must not be a dangling pointer."
> > 
> > We still need the same in `realloc` though.  
> 
> I don't agree with this reading. If you allocate something with `alloc`
> and it doesn't return an error then you should be able to feed it to
> `free`. Whether the allocator does actual allocation when size is zero
> or return a dangling pointer shouldn't matter to the caller.
> 
> The fact you `Kmalloc` returns a dangling pointer for ZST is an
> implementation detail and the caller shouldn't care (and it also
> couldn't check whether it's a dangling pointer). Nothing in your
> `alloc` doc mention about dangling pointer return for zero-sized alloc
> at all.
> 
> >   
> > > 
> > > I can see 3 ways of fixing this:
> > > 1. Reject ZSTs that have larger alignment than 16 and fix the realloc
> > > implementation to return suitable aligned ZST pointer. I don't
> > > particularly like the idea of allocating ZST can fail though.
> > > 2. Say ZST must be handled by the caller, and make alloc function
> > > unsafe. This means that we essentially revert to the `GlobalAlloc`
> > > design of Rust, and all callers have to check for ZST.
> > > 3. Accept the `old_layout` and use it to check whether the allocation
> > > is ZST allocation.
> > > 
> > > My personal preference is 3.    
> > 
> > There is also 4.
> > 
> > Let `alloc` and `realloc` return a properly aligned dangling pointer for
> > `size == 0` and don't accept dangling pointers in `realloc` and `free`.  
> 
> I'll consider the API design to be bad if I can't pass allocated pointer to
> free. If caller needs to handle ZST specially then we might as well
> just ban it completely.
> 
> > And 5.
> > 
> > Reject the combination of `None` and `size == 0` entirely, as earlier proposed
> > by Benno.
> > 
> > I'm fine with both, 4. and 5. with a slight preference for 4.
> > 
> > I'd also go along with 1., as a mix of 4. and 5.
> > 
> > I really don't like making `alloc` unsafe, and I really don't want to have
> > `old_layout` in `free`. Please let's not discuss this again. :-)  
> 
> I don't buy it.
> 
> Your argument for having `old_layout` is so that the caller doesn't
> need to care about the size. But as demonstrated the caller *does* need
> to care about whether the size is zero.
> 
> Our previous discussion doesn't cover the particular case of ZST and
> you said that it reason arise that we need this extra parameter, then
> it could be added. It feels to me that sane behaviour when it comes
> to ZST allocation is a very good reason.

Just to add that if you *really* want to avoid the old layout param,
another approach accceptable to me is to have a `NonZeroLayout` and have
`alloc` only accept that.

Either `Allocator` trait handles ZST well or it refuses to handle them
at all. No "it works for alloc but not for free" please.

Best,
Gary

> > >     
> > > > 
> > > > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > > > Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> > > > ---
> > > >  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
> > > >  1 file changed, 112 insertions(+)  
>

Re: [PATCH v7 01/26] rust: alloc: add `Allocator` trait

Posted by Danilo Krummrich 1 month, 4 weeks ago

On Sun, Sep 15, 2024 at 09:08:27PM +0100, Gary Guo wrote:
> Just to add that if you *really* want to avoid the old layout param,
> another approach accceptable to me is to have a `NonZeroLayout` and have
> `alloc` only accept that.

That is a good idea, I like that.

> 
> Either `Allocator` trait handles ZST well or it refuses to handle them
> at all. No "it works for alloc but not for free" please.

I agree with that, then we shouldn't handle them at all.

> 
> Best,
> Gary
> 
> > > >     
> > > > > 
> > > > > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > > > > Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> > > > > ---
> > > > >  rust/kernel/alloc.rs | 112 +++++++++++++++++++++++++++++++++++++++++++
> > > > >  1 file changed, 112 insertions(+)  
> > 
>

[RFC PATCH] rust: alloc: pass `old_layout` to `Allocator`

Posted by Danilo Krummrich 1 month, 3 weeks ago

Since this came up a few times, this patch shows how the implementation
looks like with an `old_layout` argument.

Signed-off-by: Danilo Krummrich <dakr@kernel.org>
---
 rust/kernel/alloc.rs           | 32 ++++++++++++--------------
 rust/kernel/alloc/allocator.rs | 28 +++++++++++++++++++----
 rust/kernel/alloc/kbox.rs      | 13 ++++++-----
 rust/kernel/alloc/kvec.rs      | 42 +++++++++++++++++++++++++++-------
 4 files changed, 78 insertions(+), 37 deletions(-)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 2170b53acd0c..78564eeb987d 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -166,7 +166,7 @@ pub unsafe trait Allocator {
     fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
         // SAFETY: Passing `None` to `realloc` is valid by it's safety requirements and asks for a
         // new memory allocation.
-        unsafe { Self::realloc(None, layout, flags) }
+        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags) }
     }
 
     /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
@@ -186,26 +186,23 @@ fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
     ///
     /// # Safety
     ///
-    /// If `ptr == Some(p)`, then `p` must point to an existing and valid memory allocation created
-    /// by this allocator. The alignment encoded in `layout` must be smaller than or equal to the
-    /// alignment requested in the previous `alloc` or `realloc` call of the same allocation.
-    ///
-    /// Additionally, `ptr` is allowed to be `None`; in this case a new memory allocation is
-    /// created.
+    /// - If `ptr == Some(p)`, then `p` must point to an existing and valid memory allocation
+    ///   created by this allocator.
+    /// - `ptr` is allowed to be `None`; in this case a new memory allocation is created.
+    /// - `old_layout` must match the `Layout` the allocation has been created with.
     ///
     /// # Guarantees
     ///
     /// This function has the same guarantees as [`Allocator::alloc`]. When `ptr == Some(p)`, then
     /// it additionally guarantees that:
     /// - the contents of the memory pointed to by `p` are preserved up to the lesser of the new
-    ///   and old size,
-    ///   and old size, i.e.
-    ///   `ret_ptr[0..min(layout.size(), old_size)] == p[0..min(layout.size(), old_size)]`, where
-    ///   `old_size` is the size of the allocation that `p` points at.
-    /// - when the return value is `Err(AllocError)`, then `p` is still valid.
+    ///   and old size, i.e. `ret_ptr[0..min(layout.size(), old_layout.size())] ==
+    ///   p[0..min(layout.size(), old_layout.size())]`.
+    /// - when the return value is `Err(AllocError)`, then `ptr` is still valid.
     unsafe fn realloc(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
+        old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError>;
 
@@ -213,14 +210,13 @@ unsafe fn realloc(
     ///
     /// # Safety
     ///
-    /// `ptr` must point to an existing and valid memory allocation created by this `Allocator` and
-    /// must not be a dangling pointer.
-    ///
-    /// The memory allocation at `ptr` must never again be read from or written to.
-    unsafe fn free(ptr: NonNull<u8>) {
+    /// - `ptr` must point to an existing and valid memory allocation created by this `Allocator`.
+    /// - `layout` must match the `Layout` the allocation has been created with.
+    /// - The memory allocation at `ptr` must never again be read from or written to.
+    unsafe fn free(ptr: NonNull<u8>, layout: Layout) {
         // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
         // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
         // smaller than or equal to the alignment previously used with this allocation.
-        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), Flags(0)) };
+        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0)) };
     }
 }
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 0b586c0361f4..07820c8c4e17 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -54,6 +54,14 @@ fn aligned_size(new_layout: Layout) -> usize {
     layout.size()
 }
 
+/// Returns a properly aligned dangling pointer from the given `layout`.
+fn zst_realloc(layout: Layout) -> NonNull<u8> {
+    let ptr = layout.align() as *mut u8;
+
+    // SAFETY: `layout.align()` (and hence `ptr`) is guaranteed to be non-zero.
+    unsafe { NonNull::new_unchecked(ptr) }
+}
+
 /// # Invariants
 ///
 /// One of the following `krealloc`, `vrealloc`, `kvrealloc`.
@@ -84,11 +92,18 @@ unsafe fn call(
         &self,
         ptr: Option<NonNull<u8>>,
         layout: Layout,
+        old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
         let size = aligned_size(layout);
         let ptr = match ptr {
-            Some(ptr) => ptr.as_ptr(),
+            Some(ptr) => {
+                if old_layout.size() == 0 {
+                    ptr::null()
+                } else {
+                    ptr.as_ptr()
+                }
+            }
             None => ptr::null(),
         };
 
@@ -106,7 +121,7 @@ unsafe fn call(
         };
 
         let ptr = if size == 0 {
-            NonNull::dangling()
+            zst_realloc(layout)
         } else {
             NonNull::new(raw_ptr).ok_or(AllocError)?
         };
@@ -124,10 +139,11 @@ unsafe impl Allocator for Kmalloc {
     unsafe fn realloc(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
+        old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // SAFETY: `ReallocFunc::call` has the same safety requirements as `Allocator::realloc`.
-        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, flags) }
+        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags) }
     }
 }
 
@@ -140,6 +156,7 @@ unsafe impl Allocator for Vmalloc {
     unsafe fn realloc(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
+        old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // TODO: Support alignments larger than PAGE_SIZE.
@@ -150,7 +167,7 @@ unsafe fn realloc(
 
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, flags) }
+        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
     }
 }
 
@@ -163,6 +180,7 @@ unsafe impl Allocator for KVmalloc {
     unsafe fn realloc(
         ptr: Option<NonNull<u8>>,
         layout: Layout,
+        old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // TODO: Support alignments larger than PAGE_SIZE.
@@ -173,6 +191,6 @@ unsafe fn realloc(
 
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, flags) }
+        unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, old_layout, flags) }
     }
 }
diff --git a/rust/kernel/alloc/kbox.rs b/rust/kernel/alloc/kbox.rs
index 6188494f040d..e9e2e94430ef 100644
--- a/rust/kernel/alloc/kbox.rs
+++ b/rust/kernel/alloc/kbox.rs
@@ -5,6 +5,7 @@
 #[allow(unused_imports)] // Used in doc comments.
 use super::allocator::{KVmalloc, Kmalloc, Vmalloc};
 use super::{AllocError, Allocator, Flags};
+use core::alloc::Layout;
 use core::fmt;
 use core::marker::PhantomData;
 use core::mem::ManuallyDrop;
@@ -233,7 +234,7 @@ pub fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>, A>, AllocError> {
         let ptr = if Self::is_zst() {
             NonNull::dangling()
         } else {
-            let layout = core::alloc::Layout::new::<MaybeUninit<T>>();
+            let layout = Layout::new::<MaybeUninit<T>>();
             let ptr = A::alloc(layout, flags)?;
 
             ptr.cast()
@@ -452,14 +453,14 @@ impl<T, A> Drop for Box<T, A>
     A: Allocator,
 {
     fn drop(&mut self) {
-        let size = core::mem::size_of_val::<T>(self);
+        let layout = Layout::for_value::<T>(self);
 
         // SAFETY: The pointer in `self.0` is guaranteed to be valid by the type invariant.
         unsafe { core::ptr::drop_in_place::<T>(self.deref_mut()) };
 
-        if size != 0 {
-            // SAFETY: As `size` is not zero, `self.0` was previously allocated with `A`.
-            unsafe { A::free(self.0.cast()) };
-        }
+        // SAFETY:
+        // - `self.0` was previously allocated with `A`.
+        // - `layout` is equal to the `Layout´ `self.0` was allocated with.
+        unsafe { A::free(self.0.cast(), layout) };
     }
 }
diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 686e969463f8..fb979013562c 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -7,6 +7,7 @@
     AllocError, Allocator, Box, Flags,
 };
 use core::{
+    alloc::Layout,
     fmt,
     marker::PhantomData,
     mem::{ManuallyDrop, MaybeUninit},
@@ -452,21 +453,28 @@ pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocEr
         // We know `cap` is <= `isize::MAX` because of the type invariants of `Self`. So the
         // multiplication by two won't overflow.
         let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
-        let layout = core::alloc::Layout::array::<T>(new_cap).map_err(|_| AllocError)?;
+        let layout = Layout::array::<T>(new_cap).map_err(|_| AllocError)?;
+
+        let old_layout = Layout::array::<T>(self.cap).map_err(|_| AllocError)?;
 
         // We need to make sure that `ptr` is either NULL or comes from a previous call to
         // `realloc_flags`. A `Vec<T, A>`'s `ptr` value is not guaranteed to be NULL and might be
         // dangling after being created with `Vec::new`. Instead, we can rely on `Vec<T, A>`'s
         // capacity to be zero if no memory has been allocated yet.
+        //
+        // Still required? In `Vec::new` we use `NonNull::dangling`, which effectively would be
+        // valid to pass to `A::realloc`, but was never created by one the `Allocator`'s functions.
         let ptr = if cap == 0 {
             None
         } else {
             Some(self.ptr.cast())
         };
 
-        // SAFETY: `ptr` is valid because it's either `None` or comes from a previous call to
-        // `A::realloc`. We also verified that the type is not a ZST.
-        let ptr = unsafe { A::realloc(ptr, layout, flags)? };
+        // SAFETY:
+        // - `ptr` is valid because it's either `None` or comes from a previous call to
+        //   `A::realloc`.
+        // - `old_layout` matches the `Layout` of the preceeding allocation, if any.
+        let ptr = unsafe { A::realloc(ptr, layout, old_layout, flags)? };
 
         self.ptr = ptr.cast();
 
@@ -528,9 +536,16 @@ fn drop(&mut self) {
         };
 
         // If `cap == 0` we never allocated any memory in the first place.
+        //
+        // Same here, theoretically, we know that `NonNull::dangling` from `Vec::new` is valid for
+        // `A::free`, but it was never created by any of the `Allocator` functions.
         if self.cap != 0 {
+            // This can never fail; since this `Layout` is equivalent to the one of the original
+            // allocation.
+            let layout = Layout::array::<T>(self.cap).unwrap();
+
             // SAFETY: `self.ptr` was previously allocated with `A`.
-            unsafe { A::free(self.ptr.cast()) };
+            unsafe { A::free(self.ptr.cast(), layout) };
         }
     }
 }
@@ -751,13 +766,17 @@ pub fn collect(self, flags: Flags) -> Vec<T, A> {
             ptr = buf.as_ptr();
         }
 
+        // This can never fail; since this `Layout` is equivalent to the one of the original
+        // allocation.
+        let old_layout = Layout::array::<T>(cap).unwrap();
+
         // This can never fail, `len` is guaranteed to be smaller than `cap`.
-        let layout = core::alloc::Layout::array::<T>(len).unwrap();
+        let layout = Layout::array::<T>(len).unwrap();
 
         // SAFETY: `buf` points to the start of the backing buffer and `len` is guaranteed to be
         // smaller than `cap`. Depending on `alloc` this operation may shrink the buffer or leaves
         // it as it is.
-        ptr = match unsafe { A::realloc(Some(buf.cast()), layout, flags) } {
+        ptr = match unsafe { A::realloc(Some(buf.cast()), layout, old_layout, flags) } {
             // If we fail to shrink, which likely can't even happen, continue with the existing
             // buffer.
             Err(_) => ptr,
@@ -846,9 +865,16 @@ fn drop(&mut self) {
         unsafe { ptr::drop_in_place(self.as_raw_mut_slice()) };
 
         // If `cap == 0` we never allocated any memory in the first place.
+        //
+        // Same here, theoretically, we know that `NonNull::dangling` from `Vec::new` is valid for
+        // `A::free`, but it was never created by any of the `Allocator` functions.
         if self.cap != 0 {
+            // This can never fail; since this `Layout` is equivalent to the one of the original
+            // allocation.
+            let layout = Layout::array::<T>(self.cap).unwrap();
+
             // SAFETY: `self.buf` was previously allocated with `A`.
-            unsafe { A::free(self.buf.cast()) };
+            unsafe { A::free(self.buf.cast(), layout) };
         }
     }
 }
-- 
2.46.1