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