[PATCH bpf 1/2] bpf: make task_vma iterator stay drained after NULL

[PATCH bpf 1/2] bpf: make task_vma iterator stay drained after NULL

[Kang Chen]

bpf_iter_task_vma_find_next() returns NULL when there is no VMA at or
after data->next_addr, but leaves next_addr unchanged.  Since the
per-VMA lock conversion, the iterator does not hold mmap_lock between
next() calls, so a VMA mapped after exhaustion can be returned by a
later next() call.

The verifier relies on the iterator contract that once next() returns
NULL it keeps returning NULL: process_iter_next_call() marks the
iterator DRAINED and treats later next() results as the constant zero.
A restarted task_vma iteration therefore hands the BPF program a kernel
pointer while the verifier believes R0 == 0.

Add a done flag to the internal iterator data and set it when
bpf_iter_task_vma_find_next() returns NULL.  Later next() calls return
NULL without walking the maple tree again.

Fixes: bee9ef4a40a2 ("bpf: switch task_vma iterator from mmap_lock to per-VMA locks")
Signed-off-by: Kang Chen <k.chen@smail.nju.edu.cn>
---
 kernel/bpf/task_iter.c | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
index 2ee7b3b3c10b..720ff7545072 100644
--- a/kernel/bpf/task_iter.c
+++ b/kernel/bpf/task_iter.c
@@ -854,6 +854,7 @@ struct bpf_iter_task_vma_kern_data {
 	struct mm_struct *mm;
 	struct vm_area_struct snapshot;
 	u64 next_addr;
+	bool done;
 };
 
 struct bpf_iter_task_vma {
@@ -931,6 +932,7 @@ __bpf_kfunc int bpf_iter_task_vma_new(struct bpf_iter_task_vma *it,
 
 	kit->data->snapshot.vm_file = NULL;
 	kit->data->next_addr = addr;
+	kit->data->done = false;
 	return 0;
 
 err_cleanup_iter:
@@ -1007,13 +1009,19 @@ __bpf_kfunc struct vm_area_struct *bpf_iter_task_vma_next(struct bpf_iter_task_v
 	if (!kit->data) /* bpf_iter_task_vma_new failed */
 		return NULL;
 
+	/* Once NULL is returned, next() has to keep returning NULL. */
+	if (kit->data->done)
+		return NULL;
+
 	snap = &kit->data->snapshot;
 
 	bpf_iter_task_vma_snapshot_reset(snap);
 
 	vma = bpf_iter_task_vma_find_next(kit->data);
-	if (!vma)
+	if (!vma) {
+		kit->data->done = true;
 		return NULL;
+	}
 
 	memcpy(snap, vma, sizeof(*snap));
 
-- 
2.34.1

[PATCH bpf 2/2] selftests/bpf: Test that a drained task_vma iterator keeps returning NULL

[Kang Chen]

Add a syscall program that walks the current task's VMAs, records the
last vm_end, drains the iterator and lets a user-space thread map a new
VMA at that address.  The program waits in bpf_loop() until the mapping
is ready, then calls bpf_iter_task_vma_next() once more and stores the
raw result.

The user-space test checks that bpf_iter_task_vma_new() succeeded, that
at least one VMA was drained, that the mapper ran (go_seen), and that
the extra next() returned NULL.  The new VMA makes the test fail on a
kernel without the drain state: the iterator would return the new VMA
again.

Signed-off-by: Kang Chen <k.chen@smail.nju.edu.cn>
---
 .../selftests/bpf/prog_tests/task_vma_iter.c  | 89 +++++++++++++++++++
 .../selftests/bpf/progs/task_vma_iter.c       | 53 +++++++++++
 2 files changed, 142 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/task_vma_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/task_vma_iter.c

diff --git a/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c b/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c
new file mode 100644
index 000000000000..bb9feb840a27
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c
@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <test_progs.h>
+#include <bpf/libbpf.h>
+#include <pthread.h>
+#include <sched.h>
+#include <sys/mman.h>
+#include "task_vma_iter.skel.h"
+
+struct map_ctx {
+	struct task_vma_iter *skel;
+	void *addr;
+	int err;
+};
+
+static void *map_new_vma(void *arg)
+{
+	struct map_ctx *ctx = arg;
+	struct task_vma_iter *skel = ctx->skel;
+	void *addr;
+
+	while (!__atomic_load_n(&skel->bss->drain_done, __ATOMIC_RELAXED))
+		sched_yield();
+
+	if (!skel->bss->new_ret) {
+		addr = (void *)(unsigned long)skel->bss->last_end;
+		ctx->addr = mmap(addr, getpagesize(), PROT_READ | PROT_WRITE,
+				 MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE,
+				 -1, 0);
+		if (ctx->addr == MAP_FAILED)
+			ctx->err = errno;
+	}
+
+	__atomic_store_n(&skel->bss->go, 1, __ATOMIC_RELEASE);
+	return NULL;
+}
+
+static void test_task_vma_iter_next_after_drain(void)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, topts);
+	struct task_vma_iter *skel;
+	struct map_ctx ctx;
+	pthread_t thread;
+	int err, fd;
+
+	skel = task_vma_iter__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "task_vma_iter__open_and_load"))
+		return;
+
+	ctx.skel = skel;
+	ctx.addr = NULL;
+	ctx.err = 0;
+
+	if (!ASSERT_OK(pthread_create(&thread, NULL, map_new_vma, &ctx),
+		       "pthread_create"))
+		goto out;
+
+	fd = bpf_program__fd(skel->progs.iter_task_vma_next_after_drain);
+	err = bpf_prog_test_run_opts(fd, &topts);
+	if (err) {
+		/* Release the mapper if the program never ran. */
+		skel->bss->new_ret = -1;
+		__atomic_store_n(&skel->bss->drain_done, 1, __ATOMIC_RELEASE);
+		__atomic_store_n(&skel->bss->go, 1, __ATOMIC_RELEASE);
+	}
+	pthread_join(thread, NULL);
+
+	if (!ASSERT_OK(err, "test_run_opts err"))
+		goto out;
+	if (!ASSERT_OK(topts.retval, "test_run_opts retval"))
+		goto out;
+
+	ASSERT_OK(ctx.err, "mmap new VMA");
+	ASSERT_OK(skel->bss->new_ret, "new_ret");
+	ASSERT_GT(skel->bss->vma_count, 0, "vma_count");
+	ASSERT_EQ(skel->bss->go_seen, 1, "go_seen");
+	/* last_end is set, so a restarted walk would find the new VMA */
+	ASSERT_EQ(skel->data->next_after_drain, 0, "next_after_drain");
+
+out:
+	if (ctx.addr && ctx.addr != MAP_FAILED)
+		munmap(ctx.addr, getpagesize());
+	task_vma_iter__destroy(skel);
+}
+
+void test_task_vma_iter(void)
+{
+	if (test__start_subtest("next_after_drain"))
+		test_task_vma_iter_next_after_drain();
+}
diff --git a/tools/testing/selftests/bpf/progs/task_vma_iter.c b/tools/testing/selftests/bpf/progs/task_vma_iter.c
new file mode 100644
index 000000000000..9159aa7fc48e
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/task_vma_iter.c
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "vmlinux.h"
+#include "bpf_experimental.h"
+#include <bpf/bpf_helpers.h>
+
+char _license[] SEC("license") = "GPL";
+
+/* Result of a bpf_iter_task_vma_next() call after the iterator was drained */
+__u64 next_after_drain = -1;
+__u64 last_end;
+__u64 drain_done;
+__u64 go;
+int new_ret;
+int go_seen;
+__u64 vma_count;
+
+static long wait_for_go(u32 index, void *ctx)
+{
+	return go ? 1 : 0;
+}
+
+SEC("syscall")
+int iter_task_vma_next_after_drain(const void *ctx)
+{
+	struct task_struct *task = bpf_get_current_task_btf();
+	struct bpf_iter_task_vma it;
+	struct vm_area_struct *vma;
+
+	new_ret = bpf_iter_task_vma_new(&it, task, 0);
+	if (new_ret) {
+		drain_done = 1;
+		go = 1;
+		goto out;
+	}
+
+	while ((vma = bpf_iter_task_vma_next(&it))) {
+		vma_count++;
+		last_end = vma->vm_end;
+	}
+
+	/* Let user space map a new VMA at last_end before the extra next(). */
+	drain_done = 1;
+	bpf_loop(1000000, wait_for_go, NULL, 0);
+	go_seen = go;
+
+	/* Store the raw value; the verifier would remove a branch on it */
+	next_after_drain = (__u64)bpf_iter_task_vma_next(&it);
+
+out:
+	/* The destructor is called even when the constructor failed. */
+	bpf_iter_task_vma_destroy(&it);
+	return 0;
+}
-- 
2.34.1

Re: [PATCH bpf 2/2] selftests/bpf: Test that a drained task_vma iterator keeps returning NULL

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 4 potential
issues to consider:
- [Medium] selftests/bpf: short bpf_loop() spin in
  iter_task_vma_next_after_drain()
- [Low] selftests/bpf: missing acquire/release ordering in map_new_vma()
- [Low] selftests/bpf: mmap() at last_end in map_new_vma() fails at
  TASK_SIZE
- [Low] selftests/bpf: standalone task_vma_iter test instead of
  extending iters.c
--

commit 6861abec39a496e31db7f94849dcb9afd420ad9b
Author: Kang Chen <k.chen@smail.nju.edu.cn>

selftests/bpf: Test that a drained task_vma iterator keeps returning NULL

This commit adds a BPF selftest verifying that calling
bpf_iter_task_vma_next() on an already drained task_vma iterator continues
to return NULL even after a new VMA is mapped at the end of the address
space.

> diff --git a/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c b/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c
> new file mode 100644
> index 0000000000000..bb9feb840a277
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/task_vma_iter.c
> @@ -0,0 +1,89 @@

[ ... ]

> +static void *map_new_vma(void *arg)
> +{
> +	struct map_ctx *ctx = arg;
> +	struct task_vma_iter *skel = ctx->skel;
> +	void *addr;
> +
> +	while (!__atomic_load_n(&skel->bss->drain_done, __ATOMIC_RELAXED))
> +		sched_yield();
> +
> +	if (!skel->bss->new_ret) {
> +		addr = (void *)(unsigned long)skel->bss->last_end;

[Severity: Low]
Can map_new_vma() observe drain_done == 1 before the stores to last_end
and new_ret are visible on weakly-ordered architectures?

iter_task_vma_next_after_drain() writes last_end, new_ret, and drain_done
with plain stores, while map_new_vma() polls drain_done with
__ATOMIC_RELAXED instead of __ATOMIC_ACQUIRE.  Similarly, map_new_vma()
sets go with __ATOMIC_RELEASE, but wait_for_go() reads go with a plain
load.  If the stores are reordered, mmap() with MAP_FIXED_NOREPLACE can
use a stale last_end address and fail the test.

> +		ctx->addr = mmap(addr, getpagesize(), PROT_READ | PROT_WRITE,
> +				 MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE,
> +				 -1, 0);
> +		if (ctx->addr == MAP_FAILED)
> +			ctx->err = errno;
> +	}

[Severity: Low]
Can mmap() at last_end fail with ENOMEM when the highest VMA ends at
TASK_SIZE?

iter_task_vma_next_after_drain() records vm_end of the final VMA in the
process.  On architectures where STACK_TOP equals TASK_SIZE, the stack
VMA's vm_end is TASK_SIZE whenever ASLR is disabled (such as
kernel.randomize_va_space=0 or under gdb) or when randomize_stack_top()
picks a zero offset, leaving no unmapped user address space at last_end.

[ ... ]

> +void test_task_vma_iter(void)
> +{
> +	if (test__start_subtest("next_after_drain"))
> +		test_task_vma_iter_next_after_drain();
> +}

[Severity: Low]
This isn't a bug, but tools/testing/selftests/bpf/prog_tests/iters.c and
tools/testing/selftests/bpf/progs/iters_task_vma.c already test the
open-coded bpf_iter_task_vma iterator.  Could this subtest be added there
instead of introducing standalone task_vma_iter files?

> diff --git a/tools/testing/selftests/bpf/progs/task_vma_iter.c b/tools/testing/selftests/bpf/progs/task_vma_iter.c
> new file mode 100644
> index 0000000000000..9159aa7fc48e4
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/task_vma_iter.c
> @@ -0,0 +1,53 @@

[ ... ]

> +	/* Let user space map a new VMA at last_end before the extra next(). */
> +	drain_done = 1;
> +	bpf_loop(1000000, wait_for_go, NULL, 0);
> +	go_seen = go;

[Severity: Medium]
Can bpf_loop(1000000, wait_for_go, NULL, 0) finish before map_new_vma()
runs and sets go?

wait_for_go() is inlined into a tight loop of 1,000,000 reads that
completes in roughly 1 to 3 ms inside bpf_prog_test_run_syscall() with
migration disabled.  test_task_vma_iter_next_after_drain() starts the
syscall program immediately after pthread_create() without waiting for
map_new_vma() to start.  On a single-CPU system (especially under
!CONFIG_PREEMPT where the kernel loop is not preempted) or under parallel
test_progs load where thread startup and mmap() take more than a few
milliseconds, bpf_loop() exhausts its iterations before go is set and fails
the go_seen assertion.

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261113024500.922824-1-k.chen@smail.nju.edu.cn?part=2