[PATCH bpf-next v2 0/7] bpf: Fix arena memory incoherence

[PATCH bpf-next v2 0/7] bpf: Fix arena memory incoherence

[Emil Tsalapatis]

Setting up arena memory for a task currently requires two operations:
Adjusting its range tree, used for tracking which memory is allocated;
and adjusting its page tables/flushing its TLB state. These operations
cannot happen atomically because their critical sections do not nest.
This lack of atomicity is the source of two bugs that can lead to
incoherence between different users of the same arena, wherein they
observe different pages for the same address.

Address the problem by more finely tracking the state of each address.
Expand the range tree used for address state tracking with a third
state, used to denote whether an address range is unavailable, either
because it is being freed or because it is being populated by a VM
fault. Use this extra state in the arena page freeing/fault logic
to ensure that operations on a single address properly serialize.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
Cc: Puranjay Mohan <puranjay@kernel.org>

v1 -> v2 (https://lore.kernel.org/bpf/20260902070239.16968-1-emil@etsalapatis.com/)

- Use __range_iter_next() in is_range_tree_set() (bot-ci)
- Move -EAGAIN handling to the patch where the errno is introduced (bot-ci)
- Add refactoring patch for (Puranjay)
- Avoid leaking unavailable range tree nodes with atomic mark_available
  operation()
- Avoid retries on res spinlock locking for -EDEADLK (Puranjay)

---

Emil Tsalapatis (7):
  bpf: Update is_range_tree_set to work for consecutive ranges
  bpf: Track availability information for ranges in range tree
  bpf: Fix arena race between page free and alloc leading to incoherency
  bpf: Add explicit state machine for arena free spans
  bpf: Atomically update PTE and range tree in arena VM fault handler
  bpf: Avoid unavailable range leakage during arena_free_pages()
  selftests/bpf: Add arena allocation race tests

 kernel/bpf/arena.c                            | 223 +++++++++++++--
 kernel/bpf/range_tree.c                       | 237 ++++++++++++---
 kernel/bpf/range_tree.h                       |   8 +-
 .../selftests/bpf/prog_tests/arena_race.c     | 270 ++++++++++++++++++
 .../testing/selftests/bpf/progs/arena_race.c  | 159 +++++++++++
 5 files changed, 839 insertions(+), 58 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/arena_race.c
 create mode 100644 tools/testing/selftests/bpf/progs/arena_race.c

 ---

 NOTE: This version has two additional patches, one NFC one a fix for an
 issue introduced in this series while fixing another. For the former, I
 think we should keep it as a separate patch to make it more clear what
 each change does. The latter is murkier - we can avoid the temporary
 regression by spreading Patch 6 into Patches 2 and 3 if need be.

-- 
2.54.0

[PATCH bpf-next v2 1/7] bpf: Update is_range_tree_set to work for consecutive ranges

[Emil Tsalapatis]

The arena range tree currently does not handle consecutive
ranges present in the tree. This is by design: Consecutive
ranges get merged into one by default. However, this design
only lets us track a single bit's worth of state for each
address range, encoded by whether the range is present in
the tree or not (i.e., is it allocated).

We require more fine-grained state tracking for each range.
This means possibly having in the tree consecutive ranges
that cannot be merged because they have different states.
However, existing code implicitly assumes that this scenario
is not possible in its logic.

Expand the logic of is_range_tree_set to handle consecutive
ranges in the tree. The logic change does not affect existing
users and amounts to a defensive check until we enable unmergeable
consecutive ranges in subsequent patches.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/range_tree.c | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/kernel/bpf/range_tree.c b/kernel/bpf/range_tree.c
index 2f28886f3ff7..fdf7ba7aefe0 100644
--- a/kernel/bpf/range_tree.c
+++ b/kernel/bpf/range_tree.c
@@ -180,12 +180,22 @@ int range_tree_clear(struct range_tree *rt, u32 start, u32 len)
 int is_range_tree_set(struct range_tree *rt, u32 start, u32 len)
 {
 	u32 last = start + len - 1;
-	struct range_node *left;
+	struct range_node *rn;
 
-	/* Is this whole range set ? */
-	left = range_it_iter_first(rt, start, last);
-	if (left && left->rn_start <= start && left->rn_last >= last)
-		return 0;
+	for ((rn = range_it_iter_first(rt, start, last)); rn != NULL;
+		rn = __range_it_iter_next(rn, start, last)) {
+		/* Make sure the range covers the start */
+		if (rn->rn_start > start)
+			return -ESRCH;
+
+		/* If it covers the entire range we're done. */
+		if (rn->rn_last >= last)
+			return 0;
+
+		start = rn->rn_last + 1;
+	}
+
+	/* No range to cover [start, last] */
 	return -ESRCH;
 }
 
-- 
2.54.0

[PATCH bpf-next v2 2/7] bpf: Track availability information for ranges in range tree

[Emil Tsalapatis]

Arena address ranges are currently encoded in a range tree: Ranges
present in the tree are free and available for allocation, while
absent ranges are allocated. However, this opens up the arena code to
subtle races between page table updates, range tree updates, and
concurrent allocations that can cause permanent inconsistencies.

Avoiding such races involves distinguishing between memory ranges that
are free and ready to be allocated, and ranges that are being freed but
should not be reused yet. Such tracking is cleanly possible through the
arena's range tree. The range tree is only consumed by arena and has no
additional future consumers, so it can be tailored towards tracking more
state. Alternatives such as deferring range freeing require more
asynchrony and per-range state tracking in the main arena code, and can
introduce additional race conditions.

Expand the tree to track whether a range present in the tree is
available for allocation. For now, all ranges are available: We add
the option to add back a range in an unavailable state, and to move
already present ranges from unavailable to available. We do not
implement other transitions, since they are not required to support
BPF arenas.

The patch adds two operations: Adding a range as unavailable, and
turning a range from unavailable to available. Unavailable ranges are
not mergable, and will be imminently be turned available by the ongoing
arena free() operation that created them. Turning ranges from
unavailable to available is a simple flag change on the range with an
optional merge with adjacent available ranges.

This diff adds -EAGAIN as a possible return value for range_tree_clear().
This is triggered by code in subsequent diffs, when there is an attempt
to use a range that has not been completely freed yet.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/arena.c      |  12 +--
 kernel/bpf/range_tree.c | 177 ++++++++++++++++++++++++++++++++++------
 kernel/bpf/range_tree.h |   4 +-
 3 files changed, 159 insertions(+), 34 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 7b6847200b43..f49b52fa8586 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -317,7 +317,7 @@ static struct bpf_map *arena_map_alloc(union bpf_attr *attr)
 		goto err_free_arena;
 
 	range_tree_init(&arena->rt);
-	err = range_tree_set(&arena->rt, 0, attr->max_entries);
+	err = range_tree_set_avail(&arena->rt, 0, attr->max_entries);
 	if (err)
 		goto err_free_scratch;
 	mutex_init(&arena->lock);
@@ -520,13 +520,13 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	/* Account into memcg of the process that created bpf_arena */
 	ret = bpf_map_alloc_pages(map, NUMA_NO_NODE, 1, &page);
 	if (ret) {
-		range_tree_set(&arena->rt, vmf->pgoff, 1);
+		range_tree_set_avail(&arena->rt, vmf->pgoff, 1);
 		goto out_sigsegv_memcg;
 	}
 
 	ret = apply_to_page_range(&init_mm, kaddr, PAGE_SIZE, apply_range_set_cb, &data);
 	if (ret) {
-		range_tree_set(&arena->rt, vmf->pgoff, 1);
+		range_tree_set_avail(&arena->rt, vmf->pgoff, 1);
 		free_pages_nolock(page, 0);
 		goto out_sigsegv_memcg;
 	}
@@ -766,7 +766,7 @@ static long arena_alloc_pages(struct bpf_arena *arena, long uaddr, long page_cnt
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 	return clear_lo32(arena->user_vm_start) + uaddr32;
 out:
-	range_tree_set(&arena->rt, pgoff + mapped, page_cnt - mapped);
+	range_tree_set_avail(&arena->rt, pgoff + mapped, page_cnt - mapped);
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 	if (mapped) {
 		flush_vmap_cache(kern_vm_start + uaddr32, mapped << PAGE_SHIFT);
@@ -881,7 +881,7 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	if (ret)
 		goto defer;
 
-	range_tree_set(&arena->rt, pgoff, page_cnt);
+	range_tree_set_avail(&arena->rt, pgoff, page_cnt);
 
 	init_llist_head(&free_pages);
 	cdata.arena = arena;
@@ -1008,7 +1008,7 @@ static void arena_free_worker(struct work_struct *work)
 		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);
+		range_tree_set_avail(&arena->rt, pgoff, page_cnt);
 	}
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 
diff --git a/kernel/bpf/range_tree.c b/kernel/bpf/range_tree.c
index fdf7ba7aefe0..456db58650bd 100644
--- a/kernel/bpf/range_tree.c
+++ b/kernel/bpf/range_tree.c
@@ -39,8 +39,15 @@ struct range_node {
 	u32 rn_start;
 	u32 rn_last; /* inclusive */
 	u32 __rn_subtree_last;
+	bool available; /* range is available for allocating. */
 };
 
+/* Is the range available for merging? */
+static inline bool range_available(struct range_node *rn)
+{
+	return rn && rn->available;
+}
+
 static struct range_node *rb_to_range_node(struct rb_node *rb)
 {
 	return rb_entry(rb, struct range_node, rb_range_size);
@@ -55,20 +62,30 @@ static u32 rn_size(struct range_node *rn)
 static inline struct range_node *__find_range(struct range_tree *rt, u32 len)
 {
 	struct rb_node *rb = rt->range_size_root.rb_root.rb_node;
-	struct range_node *best = NULL;
+	struct rb_node *best = NULL;
+	struct range_node *rn;
 
 	while (rb) {
-		struct range_node *rn = rb_to_range_node(rb);
+		rn = rb_to_range_node(rb);
 
 		if (len <= rn_size(rn)) {
-			best = rn;
+			best = rb;
 			rb = rb->rb_right;
 		} else {
 			rb = rb->rb_left;
 		}
 	}
 
-	return best;
+	/* Filter unavailable ranges. */
+	while (best) {
+		rn = rb_to_range_node(best);
+		if (range_available(rn))
+			return rn;
+
+		best = rb_prev(best);
+	}
+
+	return NULL;
 }
 
 s64 range_tree_find(struct range_tree *rt, u32 len)
@@ -135,10 +152,19 @@ range_it_iter_first(struct range_tree *rt, u32 start, u32 last)
 /* Clear the range in this range tree */
 int range_tree_clear(struct range_tree *rt, u32 start, u32 len)
 {
+	u32 first = start;
 	u32 last = start + len - 1;
 	struct range_node *new_rn;
 	struct range_node *rn;
 
+	/* Scan for unavailable ranges and try again if so. */
+	while ((rn = range_it_iter_first(rt, first, last))) {
+		if (!range_available(rn))
+			return -EAGAIN;
+
+		first = rn->rn_last + 1;
+	}
+
 	while ((rn = range_it_iter_first(rt, start, last))) {
 		if (rn->rn_start < start && rn->rn_last > last) {
 			u32 old_last = rn->rn_last;
@@ -153,6 +179,7 @@ int range_tree_clear(struct range_tree *rt, u32 start, u32 len)
 						NUMA_NO_NODE);
 			if (!new_rn)
 				return -ENOMEM;
+			new_rn->available = rn->available;
 			new_rn->rn_start = last + 1;
 			new_rn->rn_last = old_last;
 			range_it_insert(new_rn, rt);
@@ -199,23 +226,12 @@ int is_range_tree_set(struct range_tree *rt, u32 start, u32 len)
 	return -ESRCH;
 }
 
-/* Set the range in this range tree */
-int range_tree_set(struct range_tree *rt, u32 start, u32 len)
+/* Do we have adjacent ranges (and do not overlap with them)? */
+static int range_get_adjacent(struct range_tree *rt, u32 start, u32 last,
+		struct range_node **leftp, struct range_node **rightp)
 {
-	u32 last = start + len - 1;
 	struct range_node *right;
 	struct range_node *left;
-	int err;
-
-	/* Is this whole range already set ? */
-	left = range_it_iter_first(rt, start, last);
-	if (left && left->rn_start <= start && left->rn_last >= last)
-		return 0;
-
-	/* Clear out everything in the range we want to set. */
-	err = range_tree_clear(rt, start, len);
-	if (err)
-		return err;
 
 	/* Do we have a left-adjacent range ? */
 	left = range_it_iter_first(rt, start - 1, start - 1);
@@ -227,34 +243,141 @@ int range_tree_set(struct range_tree *rt, u32 start, u32 len)
 	if (right && right->rn_start != last + 1)
 		return -EFAULT;
 
-	if (left && right) {
+	*leftp = left;
+	*rightp = right;
+
+	return 0;
+}
+
+/*
+ * Merge with adjacent available ranges if possible. The new [start, last]
+ * has already been confirmed to be adjacent with left/right by the caller.
+ */
+static int range_tree_merge(struct range_tree *rt, u32 start, u32 last,
+		struct range_node *left, struct range_node *right)
+{
+	if (range_available(left) && range_available(right)) {
 		/* Combine left and right adjacent ranges */
 		range_it_remove(left, rt);
 		range_it_remove(right, rt);
 		left->rn_last = right->rn_last;
 		range_it_insert(left, rt);
 		kfree_nolock(right);
-	} else if (left) {
+	} else if (range_available(left)) {
 		/* Combine with the left range */
 		range_it_remove(left, rt);
 		left->rn_last = last;
 		range_it_insert(left, rt);
-	} else if (right) {
+	} else if (range_available(right)) {
 		/* Combine with the right range */
 		range_it_remove(right, rt);
 		right->rn_start = start;
 		range_it_insert(right, rt);
 	} else {
-		left = kmalloc_nolock(sizeof(struct range_node), __GFP_ACCOUNT, NUMA_NO_NODE);
-		if (!left)
-			return -ENOMEM;
-		left->rn_start = start;
-		left->rn_last = last;
-		range_it_insert(left, rt);
+		/* No merge available. */
+		return -ENOENT;
 	}
+
 	return 0;
 }
 
+/* Make a range available, possibly merging. */
+int range_tree_make_avail(struct range_tree *rt, u32 start, u32 len)
+{
+	u32 last = start + len - 1;
+	struct range_node *rn;
+	struct range_node *right;
+	struct range_node *left;
+	int err;
+
+	/*
+	 * Confirm the range exists is unavailable,
+	 * and fits the requested range exactly.
+	 */
+	rn = range_it_iter_first(rt, start, last);
+	if (!rn || rn->available)
+		return -EINVAL;
+
+	if (rn->rn_start != start || rn->rn_last != last)
+		return -EINVAL;
+
+	err = range_get_adjacent(rt, start, last, &left, &right);
+	if (err)
+		return err;
+
+	/* If no merging required, just make available. */
+	if (!range_available(left) && !range_available(right)) {
+		rn->available = true;
+		return 0;
+	}
+
+	/* Can merge, remove the range already. */
+	start = rn->rn_start;
+	last = rn->rn_last;
+	range_it_remove(rn, rt);
+	kfree_nolock(rn);
+
+	return range_tree_merge(rt, start, last, left, right);
+}
+
+/* Set the range in this range tree */
+static int range_tree_set(struct range_tree *rt, u32 start, u32 len, bool available)
+{
+	u32 last = start + len - 1;
+	struct range_node *right;
+	struct range_node *left;
+	int err;
+
+	/* Is this whole range already set ? */
+	left = range_it_iter_first(rt, start, last);
+	if (left && left->rn_start <= start && left->rn_last >= last &&
+	    range_available(left) && available)
+		return 0;
+
+	/* Clear out everything in the range we want to set. */
+	err = range_tree_clear(rt, start, len);
+	if (err)
+		return err;
+
+	/* Get adjacent ranges and check for overlaps. */
+	err = range_get_adjacent(rt, start, last, &left, &right);
+	if (err)
+		return err;
+
+	/*
+	 * If the range is not available for allocation, don't merge.
+	 * Unavailable ranges are in the process of being freed and should
+	 * be imminently marked available, so merging them with other
+	 * unavailable ranges will just lead to splitting the range back
+	 * almost immediately.
+	 */
+	if (available) {
+		err = range_tree_merge(rt, start, last, left, right);
+		if (!err)
+			return 0;
+	}
+
+	left = kmalloc_nolock(sizeof(struct range_node), __GFP_ACCOUNT, NUMA_NO_NODE);
+	if (!left)
+		return -ENOMEM;
+	left->available = available;
+	left->rn_start = start;
+	left->rn_last = last;
+	range_it_insert(left, rt);
+
+	return 0;
+}
+
+int range_tree_set_avail(struct range_tree *rt, u32 start, u32 len)
+{
+	return range_tree_set(rt, start, len, true);
+}
+
+int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len)
+{
+	return range_tree_set(rt, start, len, false);
+}
+
 void range_tree_destroy(struct range_tree *rt)
 {
 	struct range_node *rn;
diff --git a/kernel/bpf/range_tree.h b/kernel/bpf/range_tree.h
index ff0b9110eb71..aa27edf451bc 100644
--- a/kernel/bpf/range_tree.h
+++ b/kernel/bpf/range_tree.h
@@ -14,7 +14,9 @@ void range_tree_init(struct range_tree *rt);
 void range_tree_destroy(struct range_tree *rt);
 
 int range_tree_clear(struct range_tree *rt, u32 start, u32 len);
-int range_tree_set(struct range_tree *rt, u32 start, u32 len);
+int range_tree_set_avail(struct range_tree *rt, u32 start, u32 len);
+int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len);
+int range_tree_make_avail(struct range_tree *rt, u32 start, u32 len);
 int is_range_tree_set(struct range_tree *rt, u32 start, u32 len);
 s64 range_tree_find(struct range_tree *rt, u32 len);
 
-- 
2.54.0

[PATCH bpf-next v2 3/7] bpf: Fix arena race between page free and alloc leading to incoherency

[Emil Tsalapatis]

Existing arena kfunc code has an underlying race condition
that can lead to writes being lost from the BPF program's
point of view:

a) A memory range gets freed by operation (1), and its range
is added back to the arena range tree.

b) A concurrent allocation (2) reallocates the range, and does
writes to it. Writes from that CPU may follow the stale TLB
entries into the pages that are about to be freed.

c) (1) invalidates the TLB. The old pages, and any writes done
to them, are now inaccessible. zap_pages() similarly removes the
mappings for userspace threads.

This can be triggered by particularly demanding BPF arena data
structures that constantly allocate and deallocate memory, like
hash table allocations.

Solve this ABA problem by preventing range reallocation until
TLB invalidation/unmapping is complete. First, mark the range
freed but unavailable. Afterwards, drop the spinlock and
flush the kernel TLB and zap user page tables. Then pick up
the lock again and mark the ranges as available once again,
completing the free operation.

Reported-by: Mykola Lysenko <nickolay.lysenko@gmail.com>
Fixes: 317460317a02 ("bpf: Introduce bpf_arena.")
Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/arena.c | 128 +++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 118 insertions(+), 10 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index f49b52fa8586..8bb011aae604 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -76,6 +76,7 @@ struct arena_free_span {
 	struct llist_node node;
 	unsigned long uaddr;
 	u32 page_cnt;
+	bool release_only;
 };
 
 u64 bpf_arena_get_kern_vm_start(struct bpf_arena *arena)
@@ -513,6 +514,14 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 		goto out_sigsegv_memcg;
 
 	ret = range_tree_clear(&arena->rt, vmf->pgoff, 1);
+	/* If a range is unavailable, try again. */
+	if (ret == -EAGAIN) {
+		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+		bpf_map_memcg_exit(old_memcg, new_memcg);
+
+		goto retry;
+	}
+
 	if (ret)
 		goto out_sigsegv_memcg;
 
@@ -542,6 +551,17 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 out_sigsegv:
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 	return VM_FAULT_SIGSEGV;
+
+retry:
+
+	/* Only for special cases (GUP/device drivers). */
+	if (!(vmf->flags & FAULT_FLAG_ALLOW_RETRY))
+		return VM_FAULT_SIGBUS;
+
+	if (!(vmf->flags & FAULT_FLAG_RETRY_NOWAIT))
+		release_fault_lock(vmf);
+
+	return VM_FAULT_RETRY;
 }
 
 static const struct vm_operations_struct arena_vm_ops = {
@@ -855,6 +875,7 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	struct arena_free_span *s;
 	struct clear_range_data cdata;
 	unsigned long flags;
+	bool release_only = false;
 	int ret = 0;
 
 	/* only aligned lower 32-bit are relevant */
@@ -881,7 +902,19 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	if (ret)
 		goto defer;
 
-	range_tree_set_avail(&arena->rt, pgoff, page_cnt);
+	ret = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
+	if (ret) {
+		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+		/*
+		 * For -EAGAIN: An overlapping fault reserves
+		 * the range before installing its PTE.
+		 */
+		if (ret == -ENOMEM || ret == -EAGAIN)
+			goto defer;
+		WARN_ON_ONCE(ret);
+		bpf_map_memcg_exit(old_memcg, new_memcg);
+		return;
+	}
 
 	init_llist_head(&free_pages);
 	cdata.arena = arena;
@@ -911,6 +944,16 @@ 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) {
+		release_only = true;
+		goto defer;
+	}
+
+	ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+	WARN_ON_ONCE(ret);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 
 	return;
@@ -928,6 +971,7 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 
 	s->page_cnt = page_cnt;
 	s->uaddr = uaddr;
+	s->release_only = release_only;
 	llist_add(&s->node, &arena->free_spans);
 	irq_work_queue(&arena->free_irq);
 }
@@ -977,12 +1021,14 @@ 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 free_pages, teardown_spans;
 	struct clear_range_data cdata;
 	struct page *page;
 	unsigned long full_uaddr;
 	long kaddr, page_cnt, pgoff;
 	unsigned long flags;
+	bool retry = false;
+	int ret;
 
 	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags)) {
 		schedule_work(work);
@@ -992,28 +1038,57 @@ static void arena_free_worker(struct work_struct *work)
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
 	init_llist_head(&free_pages);
+	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->release_only) {
+			ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
+			WARN_ON_ONCE(ret);
+			kfree_nolock(s);
+			continue;
+		}
+
+		kaddr = arena_vm_start + s->uaddr;
+
+		ret = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
+		if (ret) {
+			/* Kick off another attempt at the end of this call. */
+			if (ret == -EAGAIN) {
+				llist_add(pos, &arena->free_spans);
+				retry = true;
+				continue;
+			}
+
+			/*
+			 * An -ENOMEM failure is the same failure mode as in
+			 * the defer: path of arena_free_pages(). Do not treat
+			 * the leak as a bug.
+			 */
+			if (ret != -ENOMEM)
+				WARN_ON_ONCE(ret);
+
+			kfree_nolock(s);
+			continue;
+		}
+
 		/* 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_avail(&arena->rt, pgoff, page_cnt);
+		__llist_add(pos, &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) {
+	/* Keep ranges unavailable until their stale translations are gone. */
+	llist_for_each_safe(pos, t, READ_ONCE(teardown_spans.first)) {
 		s = llist_entry(pos, struct arena_free_span, node);
 		page_cnt = s->page_cnt;
 		full_uaddr = clear_lo32(user_vm_start) + s->uaddr;
@@ -1024,8 +1099,6 @@ static void arena_free_worker(struct work_struct *work)
 
 		/* remove pages from user vmas */
 		zap_pages(arena, full_uaddr, page_cnt);
-
-		kfree_nolock(s);
 	}
 
 	/* free all pages collected by apply_to_existing_page_range() in the first loop */
@@ -1034,7 +1107,42 @@ static void arena_free_worker(struct work_struct *work)
 		__free_page(page);
 	}
 
+	if (!llist_empty(&teardown_spans)) {
+		if (raw_res_spin_lock_irqsave(&arena->spinlock, flags)) {
+			llist_for_each_safe(pos, t, __llist_del_all(&teardown_spans)) {
+				s = llist_entry(pos, struct arena_free_span, node);
+				s->release_only = true;
+				llist_add(pos, &arena->free_spans);
+			}
+
+			schedule_work(work);
+			bpf_map_memcg_exit(old_memcg, new_memcg);
+			return;
+		}
+
+		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;
+			pgoff = compute_pgoff(arena, s->uaddr);
+			/*
+			 * This range tree operation does not allocate memory,
+			 * and so should never fail regardless of contention
+			 * or memory pressure. This is in contrast to regular
+			 * inserts that _can_ fail under memory pressure and
+			 * force us to defer the free.
+			 */
+			ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
+			WARN_ON_ONCE(ret);
+			kfree_nolock(s);
+		}
+		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+	}
+
 	bpf_map_memcg_exit(old_memcg, new_memcg);
+
+	/* Retry if any region was unavailable for free. */
+	if (retry)
+		schedule_work(work);
 }
 
 static void arena_free_irq(struct irq_work *iw)
-- 
2.54.0

[PATCH bpf-next v2 4/7] bpf: Add explicit state machine for arena free spans

[Emil Tsalapatis]

The arena_free_worker() call currently tracks the status
of each arena_free_span by placing it into a separate
list. This requires multiple list manipulation calls
during the freeing operation, complicating the code for
no reason.

Add an explicit state machine for span state. This
avoids encoding the states within temporary list membership,
and also allows for safely reschedulign freeing work for each
span separately.

Suggested-by: Puranjay Mohan <puranjay@kernel.org>
Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/arena.c | 108 +++++++++++++++++++++++++++------------------
 1 file changed, 65 insertions(+), 43 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 8bb011aae604..ad58aeaa0c87 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -72,11 +72,18 @@ struct bpf_arena {
 static void arena_free_worker(struct work_struct *work);
 static void arena_free_irq(struct irq_work *iw);
 
+enum arena_free_span_state {
+	ARENA_FREE_SPAN_NOT_STARTED,	/* Freeing not started */
+	ARENA_FREE_SPAN_UNAVAIL,	/* Region cleared & unavailable */
+	ARENA_FREE_SPAN_ZAPPED,		/* Region zapped and cleared */
+	ARENA_FREE_SPAN_FAILED,		/* Freeing operation cannot continue */
+};
+
 struct arena_free_span {
 	struct llist_node node;
 	unsigned long uaddr;
 	u32 page_cnt;
-	bool release_only;
+	enum arena_free_span_state state;
 };
 
 u64 bpf_arena_get_kern_vm_start(struct bpf_arena *arena)
@@ -971,7 +978,7 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 
 	s->page_cnt = page_cnt;
 	s->uaddr = uaddr;
-	s->release_only = release_only;
+	s->state = release_only ? ARENA_FREE_SPAN_ZAPPED : ARENA_FREE_SPAN_NOT_STARTED;
 	llist_add(&s->node, &arena->free_spans);
 	irq_work_queue(&arena->free_irq);
 }
@@ -1021,7 +1028,7 @@ 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, teardown_spans;
+	struct llist_head free_pages;
 	struct clear_range_data cdata;
 	struct page *page;
 	unsigned long full_uaddr;
@@ -1038,7 +1045,6 @@ static void arena_free_worker(struct work_struct *work)
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
 	init_llist_head(&free_pages);
-	init_llist_head(&teardown_spans);
 	cdata.arena = arena;
 	cdata.free_pages = &free_pages;
 	arena_vm_start = bpf_arena_get_kern_vm_start(arena);
@@ -1047,26 +1053,18 @@ static void arena_free_worker(struct work_struct *work)
 	list = llist_del_all(&arena->free_spans);
 	llist_for_each_safe(pos, t, list) {
 		s = llist_entry(pos, struct arena_free_span, node);
-		page_cnt = s->page_cnt;
-		pgoff = compute_pgoff(arena, s->uaddr);
-
-		if (s->release_only) {
-			ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
-			WARN_ON_ONCE(ret);
-			kfree_nolock(s);
+		if (s->state != ARENA_FREE_SPAN_NOT_STARTED)
 			continue;
-		}
 
+		page_cnt = s->page_cnt;
+		pgoff = compute_pgoff(arena, s->uaddr);
 		kaddr = arena_vm_start + s->uaddr;
 
 		ret = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
 		if (ret) {
 			/* Kick off another attempt at the end of this call. */
-			if (ret == -EAGAIN) {
-				llist_add(pos, &arena->free_spans);
-				retry = true;
+			if (ret == -EAGAIN)
 				continue;
-			}
 
 			/*
 			 * An -ENOMEM failure is the same failure mode as in
@@ -1076,20 +1074,24 @@ static void arena_free_worker(struct work_struct *work)
 			if (ret != -ENOMEM)
 				WARN_ON_ONCE(ret);
 
-			kfree_nolock(s);
+			s->state = ARENA_FREE_SPAN_FAILED;
 			continue;
 		}
 
+		s->state = ARENA_FREE_SPAN_UNAVAIL;
+
 		/* 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);
-		__llist_add(pos, &teardown_spans);
 	}
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 
 	/* Keep ranges unavailable until their stale translations are gone. */
-	llist_for_each_safe(pos, t, READ_ONCE(teardown_spans.first)) {
+	llist_for_each_safe(pos, t, list) {
 		s = llist_entry(pos, struct arena_free_span, node);
+		if (s->state != ARENA_FREE_SPAN_UNAVAIL)
+			continue;
+
 		page_cnt = s->page_cnt;
 		full_uaddr = clear_lo32(user_vm_start) + s->uaddr;
 		kaddr = arena_vm_start + s->uaddr;
@@ -1099,6 +1101,12 @@ static void arena_free_worker(struct work_struct *work)
 
 		/* remove pages from user vmas */
 		zap_pages(arena, full_uaddr, page_cnt);
+
+		/*
+		 * Used to avoid zapping twice if we fail the lock acquisition
+		 * below and rerun the span through this function.
+		 */
+		s->state = ARENA_FREE_SPAN_ZAPPED;
 	}
 
 	/* free all pages collected by apply_to_existing_page_range() in the first loop */
@@ -1107,40 +1115,54 @@ static void arena_free_worker(struct work_struct *work)
 		__free_page(page);
 	}
 
-	if (!llist_empty(&teardown_spans)) {
-		if (raw_res_spin_lock_irqsave(&arena->spinlock, flags)) {
-			llist_for_each_safe(pos, t, __llist_del_all(&teardown_spans)) {
-				s = llist_entry(pos, struct arena_free_span, node);
-				s->release_only = true;
-				llist_add(pos, &arena->free_spans);
+	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);
+
+			if (s->state == ARENA_FREE_SPAN_FAILED) {
+				kfree_nolock(s);
+				continue;
 			}
 
-			schedule_work(work);
-			bpf_map_memcg_exit(old_memcg, new_memcg);
-			return;
+			llist_add(pos, &arena->free_spans);
+			retry = true;
 		}
+		goto done;
+	}
 
-		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;
-			pgoff = compute_pgoff(arena, s->uaddr);
-			/*
-			 * This range tree operation does not allocate memory,
-			 * and so should never fail regardless of contention
-			 * or memory pressure. This is in contrast to regular
-			 * inserts that _can_ fail under memory pressure and
-			 * force us to defer the free.
-			 */
-			ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
-			WARN_ON_ONCE(ret);
+	llist_for_each_safe(pos, t, list) {
+		s = llist_entry(pos, struct arena_free_span, node);
+
+		/* Remove the spans of failed allocations. */
+		if (s->state == ARENA_FREE_SPAN_FAILED) {
 			kfree_nolock(s);
+			continue;
 		}
-		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+
+		if (s->state == ARENA_FREE_SPAN_NOT_STARTED) {
+			llist_add(pos, &arena->free_spans);
+			retry = true;
+			continue;
+		}
+
+		page_cnt = s->page_cnt;
+		pgoff = compute_pgoff(arena, s->uaddr);
+		/*
+		 * This range tree operation does not allocate memory,
+		 * and so should never fail regardless of contention
+		 * or memory pressure. This is in contrast to regular
+		 * inserts that _can_ fail under memory pressure and
+		 * force us to defer the free.
+		 */
+		ret = range_tree_make_avail(&arena->rt, pgoff, page_cnt);
+		WARN_ON_ONCE(ret);
+		kfree_nolock(s);
 	}
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 
+done:
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 
-	/* Retry if any region was unavailable for free. */
 	if (retry)
 		schedule_work(work);
 }
-- 
2.54.0

[PATCH bpf-next v2 5/7] bpf: Atomically update PTE and range tree in arena VM fault handler

[Emil Tsalapatis]

The arena allocation code currently has a race in the fault handler
that can cause userspace threads to write to the wrong arena page.

a) The fault handler removes a range from the range tree to mark the
addresses as allocated, then installs a page A into the kernel page
tables.

b) A concurrent free/reallocation removes A and installs a page B.

c) The fault handler still goes ahead with installing page A in the
page table. The kernel sees page B, while userspace sees page A.

There is no way to protect the PTE installation and the range tree
modification simultaneously, because we cannot nest the synchronization
primitives for their respective critical sections. PTE allocation
may require allocations due to PTE reclamation, and its spinlock
becomes sleepable under PREEMPT_RT. Thus we cannot do this operation
while holding the range tree spinlock. There is no public API for
manually taking this spinlock, so we nest the range tree operation
inside it.

Solve this issue by adjusting the range tree in two steps. First,
mark the address of the page being allocated as unavailable. Then
drop the range spinlock, insert the PTE, take the range spinlock
again, and fully remove it from the range tree. Concurrent free
operations get serialized to before the fault handler, while
it is not possible to allocate the page once it has been reserved.
Concurrent page fault handler calls retry until the page is fully
allocated by the original call.

Also return VM_FAULT_RETRY for transient allocation failures. These
are a) faulting on pages that are temporarily marked unavailable in
the range tree and b) rqspinlock acquisition failures.

Fixes: b8467290edab ("bpf: arena: make arena kfuncs any context safe")
Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/arena.c      | 41 +++++++++++++++++++++++++++++++----------
 kernel/bpf/range_tree.c | 14 ++++++++++++++
 kernel/bpf/range_tree.h |  1 +
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index ad58aeaa0c87..8594b76dae61 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -492,18 +492,18 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	struct page *page;
 	long kbase, kaddr;
 	unsigned long flags;
+	vm_fault_t ret_fault;
 	int ret;
 
 	kbase = bpf_arena_get_kern_vm_start(arena);
 	kaddr = kbase + (u32)(vmf->address);
 
-	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags))
-		/*
-		 * A failed lock means a possible deadlock was detected. Don't
-		 * return VM_FAULT_RETRY: this handler never took mmap_lock, but
-		 * the fault path would re-take it on retry and deadlock. Fail.
-		 */
+	ret = raw_res_spin_lock_irqsave(&arena->spinlock, flags);
+	/* If we are deadlocking somehow, no way to ensure forward progress. */
+	if (ret == -EDEADLK)
 		return VM_FAULT_SIGBUS;
+	if (ret)
+		goto retry;
 
 	page = vmalloc_to_page((void *)kaddr);
 	if (page) {
@@ -549,10 +549,31 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	flush_vmap_cache(kaddr, PAGE_SIZE);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 out:
-	page_ref_add(page, 1);
+	/* Reserve the page while installing its user PTE without the arena lock. */
+	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
+	ret = range_tree_set_unavail(&arena->rt, vmf->pgoff, 1);
+	bpf_map_memcg_exit(old_memcg, new_memcg);
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
-	vmf->page = page;
-	return 0;
+	if (ret) {
+		if (ret == -EAGAIN)
+			goto retry;
+		return VM_FAULT_OOM;
+	}
+
+	ret_fault = vmf_insert_page(vmf->vma, vmf->address, page);
+
+	while ((ret = raw_res_spin_lock_irqsave(&arena->spinlock, flags))) {
+		/* If we somehow deadlocked stop trying to take the lock. */
+		if (ret == -EDEADLK)
+			return VM_FAULT_SIGBUS;
+
+		cond_resched();
+	}
+
+	ret = range_tree_remove_unavail(&arena->rt, vmf->pgoff, 1);
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+	WARN_ON_ONCE(ret);
+	return ret_fault;
 out_sigsegv_memcg:
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 out_sigsegv:
@@ -648,7 +669,7 @@ static int arena_map_mmap(struct bpf_map *map, struct vm_area_struct *vma)
 	 * of user_vm_start. Set VM_DONTCOPY to prevent arena VMA from
 	 * being copied into the child process on fork.
 	 */
-	vm_flags_set(vma, VM_DONTEXPAND | VM_DONTCOPY);
+	vm_flags_set(vma, VM_DONTEXPAND | VM_DONTCOPY | VM_MIXEDMAP);
 	vma->vm_ops = &arena_vm_ops;
 	return 0;
 }
diff --git a/kernel/bpf/range_tree.c b/kernel/bpf/range_tree.c
index 456db58650bd..46a32623efb2 100644
--- a/kernel/bpf/range_tree.c
+++ b/kernel/bpf/range_tree.c
@@ -378,6 +378,20 @@ int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len)
 	return range_tree_set(rt, start, len, false);
 }
 
+int range_tree_remove_unavail(struct range_tree *rt, u32 start, u32 len)
+{
+	u32 last = start + len - 1;
+	struct range_node *rn;
+
+	rn = range_it_iter_first(rt, start, last);
+	if (!rn || rn->available || rn->rn_start != start || rn->rn_last != last)
+		return -EINVAL;
+
+	range_it_remove(rn, rt);
+	kfree_nolock(rn);
+	return 0;
+}
+
 void range_tree_destroy(struct range_tree *rt)
 {
 	struct range_node *rn;
diff --git a/kernel/bpf/range_tree.h b/kernel/bpf/range_tree.h
index aa27edf451bc..4b12ef51cc0b 100644
--- a/kernel/bpf/range_tree.h
+++ b/kernel/bpf/range_tree.h
@@ -16,6 +16,7 @@ void range_tree_destroy(struct range_tree *rt);
 int range_tree_clear(struct range_tree *rt, u32 start, u32 len);
 int range_tree_set_avail(struct range_tree *rt, u32 start, u32 len);
 int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len);
+int range_tree_remove_unavail(struct range_tree *rt, u32 start, u32 len);
 int range_tree_make_avail(struct range_tree *rt, u32 start, u32 len);
 int is_range_tree_set(struct range_tree *rt, u32 start, u32 len);
 s64 range_tree_find(struct range_tree *rt, u32 len);
-- 
2.54.0

[PATCH bpf-next v2 6/7] bpf: Avoid unavailable range leakage during arena_free_pages()

[Emil Tsalapatis]

The arena_free_pages() call is designed to leak the attempted
freed page when it is unable to either directly complete or
defer the operation. This is valid for ranges for which freeing
has not started. For those already marked as unavailable in the
range tree, leaving them permanently unavailable breaks the
assumption that unavailable ranges denote an operation currently
in progress.

Atomically turn the range available instead. The only point a
range can be leaked as unavailable is when failing to regrab the
arena spinlock after zapping the range. At that point, the only
operation left to complete the free is to mark the range available.
The spinlock is required to merge the range with adjacent free ones
to avoid range space fragmentation, but this is not necessary for
correctness or safety. Just flip the range node back to available
without the spinlock then.

This is safe because unavailable nodes are not affected by operations
done to the other nodes, and vice versa for changing the node's available
state. We can thus treat pointers to newly created unavailable nodes
as transient references we drop when we mark the node available.
We do not make the layout of the range_node part of the tree's public
API.

The new operation causes fragmentation in the range tree. This is
preferable to leaking the range, even setting aside the correctness
issue.

Fixes: 226d4d4d1272 ("bpf: Fix arena race between page free and alloc leading to incoherency")
Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 kernel/bpf/arena.c      | 42 +++++++++++++++++++++++++++++++----------
 kernel/bpf/range_tree.c | 36 ++++++++++++++++++++++++++++++-----
 kernel/bpf/range_tree.h |  5 ++++-
 3 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 8594b76dae61..299525f93bbd 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -489,6 +489,7 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	struct bpf_map *map = vmf->vma->vm_file->private_data;
 	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
 	struct mem_cgroup *new_memcg, *old_memcg;
+	struct range_node *unavail_node;
 	struct page *page;
 	long kbase, kaddr;
 	unsigned long flags;
@@ -551,10 +552,11 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 out:
 	/* Reserve the page while installing its user PTE without the arena lock. */
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
-	ret = range_tree_set_unavail(&arena->rt, vmf->pgoff, 1);
+	unavail_node = range_tree_set_unavail(&arena->rt, vmf->pgoff, 1);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
-	if (ret) {
+	if (IS_ERR(unavail_node)) {
+		ret = PTR_ERR(unavail_node);
 		if (ret == -EAGAIN)
 			goto retry;
 		return VM_FAULT_OOM;
@@ -895,6 +897,7 @@ static void zap_pages(struct bpf_arena *arena, long uaddr, long page_cnt)
 static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt, bool sleepable)
 {
 	struct mem_cgroup *new_memcg, *old_memcg;
+	struct range_node *unavail_node = NULL;
 	u64 full_uaddr, uaddr_end;
 	long kaddr, pgoff;
 	struct page *page;
@@ -930,8 +933,9 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	if (ret)
 		goto defer;
 
-	ret = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
-	if (ret) {
+	unavail_node = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
+	if (IS_ERR(unavail_node)) {
+		ret = PTR_ERR(unavail_node);
 		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 		/*
 		 * For -EAGAIN: An overlapping fault reserves
@@ -989,13 +993,29 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 defer:
 	s = kmalloc_nolock(sizeof(struct arena_free_span), __GFP_ACCOUNT, -1);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
-	if (!s)
+	if (!s) {
+		/*
+		 * Directly mark the region available. Unavailable regions must
+		 * always be transient, and a permanent one breaks the assumptions
+		 * made in the allocation/faulting code. The operation is safe because
+		 * unavailable nodes are not modified by the range tree code without a reference
+		 * to the node. The modification from unavailable to available also does
+		 * not require touching anything in the tree apart from the node itself,
+		 * so it can be done locklessly. The downside is fragmentation in the tree,
+		 * since we cannot merge with neighbors, but this is a) safe and b) better
+		 * than leaking the range.
+		 */
+		if (release_only)
+			range_node_mark_available(unavail_node);
+
 		/*
-		 * If allocation fails in non-sleepable context, pages are intentionally left
-		 * inaccessible (leaked) until the arena is destroyed. Cleanup or retries are not
-		 * possible here, so we intentionally omit them for safety.
+		 * If we haven't even zapped the pages, intentionally leave them
+		 * inaccessible (leaked) until the arena is destroyed. Cleanup or
+		 * retries are not possible here, so we intentionally omit them for safety.
 		 */
+
 		return;
+	}
 
 	s->page_cnt = page_cnt;
 	s->uaddr = uaddr;
@@ -1048,6 +1068,7 @@ static void arena_free_worker(struct work_struct *work)
 	struct mem_cgroup *new_memcg, *old_memcg;
 	struct llist_node *list, *pos, *t;
 	struct arena_free_span *s;
+	struct range_node *unavail_node;
 	u64 arena_vm_start, user_vm_start;
 	struct llist_head free_pages;
 	struct clear_range_data cdata;
@@ -1081,8 +1102,9 @@ static void arena_free_worker(struct work_struct *work)
 		pgoff = compute_pgoff(arena, s->uaddr);
 		kaddr = arena_vm_start + s->uaddr;
 
-		ret = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
-		if (ret) {
+		unavail_node = range_tree_set_unavail(&arena->rt, pgoff, page_cnt);
+		if (IS_ERR(unavail_node)) {
+			ret = PTR_ERR(unavail_node);
 			/* Kick off another attempt at the end of this call. */
 			if (ret == -EAGAIN)
 				continue;
diff --git a/kernel/bpf/range_tree.c b/kernel/bpf/range_tree.c
index 46a32623efb2..62ebf4df51cb 100644
--- a/kernel/bpf/range_tree.c
+++ b/kernel/bpf/range_tree.c
@@ -3,6 +3,7 @@
 #include <linux/interval_tree_generic.h>
 #include <linux/slab.h>
 #include <linux/bpf.h>
+#include <linux/err.h>
 #include "range_tree.h"
 
 /*
@@ -45,7 +46,21 @@ struct range_node {
 /* Is the range available for merging? */
 static inline bool range_available(struct range_node *rn)
 {
-	return rn && rn->available;
+	/* Pairs with smp_store_release() in range_node_mark_available(). */
+	return rn && smp_load_acquire(&rn->available);
+}
+
+/*
+ * Mark a node as available. Designed to be used locklessly
+ * as a last-ditch effort to avoid leaking unavailable range
+ * tree nodes when all attempts to directly or indirectly
+ * free an arena region has failed. See arena_free_pages()
+ * for more info.
+ */
+void range_node_mark_available(struct range_node *rn)
+{
+	/* Pairs with smp_load_acquire() in range_available(). */
+	smp_store_release(&rn->available, true);
 }
 
 static struct range_node *rb_to_range_node(struct rb_node *rb)
@@ -321,7 +336,8 @@ int range_tree_make_avail(struct range_tree *rt, u32 start, u32 len)
 }
 
 /* Set the range in this range tree */
-static int range_tree_set(struct range_tree *rt, u32 start, u32 len, bool available)
+static int range_tree_set(struct range_tree *rt, u32 start, u32 len, bool available,
+			  struct range_node **new_rn)
 {
 	u32 last = start + len - 1;
 	struct range_node *right;
@@ -364,18 +380,28 @@ static int range_tree_set(struct range_tree *rt, u32 start, u32 len, bool availa
 	left->rn_start = start;
 	left->rn_last = last;
 	range_it_insert(left, rt);
+	if (new_rn)
+		*new_rn = left;
 
 	return 0;
 }
 
 int range_tree_set_avail(struct range_tree *rt, u32 start, u32 len)
 {
-	return range_tree_set(rt, start, len, true);
+	return range_tree_set(rt, start, len, true, NULL);
 }
 
-int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len)
+struct range_node *range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len)
 {
-	return range_tree_set(rt, start, len, false);
+	struct range_node *rn = NULL;
+	int err;
+
+	err = range_tree_set(rt, start, len, false, &rn);
+	if (err)
+		return ERR_PTR(err);
+	if (WARN_ON_ONCE(!rn))
+		return ERR_PTR(-EINVAL);
+	return rn;
 }
 
 int range_tree_remove_unavail(struct range_tree *rt, u32 start, u32 len)
diff --git a/kernel/bpf/range_tree.h b/kernel/bpf/range_tree.h
index 4b12ef51cc0b..79295abe3684 100644
--- a/kernel/bpf/range_tree.h
+++ b/kernel/bpf/range_tree.h
@@ -10,12 +10,15 @@ struct range_tree {
 	struct rb_root_cached range_size_root;
 };
 
+struct range_node;
+
 void range_tree_init(struct range_tree *rt);
 void range_tree_destroy(struct range_tree *rt);
 
 int range_tree_clear(struct range_tree *rt, u32 start, u32 len);
 int range_tree_set_avail(struct range_tree *rt, u32 start, u32 len);
-int range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len);
+struct range_node *range_tree_set_unavail(struct range_tree *rt, u32 start, u32 len);
+void range_node_mark_available(struct range_node *rn);
 int range_tree_remove_unavail(struct range_tree *rt, u32 start, u32 len);
 int range_tree_make_avail(struct range_tree *rt, u32 start, u32 len);
 int is_range_tree_set(struct range_tree *rt, u32 start, u32 len);
-- 
2.54.0

[PATCH bpf-next v2 7/7] selftests/bpf: Add arena allocation race tests

[Emil Tsalapatis]

Add selftests to handle arena page allocation-related races.
Ensure that concurrent frees and nonsleepable/sleepable page
allocations, as well as allocations from userspace, do not
lead to inconsistent or lost data.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 .../selftests/bpf/prog_tests/arena_race.c     | 270 ++++++++++++++++++
 .../testing/selftests/bpf/progs/arena_race.c  | 159 +++++++++++
 2 files changed, 429 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/arena_race.c
 create mode 100644 tools/testing/selftests/bpf/progs/arena_race.c

diff --git a/tools/testing/selftests/bpf/prog_tests/arena_race.c b/tools/testing/selftests/bpf/prog_tests/arena_race.c
new file mode 100644
index 000000000000..c2da611421f3
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/arena_race.c
@@ -0,0 +1,270 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#define _GNU_SOURCE
+#include <sched.h>
+#include <sys/syscall.h>
+#include <test_progs.h>
+
+#include "arena_race.skel.h"
+
+struct free_thread_ctx {
+	struct arena_race *skel;
+	int err;
+	__u32 retval;
+};
+
+struct fault_thread_ctx {
+	__u64 *addr;
+	int stop;
+};
+
+static int run_prog(struct bpf_program *prog, const char *name)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	int err;
+
+	err = bpf_prog_test_run_opts(bpf_program__fd(prog), &opts);
+	return ASSERT_OK(err, name) && ASSERT_OK(opts.retval, name) ? 0 : -1;
+}
+
+/* Trigger the sleepable free page path. */
+static void *run_free_thread(void *arg)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	struct free_thread_ctx *ctx = arg;
+
+	ctx->skel->bss->target_tid = sys_gettid();
+	ctx->err = bpf_prog_test_run_opts(
+		bpf_program__fd(ctx->skel->progs.free_page), &opts);
+	ctx->retval = opts.retval;
+	return NULL;
+}
+
+/* Continuously fault in the address. */
+static void *fault_reader_thread(void *arg)
+{
+	struct fault_thread_ctx *ctx = arg;
+
+	while (!READ_ONCE(ctx->stop))
+		(void)READ_ONCE(*ctx->addr);
+	return NULL;
+}
+
+static int wait_for(int *p)
+{
+	__u64 deadline = get_time_ns() + 5ULL * 1000 * 1000 * 1000;
+
+	while (!READ_ONCE(*p)) {
+		if (get_time_ns() > deadline)
+			return -ETIMEDOUT;
+	}
+	return 0;
+}
+
+static struct arena_race *setup_arena(__u64 **addr, bool trace_flush)
+{
+	struct arena_race *skel;
+	size_t arena_sz;
+	char *base;
+	int err;
+
+	skel = arena_race__open();
+	if (!ASSERT_OK_PTR(skel, "open"))
+		return NULL;
+	err = bpf_program__set_autoload(skel->progs.trace_flush, trace_flush);
+	if (!ASSERT_OK(err, "trace_flush_autoload"))
+		goto err_out;
+
+	err = arena_race__load(skel);
+	if (!ASSERT_OK(err, "load"))
+		goto err_out;
+	err = arena_race__attach(skel);
+	if (!ASSERT_OK(err, "attach"))
+		goto err_out;
+
+	if (run_prog(skel->progs.alloc_old, "alloc_old"))
+		goto err_out;
+	if (skel->bss->skip) {
+		test__skip();
+		goto err_out;
+	}
+
+	base = bpf_map__initial_value(skel->maps.arena, &arena_sz);
+	if (!ASSERT_OK_PTR(base, "arena_base"))
+		goto err_out;
+	*addr = (__u64 *)(base + getpagesize());
+	if (!ASSERT_EQ((unsigned long)skel->bss->ptr, (unsigned long)*addr,
+		       "arena_ptr"))
+		goto err_out;
+	return skel;
+
+err_out:
+	arena_race__destroy(skel);
+	return NULL;
+}
+
+static void test_free_before_flush(bool deferred)
+{
+	struct free_thread_ctx ctx = {};
+	struct arena_race *skel;
+	pthread_t thread;
+	__u64 *addr;
+	bool thread_created = false, flush_seen = false, completed = false;
+	int err;
+
+	if (libbpf_find_vmlinux_btf_id("flush_tlb_kernel_range",
+					BPF_TRACE_FENTRY) <= 0) {
+		printf("%s:SKIP: flush_tlb_kernel_range is not an fentry target\n",
+		       __func__);
+		test__skip();
+		return;
+	}
+
+	skel = setup_arena(&addr, true);
+	if (!skel)
+		return;
+
+	/* Pause during a TLB flush to widen the race window. */
+	skel->bss->pause_on_flush = 1;
+
+	if (deferred) {
+		/*
+		 * Test the nonsleepable free path that gets
+		 * deferred to a worker in the kernel. We do
+		 * so by triggering the arena operation from
+		 * a nonsleepable tracepoint context.
+		 */
+		skel->bss->trigger_pid_tgid =
+			((__u64)getpid() << 32) | (__u32)sys_gettid();
+		skel->bss->trigger_syscall = SYS_getpgid;
+		skel->bss->deferred_free = 1;
+		if (!ASSERT_GE(syscall(SYS_getpgid, 0), 0, "deferred_free"))
+			goto release;
+	} else {
+		/*
+		 * Test the sleepable arena free path through a
+		 * syscall test prog.
+		 */
+		ctx.skel = skel;
+		err = pthread_create(&thread, NULL, run_free_thread, &ctx);
+		if (!ASSERT_OK(err, "pthread_create")) {
+			skel->bss->release = 1;
+			goto out;
+		}
+		thread_created = true;
+	}
+
+	/* Wait until the worker thread triggers a flush. */
+	err = wait_for(&skel->bss->flush_entered);
+	if (!ASSERT_OK(err, "flush_entered"))
+		goto release;
+
+	flush_seen = true;
+
+	/* Force a reallocation during the flush. */
+	run_prog(skel->progs.try_realloc, "realloc_before_flush");
+	ASSERT_NULL(skel->bss->realloc_ptr, "realloc_before_flush");
+
+release:
+	skel->bss->release = 1;
+	if (thread_created) {
+		ASSERT_OK(pthread_join(thread, NULL), "pthread_join");
+		thread_created = false;
+		completed = ASSERT_OK(ctx.err, "free_run") &&
+			    ASSERT_OK(ctx.retval, "free_retval");
+	} else if (skel->bss->target_tid) {
+		err = wait_for(&skel->bss->worker_exited);
+		completed = ASSERT_OK(err, "worker_exited");
+	}
+	ASSERT_FALSE(skel->bss->timed_out, "flush_timed_out");
+
+	if (flush_seen && completed &&
+	    !run_prog(skel->progs.try_realloc, "realloc_after_flush")) {
+		ASSERT_EQ((unsigned long)skel->bss->realloc_ptr,
+			  (unsigned long)addr, "realloc_after_flush");
+		ASSERT_EQ(*addr, skel->rodata->new_marker, "new_marker");
+	}
+out:
+	arena_race__destroy(skel);
+}
+
+/*
+ * Force a race between a faulting thread in userspace and a
+ * free operation on the arena.
+ */
+static void test_fault_free_realloc(void)
+{
+	struct fault_thread_ctx fault = {};
+	struct arena_race *skel;
+	pthread_t fault_thread;
+	bool fault_created = false;
+	__u64 *addr;
+	__u64 expected, value;
+	int err, i;
+
+	skel = setup_arena(&addr, false);
+	if (!skel)
+		return;
+
+	skel->bss->realloc_after_free = 1;
+
+	fault.addr = addr;
+	err = pthread_create(&fault_thread, NULL, fault_reader_thread, &fault);
+	if (!ASSERT_OK(err, "pthread_create_fault"))
+		goto out;
+	fault_created = true;
+
+	for (i = 0; i < 1000; i++) {
+		LIBBPF_OPTS(bpf_test_run_opts, opts);
+
+		err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.free_page),
+					     &opts);
+		if (err) {
+			ASSERT_OK(err, "free_realloc");
+			break;
+		}
+		if (opts.retval) {
+			ASSERT_OK(opts.retval, "free_realloc");
+			break;
+		}
+		value = *addr;
+		expected = skel->bss->current_marker;
+		if (value != expected) {
+			ASSERT_EQ(value, expected, "marker_after_realloc");
+			break;
+		}
+	}
+
+	WRITE_ONCE(fault.stop, 1);
+	if (fault_created) {
+		ASSERT_OK(pthread_join(fault_thread, NULL), "pthread_join_fault");
+		fault_created = false;
+	}
+out:
+	WRITE_ONCE(fault.stop, 1);
+	if (fault_created)
+		pthread_join(fault_thread, NULL);
+	arena_race__destroy(skel);
+}
+
+void serial_test_arena_race(void)
+{
+	cpu_set_t cpuset;
+	int err;
+
+	err = sched_getaffinity(0, sizeof(cpuset), &cpuset);
+	if (!ASSERT_OK(err, "sched_getaffinity"))
+		return;
+	if (CPU_COUNT(&cpuset) < 2) {
+		printf("%s:SKIP: at least two runnable CPUs are required\n", __func__);
+		test__skip();
+		return;
+	}
+
+	if (test__start_subtest("free_before_flush"))
+		test_free_before_flush(false);
+	if (test__start_subtest("deferred_free_before_flush"))
+		test_free_before_flush(true);
+	if (test__start_subtest("fault_free_realloc"))
+		test_fault_free_realloc();
+}
diff --git a/tools/testing/selftests/bpf/progs/arena_race.c b/tools/testing/selftests/bpf/progs/arena_race.c
new file mode 100644
index 000000000000..56ebfd727324
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/arena_race.c
@@ -0,0 +1,159 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#define BPF_NO_KFUNC_PROTOTYPES
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_experimental.h"
+#include <bpf_arena_common.h>
+
+const volatile __u64 old_marker = 0x1111222233334444ULL;
+const volatile __u64 new_marker = 0x5555666677778888ULL;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARENA);
+	__uint(map_flags, BPF_F_MMAPABLE);
+	__uint(max_entries, 4);
+#ifdef __TARGET_ARCH_arm64
+	__ulong(map_extra, 0x1ull << 32);
+#else
+	__ulong(map_extra, 0x1ull << 44);
+#endif
+} arena SEC(".maps");
+
+bool skip;
+
+void __arena *ptr;
+void __arena *realloc_ptr;
+bool realloc_after_free;
+__u64 current_marker;
+__u64 marker_seq;
+
+int target_tid;
+int pause_on_flush;
+int flush_entered;
+int release;
+int timed_out;
+
+__u64 trigger_pid_tgid;
+long trigger_syscall;
+int deferred_free;
+int worker_armed;
+int worker_exited;
+
+static __always_inline void wait_for_release(void)
+{
+	while (!*(volatile int *)&release && can_loop)
+		;
+	if (!*(volatile int *)&release)
+		timed_out = 1;
+}
+
+SEC("syscall")
+int alloc_old(void *ctx)
+{
+#if defined(__BPF_FEATURE_ADDR_SPACE_CAST) || defined(BPF_ARENA_FORCE_ASM)
+	__u64 __arena *p;
+	char __arena *base = arena_base(&arena);
+
+	realloc_ptr = NULL;
+	ptr = bpf_arena_alloc_pages(&arena, base + __PAGE_SIZE, 1,
+				    NUMA_NO_NODE, 0);
+	if (!ptr)
+		return 1;
+	p = (__u64 __arena *)ptr;
+	*p = old_marker;
+	current_marker = old_marker;
+#else
+	skip = true;
+#endif
+	return 0;
+}
+
+SEC("syscall")
+int free_page(void *ctx)
+{
+#if defined(__BPF_FEATURE_ADDR_SPACE_CAST) || defined(BPF_ARENA_FORCE_ASM)
+	__u64 __arena *p;
+	__u64 marker;
+
+	if (!ptr)
+		return 1;
+	bpf_arena_free_pages(&arena, ptr, 1);
+	if (!realloc_after_free)
+		return 0;
+
+	marker = new_marker + ++marker_seq;
+	realloc_ptr = bpf_arena_alloc_pages(&arena, ptr, 1, NUMA_NO_NODE, 0);
+	if (realloc_ptr)
+		ptr = realloc_ptr;
+	else
+		realloc_ptr = ptr;
+	p = (__u64 __arena *)ptr;
+	*p = marker;
+	current_marker = marker;
+#endif
+	return 0;
+}
+
+SEC("syscall")
+int try_realloc(void *ctx)
+{
+#if defined(__BPF_FEATURE_ADDR_SPACE_CAST) || defined(BPF_ARENA_FORCE_ASM)
+	__u64 __arena *p;
+
+	realloc_ptr = bpf_arena_alloc_pages(&arena, ptr, 1, NUMA_NO_NODE, 0);
+	if (realloc_ptr) {
+		ptr = realloc_ptr;
+		p = (__u64 __arena *)realloc_ptr;
+		*p = new_marker;
+		current_marker = new_marker;
+	}
+#endif
+	return 0;
+}
+
+SEC("tp_btf/sys_enter")
+int BPF_PROG(deferred_free_prog, struct pt_regs *regs, long id)
+{
+#if defined(__BPF_FEATURE_ADDR_SPACE_CAST) || defined(BPF_ARENA_FORCE_ASM)
+	if (!deferred_free || bpf_get_current_pid_tgid() != trigger_pid_tgid ||
+	    id != trigger_syscall)
+		return 0;
+
+	deferred_free = 0;
+	/* The worker can run on another CPU before the kfunc returns. */
+	worker_armed = 1;
+	bpf_arena_free_pages(&arena, ptr, 1);
+#endif
+	return 0;
+}
+
+SEC("fentry/arena_free_worker")
+int BPF_PROG(trace_free_worker, struct work_struct *work)
+{
+	if (worker_armed && !target_tid)
+		target_tid = (__u32)bpf_get_current_pid_tgid();
+	return 0;
+}
+
+SEC("fexit/arena_free_worker")
+int BPF_PROG(trace_free_worker_ret, struct work_struct *work)
+{
+	if ((__u32)bpf_get_current_pid_tgid() == target_tid)
+		worker_exited = 1;
+	return 0;
+}
+
+SEC("?fentry/flush_tlb_kernel_range")
+int BPF_PROG(trace_flush, unsigned long start, unsigned long end)
+{
+	if (!pause_on_flush ||
+	    (__u32)bpf_get_current_pid_tgid() != target_tid)
+		return 0;
+	flush_entered = 1;
+	wait_for_release();
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.54.0