Re: [PATCH v6 4/5] slab: Introduce kmalloc_flex() and family

[Lorenzo Stoakes]

On Wed, Dec 03, 2025 at 03:30:34PM -0800, Kees Cook wrote:
> As done for kmalloc_obj*(), introduce a type-aware allocator for flexible
> arrays, which may also have "counted_by" annotations:
>
> 	ptr = kmalloc(struct_size(ptr, flex_member, count), gfp);
>
> becomes:
>
> 	ptr = kmalloc_flex(*ptr, flex_member, count, gfp);
>
> The internal use of __flex_counter() allows for automatically setting
> the counter member of a struct's flexible array member when it has
> been annotated with __counted_by(), avoiding any missed early size
> initializations while __counted_by() annotations are added to the
> kernel. Additionally, this also checks for "too large" allocations based
> on the type size of the counter variable. For example:
>
> 	if (count > type_max(ptr->flex_counter))
> 		fail...;
> 	size = struct_size(ptr, flex_member, count);
> 	ptr = kmalloc(size, gfp);
> 	ptr->flex_counter = count;
>
> becomes (n.b. unchanged from earlier example):
>
> 	ptr = kmalloc_flex(*ptr, flex_member, count, gfp);
> 	ptr->flex_count = count;
>
> Note that manual initialization of the flexible array counter is still
> required (at some point) after allocation as not all compiler versions
> support the __counted_by annotation yet. But doing it internally makes
> sure they cannot be missed when __counted_by _is_ available, meaning
> that the bounds checker will not trip due to the lack of "early enough"
> initializations that used to work before enabling the stricter bounds
> checking. For example:
>
> 	ptr = kmalloc_flex(*ptr, flex_member, count, gfp);
> 	fill(ptr->flex, count);
> 	ptr->flex_count = count;
>
> This works correctly before adding a __counted_by annotation (since
> nothing is checking ptr->flex accesses against ptr->flex_count). After
> adding the annotation, the bounds sanitizer would trip during fill()
> because ptr->flex_count wasn't set yet. But with kmalloc_flex() setting
> ptr->flex_count internally at allocation time, the existing code works
> without needing to move the ptr->flex_count assignment before the call
> to fill(). (This has been a stumbling block for __counted_by adoption.)
>
> Signed-off-by: Kees Cook <kees@xxxxxxxxxx>
> ---
> Cc: Jonathan Corbet <corbet@xxxxxxx>
> Cc: Vlastimil Babka <vbabka@xxxxxxx>
> Cc: Andrew Morton <akpm@xxxxxxxxxxxxxxxxxxxx>
> Cc: Christoph Lameter <cl@xxxxxxxxxx>
> Cc: David Rientjes <rientjes@xxxxxxxxxx>
> Cc: Roman Gushchin <roman.gushchin@xxxxxxxxx>
> Cc: Harry Yoo <harry.yoo@xxxxxxxxxx>
> Cc: "Gustavo A. R. Silva" <gustavoars@xxxxxxxxxx>
> Cc: <workflows@xxxxxxxxxxxxxxx>
> Cc: <linux-doc@xxxxxxxxxxxxxxx>
> Cc: <linux-mm@xxxxxxxxx>
> Cc: <linux-hardening@xxxxxxxxxxxxxxx>
> ---
>  Documentation/process/deprecated.rst |  7 ++++
>  include/linux/slab.h                 | 48 ++++++++++++++++++++++++++++
>  2 files changed, 55 insertions(+)
>
> diff --git a/Documentation/process/deprecated.rst b/Documentation/process/deprecated.rst
> index 91c628fa2d59..fed56864d036 100644
> --- a/Documentation/process/deprecated.rst
> +++ b/Documentation/process/deprecated.rst
> @@ -387,6 +387,7 @@ allocations. For example, these open coded assignments::
>  	ptr = kzalloc(sizeof(*ptr), gfp);
>  	ptr = kmalloc_array(count, sizeof(*ptr), gfp);
>  	ptr = kcalloc(count, sizeof(*ptr), gfp);
> +	ptr = kmalloc(struct_size(ptr, flex_member, count), gfp);
>  	ptr = kmalloc(sizeof(struct foo, gfp);
>
>  become, respectively::
> @@ -395,4 +396,10 @@ become, respectively::
>  	ptr = kzalloc_obj(*ptr, gfp);
>  	ptr = kmalloc_objs(*ptr, count, gfp);
>  	ptr = kzalloc_objs(*ptr, count, gfp);
> +	ptr = kmalloc_flex(*ptr, flex_member, count, gfp);
>  	__auto_type ptr = kmalloc_obj(struct foo, gfp);
> +
> +If `ptr->flex_member` is annotated with __counted_by(), the allocation
> +will automatically fail if `count` is larger than the maximum
> +representable value that can be stored in the counter member associated
> +with `flex_member`.
> diff --git a/include/linux/slab.h b/include/linux/slab.h
> index 726457daedbd..2656ea610b68 100644
> --- a/include/linux/slab.h
> +++ b/include/linux/slab.h
> @@ -982,6 +982,33 @@ void *kmalloc_nolock_noprof(size_t size, gfp_t gfp_flags, int node);
>  	(TYPE *)KMALLOC(__obj_size, GFP);				\
>  })
>
> +/**
> + * __alloc_flex - Allocate an object that has a trailing flexible array
> + * @KMALLOC: kmalloc wrapper function to use for allocation.
> + * @GFP: GFP flags for the allocation.
> + * @TYPE: type of structure to allocate space for.
> + * @FAM: The name of the flexible array member of @TYPE structure.
> + * @COUNT: how many @FAM elements to allocate space for.
> + *
> + * Returns: Newly allocated pointer to @TYPE with @COUNT-many trailing
> + * @FAM elements, or NULL on failure or if @COUNT cannot be represented
> + * by the member of @TYPE that counts the @FAM elements (annotated via
> + * __counted_by()).
> + */
> +#define __alloc_flex(KMALLOC, GFP, TYPE, FAM, COUNT)			\
> +({									\
> +	const size_t __count = (COUNT);					\
> +	const size_t __obj_size = struct_size_t(TYPE, FAM, __count);	\
> +	TYPE *__obj_ptr;						\
> +	if (WARN_ON_ONCE(overflows_flex_counter_type(TYPE, FAM,	__count))) \
> +		__obj_ptr = NULL;					\

I wonder if this is correct? Because overflows_flex_counter_type() is:

#define overflows_flex_counter_type(TYPE, FAM, COUNT)		\
	(!overflows_type(COUNT, typeof_flex_counter(((TYPE *)NULL)->FAM)))

I.e. returns true if overflows_type() returns false, and whose comment states:

 * Returns: true if @COUNT can be represented in the @FAM's counter. When
 * @FAM is not annotated with __counted_by(), always returns true.

So we warn on when the @COUNT _can_ be represented in @FAM's counter?

I'm seeing the warning for the following:

struct multi_remaps {
	unsigned int nr, cap;
	long offsets[] __counted_by(cap);
};
...
unsigned int cap;
...
multi = kmalloc_flex(*multi, offsets, cap, gfp);

Surely that ! should not be there?

Annnd now I typed that I realise that Linus fixed this up in mainline and I was
working with a stale version of this file :))

Anyway, I see that the comment isn't fixed up, so I think that's something we
should patch, like:

 * Returns: true if @COUNT can be represented in the @FAM's counter. When
 * @FAM is not annotated with __counted_by(), always returns true.

->

 * Returns: true if @COUNT cannot be represented in the @FAM's counter. When
 * @FAM is not annotated with __counted_by(), always returns false.

Cheers, Lorenzo