bpf: Introduce task_vma open-coded iterator kfuncs

This patch adds kfuncs bpf_iter_task_vma_{new,next,destroy} which allow
creation and manipulation of struct bpf_iter_task_vma in open-coded
iterator style. BPF programs can use these kfuncs directly or through
bpf_for_each macro for natural-looking iteration of all task vmas.

The implementation borrows heavily from bpf_find_vma helper's locking -
differing only in that it holds the mmap_read lock for all iterations
while the helper only executes its provided callback on a maximum of 1
vma. Aside from locking, struct vma_iterator and vma_next do all the
heavy lifting.

A pointer to an inner data struct, struct bpf_iter_task_vma_data, is the
only field in struct bpf_iter_task_vma. This is because the inner data
struct contains a struct vma_iterator (not ptr), whose size is likely to
change under us. If bpf_iter_task_vma_kern contained vma_iterator directly
such a change would require change in opaque bpf_iter_task_vma struct's
size. So better to allocate vma_iterator using BPF allocator, and since
that alloc must already succeed, might as well allocate all iter fields,
thereby freezing struct bpf_iter_task_vma size.

Signed-off-by: Dave Marchevsky <davemarchevsky@fb.com>
Signed-off-by: Andrii Nakryiko <andrii@kernel.org>
Acked-by: Andrii Nakryiko <andrii@kernel.org>
Link: https://lore.kernel.org/bpf/20231013204426.1074286-4-davemarchevsky@fb.com

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index d2840dd5b00d71660593ccd9a0df068f1f4d127f..62a53ebfedf9c08619f4474ee74a8844ce520164 100644 (file)
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -2552,6 +2552,9 @@ BTF_ID_FLAGS(func, bpf_dynptr_slice_rdwr, KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_iter_num_new, KF_ITER_NEW)
 BTF_ID_FLAGS(func, bpf_iter_num_next, KF_ITER_NEXT | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_iter_num_destroy, KF_ITER_DESTROY)
+BTF_ID_FLAGS(func, bpf_iter_task_vma_new, KF_ITER_NEW | KF_RCU)
+BTF_ID_FLAGS(func, bpf_iter_task_vma_next, KF_ITER_NEXT | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_iter_task_vma_destroy, KF_ITER_DESTROY)
 BTF_ID_FLAGS(func, bpf_dynptr_adjust)
 BTF_ID_FLAGS(func, bpf_dynptr_is_null)
 BTF_ID_FLAGS(func, bpf_dynptr_is_rdonly)

diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
index 7473068ed313e4b4cd7dde4c0b2b5dabe18acc46..fef17628341fa1c3fd8250aaa7ba585cfc738252 100644 (file)
--- a/kernel/bpf/task_iter.c
+++ b/kernel/bpf/task_iter.c
@@ -7,7 +7,9 @@
 #include <linux/fs.h>
 #include <linux/fdtable.h>
 #include <linux/filter.h>
+#include <linux/bpf_mem_alloc.h>
 #include <linux/btf_ids.h>
+#include <linux/mm_types.h>
 #include "mmap_unlock_work.h"
 
 static const char * const iter_task_type_names[] = {
@@ -803,6 +805,95 @@ const struct bpf_func_proto bpf_find_vma_proto = {
        .arg5_type      = ARG_ANYTHING,
 };
 
+struct bpf_iter_task_vma_kern_data {
+       struct task_struct *task;
+       struct mm_struct *mm;
+       struct mmap_unlock_irq_work *work;
+       struct vma_iterator vmi;
+};
+
+struct bpf_iter_task_vma {
+       /* opaque iterator state; having __u64 here allows to preserve correct
+        * alignment requirements in vmlinux.h, generated from BTF
+        */
+       __u64 __opaque[1];
+} __attribute__((aligned(8)));
+
+/* Non-opaque version of bpf_iter_task_vma */
+struct bpf_iter_task_vma_kern {
+       struct bpf_iter_task_vma_kern_data *data;
+} __attribute__((aligned(8)));
+
+__diag_push();
+__diag_ignore_all("-Wmissing-prototypes",
+                 "Global functions as their definitions will be in vmlinux BTF");
+
+__bpf_kfunc int bpf_iter_task_vma_new(struct bpf_iter_task_vma *it,
+                                     struct task_struct *task, u64 addr)
+{
+       struct bpf_iter_task_vma_kern *kit = (void *)it;
+       bool irq_work_busy = false;
+       int err;
+
+       BUILD_BUG_ON(sizeof(struct bpf_iter_task_vma_kern) != sizeof(struct bpf_iter_task_vma));
+       BUILD_BUG_ON(__alignof__(struct bpf_iter_task_vma_kern) != __alignof__(struct bpf_iter_task_vma));
+
+       /* is_iter_reg_valid_uninit guarantees that kit hasn't been initialized
+        * before, so non-NULL kit->data doesn't point to previously
+        * bpf_mem_alloc'd bpf_iter_task_vma_kern_data
+        */
+       kit->data = bpf_mem_alloc(&bpf_global_ma, sizeof(struct bpf_iter_task_vma_kern_data));
+       if (!kit->data)
+               return -ENOMEM;
+
+       kit->data->task = get_task_struct(task);
+       kit->data->mm = task->mm;
+       if (!kit->data->mm) {
+               err = -ENOENT;
+               goto err_cleanup_iter;
+       }
+
+       /* kit->data->work == NULL is valid after bpf_mmap_unlock_get_irq_work */
+       irq_work_busy = bpf_mmap_unlock_get_irq_work(&kit->data->work);
+       if (irq_work_busy || !mmap_read_trylock(kit->data->mm)) {
+               err = -EBUSY;
+               goto err_cleanup_iter;
+       }
+
+       vma_iter_init(&kit->data->vmi, kit->data->mm, addr);
+       return 0;
+
+err_cleanup_iter:
+       if (kit->data->task)
+               put_task_struct(kit->data->task);
+       bpf_mem_free(&bpf_global_ma, kit->data);
+       /* NULL kit->data signals failed bpf_iter_task_vma initialization */
+       kit->data = NULL;
+       return err;
+}
+
+__bpf_kfunc struct vm_area_struct *bpf_iter_task_vma_next(struct bpf_iter_task_vma *it)
+{
+       struct bpf_iter_task_vma_kern *kit = (void *)it;
+
+       if (!kit->data) /* bpf_iter_task_vma_new failed */
+               return NULL;
+       return vma_next(&kit->data->vmi);
+}
+
+__bpf_kfunc void bpf_iter_task_vma_destroy(struct bpf_iter_task_vma *it)
+{
+       struct bpf_iter_task_vma_kern *kit = (void *)it;
+
+       if (kit->data) {
+               bpf_mmap_unlock_mm(kit->data->work, kit->data->mm);
+               put_task_struct(kit->data->task);
+               bpf_mem_free(&bpf_global_ma, kit->data);
+       }
+}
+
+__diag_pop();
+
 DEFINE_PER_CPU(struct mmap_unlock_irq_work, mmap_unlock_work);
 
 static void do_mmap_read_unlock(struct irq_work *entry)