Re: [RFC PATCH 06/13] mm/kwatch: add lockless per-task context pool

[sashiko-bot@kernel.org]

Thank you for your contribution! Sashiko AI review found 5 potential issue(s) to consider:
- [High] NMI/IRQ reentrancy during context allocation and release causes shared context corruption and Use-After-Free.
- [High] `kwatch_tsk_ctx_free` frees the context pool without RCU synchronization, causing a Use-After-Free.
- [High] `kwatch_tsk_ctx_prealloc` unconditionally updates global size bounds without reallocating, causing Out-Of-Bounds accesses.
- [High] `kwatch_ctx_pool_size` truncates large inputs to 0, causing `ZERO_SIZE_PTR` dereferences and shift-by-64 Undefined Behavior.
- [High] The context pool stores un-refcounted task pointers, leading to state hijacking if a task exits abruptly (ABA problem).
--

--- Patch [6]: [RFC PATCH 06/13] mm/kwatch: add lockless per-task context pool ---
commit 70e01018388f64f2db26eb6d33251317f890288b
Author: Jinchao Wang <wangjinchao600@gmail.com>

mm/kwatch: add lockless per-task context pool

This patch introduces a lockless per-task context pool for kwatch to 
track task window states in NMI/kprobe contexts without allocating or 
taking locks. It uses an open-addressing array hashed by task_struct 
pointer, claiming slots with cmpxchg() and releasing with 
smp_store_release().

> diff --git a/mm/kwatch/task_ctx.c b/mm/kwatch/task_ctx.c
> new file mode 100644
> index 0000000000000..f8e582f0dcfe9
> --- /dev/null
> +++ b/mm/kwatch/task_ctx.c
> @@ -0,0 +1,105 @@
> +// SPDX-License-Identifier: GPL-2.0
> +#include <linux/slab.h>
> +#include <linux/hash.h>
> +#include <linux/sched.h>
> +#include <linux/log2.h>
> +#include "kwatch.h"
> +
> +static u16 kwatch_ctx_pool_size;
> +static u16 kwatch_ctx_pool_mask;
> +
> +static struct kwatch_tsk_ctx *kwatch_ctx_pool;
> +
> +int kwatch_tsk_ctx_prealloc(u16 max_concurrency)
> +{
> +	if (!max_concurrency)
> +		max_concurrency = 256;
> +
> +	kwatch_ctx_pool_size = roundup_pow_of_two(max_concurrency);
> +	kwatch_ctx_pool_mask = kwatch_ctx_pool_size - 1;

[Severity: High]
Does this code silently truncate the pool size and cause undefined behavior?

If max_concurrency > 32768, roundup_pow_of_two() will return 65536, which gets 
truncated to 0 when assigned to the u16 kwatch_ctx_pool_size. 

If this happens, kcalloc(0, ...) bypasses the allocation check by returning 
ZERO_SIZE_PTR (0x10). Later, calling ilog2(0) in kwatch_tsk_ctx_get() evaluates 
to -1, which results in a shift by 65 in hash_ptr(). 

Also, if max_concurrency is 1, ilog2(1) evaluates to 0, which would result in 
a shift by 64 in hash_ptr(). Both cases trigger undefined behavior.

> +
> +	if (unlikely(!kwatch_ctx_pool)) {

[Severity: High]
Can calling this function multiple times lead to out-of-bounds accesses?

If kwatch_tsk_ctx_prealloc() is called again with a larger max_concurrency, 
kwatch_ctx_pool_size and kwatch_ctx_pool_mask will be updated to the larger 
values. However, since kwatch_ctx_pool is already non-NULL, the array won't be 
reallocated. Subsequent lookups will use the new, larger mask to calculate 
indices, potentially accessing unallocated memory past the end of the original 
array.

> +		kwatch_ctx_pool = kcalloc(kwatch_ctx_pool_size,
> +					  sizeof(struct kwatch_tsk_ctx),
> +					  GFP_KERNEL);
> +		if (!kwatch_ctx_pool)
> +			return -ENOMEM;
> +	}
> +	return 0;
> +}
> +
> +struct kwatch_tsk_ctx *kwatch_tsk_ctx_get(bool can_alloc)
> +{
> +	int start_idx, i, idx;
> +	struct task_struct *t;
> +
> +	if (unlikely(!kwatch_ctx_pool))
> +		return NULL;
> +
> +	start_idx = hash_ptr(current, ilog2(kwatch_ctx_pool_size));
> +
> +	for (i = 0; i < kwatch_ctx_pool_size; i++) {
> +		idx = (start_idx + i) & kwatch_ctx_pool_mask;
> +		t = READ_ONCE(kwatch_ctx_pool[idx].task);
> +		if (t == current)

[Severity: High]
Could this raw address comparison match a completely different task?

Since the task pointers aren't refcounted, if a task is killed (e.g., via 
SIGKILL) while inside a watched function, it exits without calling 
kwatch_tsk_ctx_put(). Its task_struct is eventually freed, but the raw memory 
address is left dangling in kwatch_ctx_pool[idx].task.

If a new task is spawned and its task_struct is allocated at the exact same 
memory address, kwatch_tsk_ctx_get() will find a false match here. The new 
task will unintentionally inherit the dead task's depth and epoch state.

> +			return &kwatch_ctx_pool[idx];
> +	}
> +
> +	if (!can_alloc)
> +		return NULL;
> +
> +	for (i = 0; i < kwatch_ctx_pool_size; i++) {
> +		idx = (start_idx + i) & kwatch_ctx_pool_mask;
> +		t = READ_ONCE(kwatch_ctx_pool[idx].task);
> +		if (!t) {
> +			if (!cmpxchg(&kwatch_ctx_pool[idx].task, NULL, current))

[Severity: High]
Can an NMI interrupting this path corrupt the context slot?

If an NMI interrupts the current task immediately after this cmpxchg() succeeds, 
the NMI handler might call kwatch_tsk_ctx_get(), find the same slot (since 
t == current), use it, and call kwatch_tsk_ctx_put().

kwatch_tsk_ctx_put() will then unconditionally set ctx->task to NULL:

    smp_store_release(&ctx->task, NULL);

When the NMI returns, the interrupted task resumes thinking it owns the slot, 
but the slot is now marked free and can be claimed by another task on another 
CPU, leading to concurrent modification of the same context.

> +				return &kwatch_ctx_pool[idx];
> +		}
> +	}
> +
> +	return NULL;
> +}

[ ... ]

> +void kwatch_tsk_ctx_free(void)
> +{
> +	kfree(kwatch_ctx_pool);

[Severity: High]
Does this code trigger a use-after-free if tasks are actively being watched?

If kwatch_tsk_ctx_free() is called on one CPU while another task (or NMI) 
is inside kwatch_tsk_ctx_get() reading kwatch_ctx_pool on another CPU, 
kfree() frees the memory immediately without waiting for an RCU grace period. 
The lockless reader will then dereference the freed memory.

> +	kwatch_ctx_pool = NULL;
> +}

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260714182243.10687-1-wangjinchao600@gmail.com?part=6