[PATCH RFCv2 bpf-next 1/4] bpf: Add support for kprobe multi wrapper attach

[Jiri Olsa <jolsa@xxxxxxxxxx>]

Adding support to attach bpf program for entry and return probe
of the same function. This is common usecase and at the moment
it requires to create two kprobe multi links.

Adding new attr.link_create.kprobe_multi.flags value that instructs
kernel to attach link program to both entry and exit probe.

It's possible to control execution of the bpf program on return
probe simply by returning zero or non zero from the entry bpf
program execution to execute or not respectively the bpf program
on return probe.

Signed-off-by: Jiri Olsa <jolsa@xxxxxxxxxx>
---
 include/uapi/linux/bpf.h       |  3 ++-
 kernel/trace/bpf_trace.c       | 24 ++++++++++++++++++------
 tools/include/uapi/linux/bpf.h |  3 ++-
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index d2e6c5fcec01..a430855c5bcd 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1247,7 +1247,8 @@ enum bpf_perf_event_type {
  * BPF_TRACE_KPROBE_MULTI attach type to create return probe.
  */
 enum {
-	BPF_F_KPROBE_MULTI_RETURN = (1U << 0)
+	BPF_F_KPROBE_MULTI_RETURN  = (1U << 0),
+	BPF_F_KPROBE_MULTI_WRAPPER = (1U << 1),
 };
 
 /* link_create.uprobe_multi.flags used in LINK_CREATE command for
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 241ddf5e3895..726a8c71f0da 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -2587,6 +2587,7 @@ struct bpf_kprobe_multi_link {
 	u32 mods_cnt;
 	struct module **mods;
 	u32 flags;
+	bool is_wrapper;
 };
 
 struct bpf_kprobe_multi_run_ctx {
@@ -2826,10 +2827,11 @@ kprobe_multi_link_handler(struct fprobe *fp, unsigned long fentry_ip,
 			  void *data)
 {
 	struct bpf_kprobe_multi_link *link;
+	int err;
 
 	link = container_of(fp, struct bpf_kprobe_multi_link, fp);
-	kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
-	return 0;
+	err = kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
+	return link->is_wrapper ? err : 0;
 }
 
 static void
@@ -2967,6 +2969,7 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
 	void __user *uaddrs;
 	u64 *cookies = NULL;
 	void __user *usyms;
+	bool is_wrapper;
 	int err;
 
 	/* no support for 32bit archs yet */
@@ -2977,9 +2980,12 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
 		return -EINVAL;
 
 	flags = attr->link_create.kprobe_multi.flags;
-	if (flags & ~BPF_F_KPROBE_MULTI_RETURN)
+	if (flags & ~(BPF_F_KPROBE_MULTI_RETURN|
+		      BPF_F_KPROBE_MULTI_WRAPPER))
 		return -EINVAL;
 
+	is_wrapper = flags & BPF_F_KPROBE_MULTI_WRAPPER;
+
 	uaddrs = u64_to_user_ptr(attr->link_create.kprobe_multi.addrs);
 	usyms = u64_to_user_ptr(attr->link_create.kprobe_multi.syms);
 	if (!!uaddrs == !!usyms)
@@ -3054,15 +3060,21 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
 	if (err)
 		goto error;
 
-	if (flags & BPF_F_KPROBE_MULTI_RETURN)
-		link->fp.exit_handler = kprobe_multi_link_exit_handler;
-	else
+	if (is_wrapper) {
 		link->fp.entry_handler = kprobe_multi_link_handler;
+		link->fp.exit_handler = kprobe_multi_link_exit_handler;
+	} else {
+		if (flags & BPF_F_KPROBE_MULTI_RETURN)
+			link->fp.exit_handler = kprobe_multi_link_exit_handler;
+		else
+			link->fp.entry_handler = kprobe_multi_link_handler;
+	}
 
 	link->addrs = addrs;
 	link->cookies = cookies;
 	link->cnt = cnt;
 	link->flags = flags;
+	link->is_wrapper = is_wrapper;
 
 	if (cookies) {
 		/*
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index d2e6c5fcec01..a430855c5bcd 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -1247,7 +1247,8 @@ enum bpf_perf_event_type {
  * BPF_TRACE_KPROBE_MULTI attach type to create return probe.
  */
 enum {
-	BPF_F_KPROBE_MULTI_RETURN = (1U << 0)
+	BPF_F_KPROBE_MULTI_RETURN  = (1U << 0),
+	BPF_F_KPROBE_MULTI_WRAPPER = (1U << 1),
 };
 
 /* link_create.uprobe_multi.flags used in LINK_CREATE command for
-- 
2.43.2