In some profiler use cases, it is necessary to map an address to the
backing file, e.g., a shared library. bpf_find_vma helper provides a
flexible way to achieve this. bpf_find_vma maps an address of a task to
the vma (vm_area_struct) for this address, and feed the vma to an callback
BPF function. The callback function is necessary here, as we need to
ensure mmap_sem is unlocked.

It is necessary to lock mmap_sem for find_vma. To lock and unlock mmap_sem
safely when irqs are disable, we use the same mechanism as stackmap with
build_id. Specifically, when irqs are disabled, the unlocked is postponed
in an irq_work. Refactor stackmap.c so that the irq_work is shared among
bpf_find_vma and stackmap helpers.

Signed-off-by: Song Liu <redacted>
---
  include/linux/bpf.h            |  1 +
  include/uapi/linux/bpf.h       | 20 ++++++++++
  kernel/bpf/mmap_unlock_work.h  | 65 +++++++++++++++++++++++++++++++
  kernel/bpf/stackmap.c          | 71 ++++++++--------------------------
  kernel/bpf/task_iter.c         | 49 +++++++++++++++++++++++
  kernel/bpf/verifier.c          | 36 +++++++++++++++++
  kernel/trace/bpf_trace.c       |  2 +
  tools/include/uapi/linux/bpf.h | 20 ++++++++++
  8 files changed, 209 insertions(+), 55 deletions(-)
  create mode 100644 kernel/bpf/mmap_unlock_work.h

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 2be6dfd68df99..df3410bff4b06 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h

@@ -2157,6 +2157,7 @@ extern const struct bpf_func_proto bpf_btf_find_by_name_kind_proto;
  extern const struct bpf_func_proto bpf_sk_setsockopt_proto;
  extern const struct bpf_func_proto bpf_sk_getsockopt_proto;
  extern const struct bpf_func_proto bpf_kallsyms_lookup_name_proto;
+extern const struct bpf_func_proto bpf_find_vma_proto;
  
  const struct bpf_func_proto *tracing_prog_func_proto(
    enum bpf_func_id func_id, const struct bpf_prog *prog);

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index ba5af15e25f5c..22fa7b74de451 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h

@@ -4938,6 +4938,25 @@ union bpf_attr {
   *		**-ENOENT** if symbol is not found.
   *
   *		**-EPERM** if caller does not have permission to obtain kernel address.
+ *
+ * long bpf_find_vma(struct task_struct *task, u64 addr, void *callback_fn, void *callback_ctx, u64 flags)
+ *	Description
+ *		Find vma of *task* that contains *addr*, call *callback_fn*
+ *		function with *task*, *vma*, and *callback_ctx*.
+ *		The *callback_fn* should be a static function and
+ *		the *callback_ctx* should be a pointer to the stack.
+ *		The *flags* is used to control certain aspects of the helper.
+ *		Currently, the *flags* must be 0.
+ *
+ *		The expected callback signature is
+ *
+ *		long (\*callback_fn)(struct task_struct \*task, struct vm_area_struct \*vma, void \*ctx);

+ *
+ *	Return
+ *		0 on success.
+ *		**-ENOENT** if *task->mm* is NULL, or no vma contains *addr*.
+ *		**-EBUSY** if failed to try lock mmap_lock.
+ *		**-EINVAL** for invalid **flags**.
   */
  #define __BPF_FUNC_MAPPER(FN)		\
  	FN(unspec),			\

@@ -5120,6 +5139,7 @@ union bpf_attr {
  	FN(trace_vprintk),		\
  	FN(skc_to_unix_sock),		\
  	FN(kallsyms_lookup_name),	\
+	FN(find_vma),			\
  	/* */
  
  /* integer value in 'imm' field of BPF_CALL instruction selects which helper

+
+const struct bpf_func_proto bpf_find_vma_proto = {
+	.func		= bpf_find_vma,
+	.ret_type	= RET_INTEGER,
+	.arg1_type	= ARG_PTR_TO_BTF_ID,
+	.arg1_btf_id	= &btf_find_vma_ids[0],
+	.arg2_type	= ARG_ANYTHING,
+	.arg3_type	= ARG_PTR_TO_FUNC,
+	.arg4_type	= ARG_PTR_TO_STACK_OR_NULL,
+	.arg5_type	= ARG_ANYTHING,
+};
+
  static int __init task_iter_init(void)
  {
  	int ret;

diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index f0dca726ebfde..a65526112924a 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c

@@ -6132,6 +6132,35 @@ static int set_timer_callback_state(struct bpf_verifier_env *env,
  	return 0;
  }
  
+BTF_ID_LIST_SINGLE(btf_set_find_vma_ids, struct, vm_area_struct)