[PATCH] bpf: arena: Flush stale TLB entries before reusing ranges

[Andrea Righi <arighi@nvidia.com>]

arena_free_pages() returns a virtual range to the range tree before it
flushes stale kernel TLB entries for that range. Once the arena lock is
released, another CPU can allocate the range and install new PTEs before
the flush completes.

An allocator running with a stale translation can then initialize the
old page instead of the newly installed page. The freeing CPU
subsequently invalidates the translation and frees the old page, leaving
the new page zeroed despite the allocator having initialized it.

This corruption was observed with scx_eevdf while repeatedly running
stress-ng --sleep 0 --timeout 10.

Keep the range unavailable while clearing its PTEs, invalidating kernel
translations, zapping user mappings, and freeing the old pages.
Reacquire the arena lock and publish the range only after teardown is
complete.

After page teardown, reacquire the arena lock to return the virtual
range to the free-range tree. Resilient spin lock acquisition can fail,
so defer this final step when that happens. Record that teardown is
already complete so the worker only publishes the range instead of
trying to tear down the pages again.

Also add a selftest that primes the old translation on multiple CPUs,
races a free against concurrent allocations, and verifies the contents
of the reallocated page. Exercise both sleepable frees and frees
deferred from a non-sleepable context.

Fixes: b8467290edab ("bpf: arena: make arena kfuncs any context safe")
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/bpf/arena.c                            |  98 ++++++++---
 .../selftests/bpf/prog_tests/arena_reuse.c    | 165 ++++++++++++++++++
 .../testing/selftests/bpf/progs/arena_reuse.c | 113 ++++++++++++
 3 files changed, 356 insertions(+), 20 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/arena_reuse.c
 create mode 100644 tools/testing/selftests/bpf/progs/arena_reuse.c

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 7b6847200b431..73a3103a401b5 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -76,6 +76,8 @@ struct arena_free_span {
 	struct llist_node node;
 	unsigned long uaddr;
 	u32 page_cnt;
+	/* PTE and TLB teardown complete. */
+	bool publish;
 };
 
 u64 bpf_arena_get_kern_vm_start(struct bpf_arena *arena)
@@ -873,15 +875,13 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
 	if (!sleepable)
-		goto defer;
+		goto defer_free;
 
 	ret = raw_res_spin_lock_irqsave(&arena->spinlock, flags);
 
 	/* Can't proceed without holding the spinlock so defer the free */
 	if (ret)
-		goto defer;
-
-	range_tree_set(&arena->rt, pgoff, page_cnt);
+		goto defer_free;
 
 	init_llist_head(&free_pages);
 	cdata.arena = arena;
@@ -890,10 +890,9 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	apply_to_existing_page_range(&init_mm, kaddr, page_cnt << PAGE_SHIFT,
 				     apply_range_clear_cb, &cdata);
 
-	/* drop the lock to do the tlb flush and zap pages */
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 
-	/* ensure no stale TLB entries */
+	/* Ensure no stale kernel translations before publishing the range. */
 	flush_tlb_kernel_range(kaddr, kaddr + (page_cnt * PAGE_SIZE));
 
 	if (page_cnt > 1)
@@ -911,11 +910,17 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 			zap_pages(arena, full_uaddr, 1);
 		__free_page(page);
 	}
+
+	ret = raw_res_spin_lock_irqsave(&arena->spinlock, flags);
+	if (ret)
+		goto defer_publish;
+	range_tree_set(&arena->rt, pgoff, page_cnt);
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 
 	return;
 
-defer:
+defer_free:
 	s = kmalloc_nolock(sizeof(struct arena_free_span), __GFP_ACCOUNT, -1);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 	if (!s)
@@ -928,6 +933,21 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 
 	s->page_cnt = page_cnt;
 	s->uaddr = uaddr;
+	s->publish = false;
+	llist_add(&s->node, &arena->free_spans);
+	irq_work_queue(&arena->free_irq);
+	return;
+
+defer_publish:
+	s = kmalloc_nolock(sizeof(struct arena_free_span), __GFP_ACCOUNT, -1);
+	bpf_map_memcg_exit(old_memcg, new_memcg);
+	if (!s)
+		/* The pages are gone; leaking the unavailable range is safe. */
+		return;
+
+	s->page_cnt = page_cnt;
+	s->uaddr = uaddr;
+	s->publish = true;
 	llist_add(&s->node, &arena->free_spans);
 	irq_work_queue(&arena->free_irq);
 }
@@ -977,7 +997,8 @@ static void arena_free_worker(struct work_struct *work)
 	struct llist_node *list, *pos, *t;
 	struct arena_free_span *s;
 	u64 arena_vm_start, user_vm_start;
-	struct llist_head free_pages;
+	struct llist_head completed_spans, free_pages;
+	struct llist_head publish_spans, teardown_spans;
 	struct clear_range_data cdata;
 	struct page *page;
 	unsigned long full_uaddr;
@@ -991,41 +1012,51 @@ static void arena_free_worker(struct work_struct *work)
 
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
+	init_llist_head(&completed_spans);
 	init_llist_head(&free_pages);
+	init_llist_head(&publish_spans);
+	init_llist_head(&teardown_spans);
 	cdata.arena = arena;
 	cdata.free_pages = &free_pages;
 	arena_vm_start = bpf_arena_get_kern_vm_start(arena);
 	user_vm_start = bpf_arena_get_user_vm_start(arena);
 
 	list = llist_del_all(&arena->free_spans);
-	llist_for_each(pos, list) {
+	llist_for_each_safe(pos, t, list) {
 		s = llist_entry(pos, struct arena_free_span, node);
 		page_cnt = s->page_cnt;
-		kaddr = arena_vm_start + s->uaddr;
 		pgoff = compute_pgoff(arena, s->uaddr);
 
+		if (s->publish) {
+			range_tree_set(&arena->rt, pgoff, page_cnt);
+			llist_add(&s->node, &completed_spans);
+			continue;
+		}
+
+		kaddr = arena_vm_start + s->uaddr;
 		/* clear ptes and collect pages in free_pages llist */
 		apply_to_existing_page_range(&init_mm, kaddr, page_cnt << PAGE_SHIFT,
 					     apply_range_clear_cb, &cdata);
-
-		range_tree_set(&arena->rt, pgoff, page_cnt);
+		llist_add(&s->node, &teardown_spans);
 	}
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 
-	/* Iterate the list again without holding spinlock to do the tlb flush and zap_pages */
-	llist_for_each_safe(pos, t, list) {
+	llist_for_each_safe(pos, t, __llist_del_all(&completed_spans)) {
+		s = llist_entry(pos, struct arena_free_span, node);
+		kfree_nolock(s);
+	}
+
+	/* Tear ranges down while they are still unavailable to allocation. */
+	llist_for_each_safe(pos, t, __llist_del_all(&teardown_spans)) {
 		s = llist_entry(pos, struct arena_free_span, node);
 		page_cnt = s->page_cnt;
-		full_uaddr = clear_lo32(user_vm_start) + s->uaddr;
 		kaddr = arena_vm_start + s->uaddr;
+		full_uaddr = clear_lo32(user_vm_start) + s->uaddr;
 
-		/* ensure no stale TLB entries */
 		flush_tlb_kernel_range(kaddr, kaddr + (page_cnt * PAGE_SIZE));
-
-		/* remove pages from user vmas */
 		zap_pages(arena, full_uaddr, page_cnt);
-
-		kfree_nolock(s);
+		s->publish = true;
+		llist_add(&s->node, &publish_spans);
 	}
 
 	/* free all pages collected by apply_to_existing_page_range() in the first loop */
@@ -1034,6 +1065,33 @@ static void arena_free_worker(struct work_struct *work)
 		__free_page(page);
 	}
 
+	list = __llist_del_all(&publish_spans);
+	if (!list)
+		goto out;
+
+	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags)) {
+		llist_for_each_safe(pos, t, list) {
+			s = llist_entry(pos, struct arena_free_span, node);
+			llist_add(&s->node, &arena->free_spans);
+		}
+		schedule_work(work);
+		goto out;
+	}
+
+	llist_for_each_safe(pos, t, list) {
+		s = llist_entry(pos, struct arena_free_span, node);
+		pgoff = compute_pgoff(arena, s->uaddr);
+		range_tree_set(&arena->rt, pgoff, s->page_cnt);
+		llist_add(&s->node, &completed_spans);
+	}
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+
+	llist_for_each_safe(pos, t, __llist_del_all(&completed_spans)) {
+		s = llist_entry(pos, struct arena_free_span, node);
+		kfree_nolock(s);
+	}
+
+out:
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 }
 
diff --git a/tools/testing/selftests/bpf/prog_tests/arena_reuse.c b/tools/testing/selftests/bpf/prog_tests/arena_reuse.c
new file mode 100644
index 0000000000000..9ac4b7ffe3b8d
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/arena_reuse.c
@@ -0,0 +1,165 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <stdatomic.h>
+#include <sys/sysinfo.h>
+#include <test_progs.h>
+
+#include "arena_reuse.skel.h"
+
+#define MAX_ALLOCATORS 15
+#define ROUNDS 20000
+
+struct reuse_ctx {
+	pthread_barrier_t start;
+	pthread_barrier_t done;
+	pthread_barrier_t next;
+	atomic_int err;
+	int prime_fd;
+};
+
+struct thread_ctx {
+	struct reuse_ctx *shared;
+	int action_fd;
+	int cpu;
+	bool allocator;
+};
+
+static int run_prog(int fd)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	int err;
+
+	err = bpf_prog_test_run_opts(fd, &opts);
+	if (err)
+		return err;
+	return opts.retval ? -EUCLEAN : 0;
+}
+
+static void *reuse_thread(void *arg)
+{
+	struct thread_ctx *ctx = arg;
+	struct reuse_ctx *shared = ctx->shared;
+	cpu_set_t cpuset;
+	int err, round;
+
+	CPU_ZERO(&cpuset);
+	CPU_SET(ctx->cpu, &cpuset);
+	err = pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+	if (err) {
+		atomic_store(&shared->err, err);
+		return NULL;
+	}
+
+	for (round = 0; round < ROUNDS; round++) {
+		if (ctx->allocator && !atomic_load(&shared->err)) {
+			err = run_prog(shared->prime_fd);
+			if (err)
+				atomic_store(&shared->err, err);
+		}
+		pthread_barrier_wait(&shared->start);
+		if (!atomic_load(&shared->err)) {
+			err = run_prog(ctx->action_fd);
+			if (err)
+				atomic_store(&shared->err, err);
+		}
+		pthread_barrier_wait(&shared->done);
+		pthread_barrier_wait(&shared->next);
+	}
+	return NULL;
+}
+
+static void test_arena_reuse(bool deferred_free)
+{
+	struct arena_reuse *skel;
+	struct reuse_ctx ctx = {};
+	struct thread_ctx tctx[MAX_ALLOCATORS + 1];
+	pthread_t threads[MAX_ALLOCATORS + 1];
+	int nr_allocators, nr_threads;
+	int check_fd, init_fd;
+	int i, err, round;
+
+	nr_allocators = MIN(get_nprocs() - 1, MAX_ALLOCATORS);
+	if (nr_allocators < 1) {
+		test__skip();
+		return;
+	}
+	nr_threads = nr_allocators + 1;
+
+	skel = arena_reuse__open();
+	if (!ASSERT_OK_PTR(skel, "open"))
+		return;
+	skel->rodata->deferred_free = deferred_free;
+	if (!ASSERT_OK(arena_reuse__load(skel), "load")) {
+		arena_reuse__destroy(skel);
+		return;
+	}
+
+	init_fd = bpf_program__fd(skel->progs.init_page);
+	check_fd = bpf_program__fd(skel->progs.check_page);
+	if (!ASSERT_OK(run_prog(init_fd), "init_page"))
+		goto out;
+
+	pthread_barrier_init(&ctx.start, NULL, nr_threads + 1);
+	pthread_barrier_init(&ctx.done, NULL, nr_threads + 1);
+	pthread_barrier_init(&ctx.next, NULL, nr_threads + 1);
+	ctx.prime_fd = bpf_program__fd(skel->progs.prime_tlb);
+	atomic_init(&ctx.err, 0);
+
+	for (i = 0; i < nr_threads; i++) {
+		tctx[i].shared = &ctx;
+		tctx[i].cpu = i;
+		tctx[i].allocator = i != nr_allocators;
+		tctx[i].action_fd = bpf_program__fd(tctx[i].allocator ?
+						   skel->progs.alloc_page :
+						   skel->progs.free_page);
+		err = pthread_create(&threads[i], NULL, reuse_thread, &tctx[i]);
+		if (!ASSERT_OK(err, "pthread_create"))
+			goto join;
+	}
+
+	for (round = 0; round < ROUNDS; round++) {
+		skel->bss->alloc_wins = 0;
+		pthread_barrier_wait(&ctx.start);
+		pthread_barrier_wait(&ctx.done);
+
+		if (atomic_load(&ctx.err))
+			goto next;
+		if (skel->bss->alloc_wins > 1) {
+			atomic_store(&ctx.err, -EUCLEAN);
+			goto next;
+		}
+		if (!skel->bss->alloc_wins) {
+			for (i = 0; i < 1000 && run_prog(init_fd); i++)
+				usleep(10);
+			if (i == 1000)
+				atomic_store(&ctx.err, -EIO);
+			goto next;
+		}
+		err = run_prog(check_fd);
+		if (err) {
+			fprintf(stderr, "arena reuse corruption at round %d\n", round);
+			atomic_store(&ctx.err, err);
+		}
+next:
+		pthread_barrier_wait(&ctx.next);
+	}
+
+join:
+	for (i = 0; i < nr_threads; i++)
+		pthread_join(threads[i], NULL);
+
+	ASSERT_OK(atomic_load(&ctx.err), "concurrent arena reuse");
+	pthread_barrier_destroy(&ctx.next);
+	pthread_barrier_destroy(&ctx.done);
+	pthread_barrier_destroy(&ctx.start);
+out:
+	arena_reuse__destroy(skel);
+}
+
+void serial_test_arena_reuse(void)
+{
+	if (test__start_subtest("sleepable"))
+		test_arena_reuse(false);
+	if (test__start_subtest("deferred"))
+		test_arena_reuse(true);
+}
diff --git a/tools/testing/selftests/bpf/progs/arena_reuse.c b/tools/testing/selftests/bpf/progs/arena_reuse.c
new file mode 100644
index 0000000000000..232a7961c3ab0
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/arena_reuse.c
@@ -0,0 +1,113 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+#include "bpf_experimental.h"
+#include <bpf_arena_common.h>
+
+#define PATTERN 0x5a17c0decafef00dULL
+#define PAGE_QWORDS (PAGE_SIZE / sizeof(__u64))
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARENA);
+	__uint(map_flags, BPF_F_MMAPABLE);
+	__uint(max_entries, 1);
+#ifdef __TARGET_ARCH_arm64
+	__ulong(map_extra, 0x1ull << 32);
+#else
+	__ulong(map_extra, 0x1ull << 44);
+#endif
+} arena SEC(".maps");
+
+void __arena *page;
+__u32 alloc_wins;
+const volatile bool deferred_free;
+
+void bpf_rcu_read_lock(void) __ksym;
+void bpf_rcu_read_unlock(void) __ksym;
+
+static __always_inline void fill_page(__u64 __arena *p)
+{
+	int i;
+
+	for (i = 0; i < PAGE_QWORDS; i++)
+		p[i] = PATTERN;
+}
+
+static __always_inline __u64 __arena *get_page(void)
+{
+	__u64 addr = (__u64)page;
+
+	if (!addr)
+		return NULL;
+	return (__u64 __arena *)((__u64)arena_base(&arena) + (__u32)addr);
+}
+
+SEC("syscall")
+int init_page(void *ctx)
+{
+	__u64 __arena *p;
+
+	p = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
+	if (!p)
+		return 1;
+	fill_page(p);
+	page = p;
+	return 0;
+}
+
+SEC("syscall")
+int prime_tlb(void *ctx)
+{
+	__u64 __arena *p = get_page();
+
+	if (!p)
+		return 1;
+	return p[0] != PATTERN;
+}
+
+SEC("syscall")
+int free_page(void *ctx)
+{
+	void __arena *p = page;
+
+	if (!p)
+		return 1;
+	if (deferred_free)
+		bpf_rcu_read_lock();
+	bpf_arena_free_pages(&arena, p, 1);
+	if (deferred_free)
+		bpf_rcu_read_unlock();
+	return 0;
+}
+
+SEC("syscall")
+int alloc_page(void *ctx)
+{
+	__u64 __arena *p;
+
+	p = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
+	if (!p)
+		return 0;
+	fill_page(p);
+	page = p;
+	__sync_fetch_and_add(&alloc_wins, 1);
+	return 0;
+}
+
+SEC("syscall")
+int check_page(void *ctx)
+{
+	__u64 __arena *p = get_page();
+	int i;
+
+	if (!p)
+		return 1;
+	for (i = 0; i < PAGE_QWORDS; i++) {
+		if (p[i] != PATTERN)
+			return i + 1;
+	}
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.55.0