Commit 4ac45468 authored by Dave Marchevsky's avatar Dave Marchevsky Committed by Andrii Nakryiko
Browse files

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: default avatarDave Marchevsky <davemarchevsky@fb.com>
Signed-off-by: default avatarAndrii Nakryiko <andrii@kernel.org>
Acked-by: default avatarAndrii Nakryiko <andrii@kernel.org>
Link: https://lore.kernel.org/bpf/20231013204426.1074286-4-davemarchevsky@fb.com
parent 45b38941

kernel/bpf/helpers.c

+3 −0
Original line number Diff line number Diff line
@@ -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)

kernel/bpf/task_iter.c

+91 −0
Original line number Diff line number Diff line
@@ -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)