[PATCHv2 bpf-next 0/9] uprobe, bpf: Add session support

[PATCHv2 bpf-next 0/9] uprobe, bpf: Add session support

[Jiri Olsa]

hi,
this patchset is adding support for session uprobe attachment
and using it through bpf link for bpf programs.

The session means that the uprobe consumer is executed on entry
and return of probed function with additional control:
  - entry callback can control execution of the return callback
  - entry and return callbacks can share data/cookie

On more details please see patch #1.

v2 changes:
  - re-implement uprobe session support [Andrii]
  - added test for mixed uprobe consumers

thanks,
jirka


---
Jiri Olsa (9):
      uprobe: Add support for session consumer
      bpf: Add support for uprobe multi session attach
      bpf: Add support for uprobe multi session context
      libbpf: Add support for uprobe multi session attach
      libbpf: Add uprobe session attach type names to attach_type_name
      selftests/bpf: Add uprobe session test
      selftests/bpf: Add uprobe session cookie test
      selftests/bpf: Add uprobe session recursive test
      selftests/bpf: Add uprobe session consumers test

 include/linux/uprobes.h                                            |  16 +++-
 include/uapi/linux/bpf.h                                           |   1 +
 kernel/bpf/syscall.c                                               |   9 ++-
 kernel/events/uprobes.c                                            | 129 ++++++++++++++++++++++++++++++---
 kernel/trace/bpf_trace.c                                           |  54 ++++++++++----
 kernel/trace/trace_uprobe.c                                        |  12 ++-
 tools/include/uapi/linux/bpf.h                                     |   1 +
 tools/lib/bpf/bpf.c                                                |   1 +
 tools/lib/bpf/libbpf.c                                             |  51 ++++++++++++-
 tools/lib/bpf/libbpf.h                                             |   4 +-
 tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c         | 333 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session.c           |  53 ++++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c |  53 ++++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c    |  48 ++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c |  44 +++++++++++
 15 files changed, 771 insertions(+), 38 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session.c
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c

[PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

Adding support for uprobe consumer to be defined as session and have
new behaviour for consumer's 'handler' and 'ret_handler' callbacks.

The session means that 'handler' and 'ret_handler' callbacks are
connected in a way that allows to:

  - control execution of 'ret_handler' from 'handler' callback
  - share data between 'handler' and 'ret_handler' callbacks

The session is enabled by setting new 'session' bool field to true
in uprobe_consumer object.

We keep count of session consumers for uprobe and allocate session_consumer
object for each in return_instance object. This allows us to store
return values of 'handler' callbacks and data pointers of shared
data between both handlers.

The session concept fits to our common use case where we do filtering
on entry uprobe and based on the result we decide to run the return
uprobe (or not).

It's also convenient to share the data between session callbacks.

The control of 'ret_handler' callback execution is done via return
value of the 'handler' callback. If it's 0 we install and execute
return uprobe, if it's 1 we do not.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/uprobes.h     |  16 ++++-
 kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
 kernel/trace/bpf_trace.c    |   6 +-
 kernel/trace/trace_uprobe.c |  12 ++--
 4 files changed, 144 insertions(+), 19 deletions(-)

diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
index f46e0ca0169c..903a860a8d01 100644
--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h
@@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
 };
 
 struct uprobe_consumer {
-	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
+	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
 	int (*ret_handler)(struct uprobe_consumer *self,
 				unsigned long func,
-				struct pt_regs *regs);
+				struct pt_regs *regs, __u64 *data);
 	bool (*filter)(struct uprobe_consumer *self,
 				enum uprobe_filter_ctx ctx,
 				struct mm_struct *mm);
 
 	struct uprobe_consumer *next;
+
+	bool			session;	/* marks uprobe session consumer */
+	unsigned int		session_id;	/* set when uprobe_consumer is registered */
 };
 
 #ifdef CONFIG_UPROBES
@@ -80,6 +83,12 @@ struct uprobe_task {
 	unsigned int			depth;
 };
 
+struct session_consumer {
+	__u64		cookie;
+	unsigned int	id;
+	int		rc;
+};
+
 struct return_instance {
 	struct uprobe		*uprobe;
 	unsigned long		func;
@@ -88,6 +97,9 @@ struct return_instance {
 	bool			chained;	/* true, if instance is nested */
 
 	struct return_instance	*next;		/* keep as stack */
+
+	int			sessions_cnt;
+	struct session_consumer	sessions[];
 };
 
 enum rp_check {
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index 2c83ba776fc7..4da410460f2a 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -63,6 +63,8 @@ struct uprobe {
 	loff_t			ref_ctr_offset;
 	unsigned long		flags;
 
+	unsigned int		sessions_cnt;
+
 	/*
 	 * The generic code assumes that it has two members of unknown type
 	 * owned by the arch-specific code:
@@ -750,11 +752,30 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
 	return uprobe;
 }
 
+static void
+uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
+{
+	static unsigned int session_id;
+
+	if (uc->session) {
+		uprobe->sessions_cnt++;
+		uc->session_id = ++session_id ?: ++session_id;
+	}
+}
+
+static void
+uprobe_consumer_unaccount(struct uprobe *uprobe, struct uprobe_consumer *uc)
+{
+	if (uc->session)
+		uprobe->sessions_cnt--;
+}
+
 static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
 {
 	down_write(&uprobe->consumer_rwsem);
 	uc->next = uprobe->consumers;
 	uprobe->consumers = uc;
+	uprobe_consumer_account(uprobe, uc);
 	up_write(&uprobe->consumer_rwsem);
 }
 
@@ -773,6 +794,7 @@ static bool consumer_del(struct uprobe *uprobe, struct uprobe_consumer *uc)
 		if (*con == uc) {
 			*con = uc->next;
 			ret = true;
+			uprobe_consumer_unaccount(uprobe, uc);
 			break;
 		}
 	}
@@ -1744,6 +1766,23 @@ static struct uprobe_task *get_utask(void)
 	return current->utask;
 }
 
+static size_t ri_size(int sessions_cnt)
+{
+	struct return_instance *ri __maybe_unused;
+
+	return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
+}
+
+static struct return_instance *alloc_return_instance(int sessions_cnt)
+{
+	struct return_instance *ri;
+
+	ri = kzalloc(ri_size(sessions_cnt), GFP_KERNEL);
+	if (ri)
+		ri->sessions_cnt = sessions_cnt;
+	return ri;
+}
+
 static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
 {
 	struct uprobe_task *n_utask;
@@ -1756,11 +1795,11 @@ static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
 
 	p = &n_utask->return_instances;
 	for (o = o_utask->return_instances; o; o = o->next) {
-		n = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
+		n = alloc_return_instance(o->sessions_cnt);
 		if (!n)
 			return -ENOMEM;
 
-		*n = *o;
+		memcpy(n, o, ri_size(o->sessions_cnt));
 		get_uprobe(n->uprobe);
 		n->next = NULL;
 
@@ -1853,9 +1892,9 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
 	utask->return_instances = ri;
 }
 
-static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
+static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
+			      struct return_instance *ri)
 {
-	struct return_instance *ri;
 	struct uprobe_task *utask;
 	unsigned long orig_ret_vaddr, trampoline_vaddr;
 	bool chained;
@@ -1874,9 +1913,11 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 		return;
 	}
 
-	ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
-	if (!ri)
-		return;
+	if (!ri) {
+		ri = alloc_return_instance(0);
+		if (!ri)
+			return;
+	}
 
 	trampoline_vaddr = get_trampoline_vaddr();
 	orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
@@ -2065,35 +2106,85 @@ static struct uprobe *find_active_uprobe(unsigned long bp_vaddr, int *is_swbp)
 	return uprobe;
 }
 
+static struct session_consumer *
+session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
+		      int session_id)
+{
+	struct session_consumer *next;
+
+	next = sc ? sc + 1 : &ri->sessions[0];
+	next->id = session_id;
+	return next;
+}
+
+static struct session_consumer *
+session_consumer_find(struct return_instance *ri, int *iter, int session_id)
+{
+	struct session_consumer *sc;
+	int idx = *iter;
+
+	for (sc = &ri->sessions[idx]; idx < ri->sessions_cnt; idx++, sc++) {
+		if (sc->id == session_id) {
+			*iter = idx + 1;
+			return sc;
+		}
+	}
+	return NULL;
+}
+
 static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
 {
 	struct uprobe_consumer *uc;
 	int remove = UPROBE_HANDLER_REMOVE;
+	struct session_consumer *sc = NULL;
+	struct return_instance *ri = NULL;
 	bool need_prep = false; /* prepare return uprobe, when needed */
 
 	down_read(&uprobe->register_rwsem);
+	if (uprobe->sessions_cnt) {
+		ri = alloc_return_instance(uprobe->sessions_cnt);
+		if (!ri)
+			goto out;
+	}
+
 	for (uc = uprobe->consumers; uc; uc = uc->next) {
+		__u64 *cookie = NULL;
 		int rc = 0;
 
+		if (uc->session) {
+			sc = session_consumer_next(ri, sc, uc->session_id);
+			cookie = &sc->cookie;
+		}
+
 		if (uc->handler) {
-			rc = uc->handler(uc, regs);
+			rc = uc->handler(uc, regs, cookie);
 			WARN(rc & ~UPROBE_HANDLER_MASK,
 				"bad rc=0x%x from %ps()\n", rc, uc->handler);
 		}
 
-		if (uc->ret_handler)
+		if (uc->session) {
+			sc->rc = rc;
+			need_prep |= !rc;
+		} else if (uc->ret_handler) {
 			need_prep = true;
+		}
 
 		remove &= rc;
 	}
 
+	/* no removal if there's at least one session consumer */
+	remove &= !uprobe->sessions_cnt;
+
 	if (need_prep && !remove)
-		prepare_uretprobe(uprobe, regs); /* put bp at return */
+		prepare_uretprobe(uprobe, regs, ri); /* put bp at return */
+	else
+		kfree(ri);
 
 	if (remove && uprobe->consumers) {
 		WARN_ON(!uprobe_is_active(uprobe));
 		unapply_uprobe(uprobe, current->mm);
 	}
+ out:
 	up_read(&uprobe->register_rwsem);
 }
 
@@ -2101,12 +2192,28 @@ static void
 handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
 {
 	struct uprobe *uprobe = ri->uprobe;
+	struct session_consumer *sc;
 	struct uprobe_consumer *uc;
+	int session_iter = 0;
 
 	down_read(&uprobe->register_rwsem);
 	for (uc = uprobe->consumers; uc; uc = uc->next) {
+		__u64 *cookie = NULL;
+
+		if (uc->session) {
+			/*
+			 * Consumers could be added and removed, but they will not
+			 * change position, so we can iterate sessions just once
+			 * and keep the last found session as base for next search.
+			 */
+			sc = session_consumer_find(ri, &session_iter, uc->session_id);
+			if (!sc || sc->rc)
+				continue;
+			cookie = &sc->cookie;
+		}
+
 		if (uc->ret_handler)
-			uc->ret_handler(uc, ri->func, regs);
+			uc->ret_handler(uc, ri->func, regs, cookie);
 	}
 	up_read(&uprobe->register_rwsem);
 }
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index cd098846e251..02d052639dfe 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3332,7 +3332,8 @@ uprobe_multi_link_filter(struct uprobe_consumer *con, enum uprobe_filter_ctx ctx
 }
 
 static int
-uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs)
+uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs,
+			  __u64 *data)
 {
 	struct bpf_uprobe *uprobe;
 
@@ -3341,7 +3342,8 @@ uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs)
 }
 
 static int
-uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, struct pt_regs *regs)
+uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, struct pt_regs *regs,
+			      __u64 *data)
 {
 	struct bpf_uprobe *uprobe;
 
diff --git a/kernel/trace/trace_uprobe.c b/kernel/trace/trace_uprobe.c
index c98e3b3386ba..7068c279a244 100644
--- a/kernel/trace/trace_uprobe.c
+++ b/kernel/trace/trace_uprobe.c
@@ -88,9 +88,11 @@ static struct trace_uprobe *to_trace_uprobe(struct dyn_event *ev)
 static int register_uprobe_event(struct trace_uprobe *tu);
 static int unregister_uprobe_event(struct trace_uprobe *tu);
 
-static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs);
+static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs,
+			     __u64 *data);
 static int uretprobe_dispatcher(struct uprobe_consumer *con,
-				unsigned long func, struct pt_regs *regs);
+				unsigned long func, struct pt_regs *regs,
+				__u64 *data);
 
 #ifdef CONFIG_STACK_GROWSUP
 static unsigned long adjust_stack_addr(unsigned long addr, unsigned int n)
@@ -1504,7 +1506,8 @@ trace_uprobe_register(struct trace_event_call *event, enum trace_reg type,
 	}
 }
 
-static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs)
+static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs,
+			     __u64 *data)
 {
 	struct trace_uprobe *tu;
 	struct uprobe_dispatch_data udd;
@@ -1534,7 +1537,8 @@ static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs)
 }
 
 static int uretprobe_dispatcher(struct uprobe_consumer *con,
-				unsigned long func, struct pt_regs *regs)
+				unsigned long func, struct pt_regs *regs,
+				__u64 *data)
 {
 	struct trace_uprobe *tu;
 	struct uprobe_dispatch_data udd;
-- 
2.45.2

[PATCHv2 bpf-next 2/9] bpf: Add support for uprobe multi session attach

[Jiri Olsa]

Adding support to attach bpf program for entry and return probe
of the same function. This is common use case which at the moment
requires to create two uprobe multi links.

Adding new BPF_TRACE_UPROBE_SESSION attach type that instructs
kernel to attach single 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 the bpf program on return
probe respectively.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/uapi/linux/bpf.h       |  1 +
 kernel/bpf/syscall.c           |  9 +++++++--
 kernel/trace/bpf_trace.c       | 25 +++++++++++++++++++------
 tools/include/uapi/linux/bpf.h |  1 +
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 35bcf52dbc65..1d93cb014884 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1116,6 +1116,7 @@ enum bpf_attach_type {
 	BPF_NETKIT_PRIMARY,
 	BPF_NETKIT_PEER,
 	BPF_TRACE_KPROBE_SESSION,
+	BPF_TRACE_UPROBE_SESSION,
 	__MAX_BPF_ATTACH_TYPE
 };
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 869265852d51..2a63a528fa3c 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -4049,10 +4049,14 @@ static int bpf_prog_attach_check_attach_type(const struct bpf_prog *prog,
 		if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI &&
 		    attach_type != BPF_TRACE_UPROBE_MULTI)
 			return -EINVAL;
+		if (prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION &&
+		    attach_type != BPF_TRACE_UPROBE_SESSION)
+			return -EINVAL;
 		if (attach_type != BPF_PERF_EVENT &&
 		    attach_type != BPF_TRACE_KPROBE_MULTI &&
 		    attach_type != BPF_TRACE_KPROBE_SESSION &&
-		    attach_type != BPF_TRACE_UPROBE_MULTI)
+		    attach_type != BPF_TRACE_UPROBE_MULTI &&
+		    attach_type != BPF_TRACE_UPROBE_SESSION)
 			return -EINVAL;
 		return 0;
 	case BPF_PROG_TYPE_SCHED_CLS:
@@ -5315,7 +5319,8 @@ static int link_create(union bpf_attr *attr, bpfptr_t uattr)
 		else if (attr->link_create.attach_type == BPF_TRACE_KPROBE_MULTI ||
 			 attr->link_create.attach_type == BPF_TRACE_KPROBE_SESSION)
 			ret = bpf_kprobe_multi_link_attach(attr, prog);
-		else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI)
+		else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI ||
+			 attr->link_create.attach_type == BPF_TRACE_UPROBE_SESSION)
 			ret = bpf_uprobe_multi_link_attach(attr, prog);
 		break;
 	default:
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 02d052639dfe..1b19c1cdb5e1 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -1645,6 +1645,17 @@ static inline bool is_kprobe_session(const struct bpf_prog *prog)
 	return prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION;
 }
 
+static inline bool is_uprobe_multi(const struct bpf_prog *prog)
+{
+	return prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI ||
+	       prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
+}
+
+static inline bool is_uprobe_session(const struct bpf_prog *prog)
+{
+	return prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
+}
+
 static const struct bpf_func_proto *
 kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 {
@@ -1662,13 +1673,13 @@ kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 	case BPF_FUNC_get_func_ip:
 		if (is_kprobe_multi(prog))
 			return &bpf_get_func_ip_proto_kprobe_multi;
-		if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
+		if (is_uprobe_multi(prog))
 			return &bpf_get_func_ip_proto_uprobe_multi;
 		return &bpf_get_func_ip_proto_kprobe;
 	case BPF_FUNC_get_attach_cookie:
 		if (is_kprobe_multi(prog))
 			return &bpf_get_attach_cookie_proto_kmulti;
-		if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
+		if (is_uprobe_multi(prog))
 			return &bpf_get_attach_cookie_proto_umulti;
 		return &bpf_get_attach_cookie_proto_trace;
 	default:
@@ -3387,7 +3398,7 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
 	if (sizeof(u64) != sizeof(void *))
 		return -EOPNOTSUPP;
 
-	if (prog->expected_attach_type != BPF_TRACE_UPROBE_MULTI)
+	if (!is_uprobe_multi(prog))
 		return -EINVAL;
 
 	flags = attr->link_create.uprobe_multi.flags;
@@ -3463,10 +3474,12 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
 
 		uprobes[i].link = link;
 
-		if (flags & BPF_F_UPROBE_MULTI_RETURN)
-			uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
-		else
+		if (!(flags & BPF_F_UPROBE_MULTI_RETURN))
 			uprobes[i].consumer.handler = uprobe_multi_link_handler;
+		if (flags & BPF_F_UPROBE_MULTI_RETURN || is_uprobe_session(prog))
+			uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
+		if (is_uprobe_session(prog))
+			uprobes[i].consumer.session = true;
 
 		if (pid)
 			uprobes[i].consumer.filter = uprobe_multi_link_filter;
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 35bcf52dbc65..1d93cb014884 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -1116,6 +1116,7 @@ enum bpf_attach_type {
 	BPF_NETKIT_PRIMARY,
 	BPF_NETKIT_PEER,
 	BPF_TRACE_KPROBE_SESSION,
+	BPF_TRACE_UPROBE_SESSION,
 	__MAX_BPF_ATTACH_TYPE
 };
 
-- 
2.45.2

[PATCHv2 bpf-next 3/9] bpf: Add support for uprobe multi session context

[Jiri Olsa]

Placing bpf_session_run_ctx layer in between bpf_run_ctx and
bpf_uprobe_multi_run_ctx, so the session data can be retrieved
from uprobe_multi link.

Plus granting session kfuncs access to uprobe session programs.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/trace/bpf_trace.c | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 1b19c1cdb5e1..d431b880ca11 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3184,7 +3184,7 @@ struct bpf_uprobe_multi_link {
 };
 
 struct bpf_uprobe_multi_run_ctx {
-	struct bpf_run_ctx run_ctx;
+	struct bpf_session_run_ctx session_ctx;
 	unsigned long entry_ip;
 	struct bpf_uprobe *uprobe;
 };
@@ -3297,10 +3297,15 @@ static const struct bpf_link_ops bpf_uprobe_multi_link_lops = {
 
 static int uprobe_prog_run(struct bpf_uprobe *uprobe,
 			   unsigned long entry_ip,
-			   struct pt_regs *regs)
+			   struct pt_regs *regs,
+			   bool is_return, void *data)
 {
 	struct bpf_uprobe_multi_link *link = uprobe->link;
 	struct bpf_uprobe_multi_run_ctx run_ctx = {
+		.session_ctx = {
+			.is_return = is_return,
+			.data = data,
+		},
 		.entry_ip = entry_ip,
 		.uprobe = uprobe,
 	};
@@ -3319,7 +3324,7 @@ static int uprobe_prog_run(struct bpf_uprobe *uprobe,
 
 	migrate_disable();
 
-	old_run_ctx = bpf_set_run_ctx(&run_ctx.run_ctx);
+	old_run_ctx = bpf_set_run_ctx(&run_ctx.session_ctx.run_ctx);
 	err = bpf_prog_run(link->link.prog, regs);
 	bpf_reset_run_ctx(old_run_ctx);
 
@@ -3349,7 +3354,7 @@ uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs,
 	struct bpf_uprobe *uprobe;
 
 	uprobe = container_of(con, struct bpf_uprobe, consumer);
-	return uprobe_prog_run(uprobe, instruction_pointer(regs), regs);
+	return uprobe_prog_run(uprobe, instruction_pointer(regs), regs, false, data);
 }
 
 static int
@@ -3359,14 +3364,15 @@ uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, s
 	struct bpf_uprobe *uprobe;
 
 	uprobe = container_of(con, struct bpf_uprobe, consumer);
-	return uprobe_prog_run(uprobe, func, regs);
+	return uprobe_prog_run(uprobe, func, regs, true, data);
 }
 
 static u64 bpf_uprobe_multi_entry_ip(struct bpf_run_ctx *ctx)
 {
 	struct bpf_uprobe_multi_run_ctx *run_ctx;
 
-	run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx, run_ctx);
+	run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx,
+			       session_ctx.run_ctx);
 	return run_ctx->entry_ip;
 }
 
@@ -3374,7 +3380,8 @@ static u64 bpf_uprobe_multi_cookie(struct bpf_run_ctx *ctx)
 {
 	struct bpf_uprobe_multi_run_ctx *run_ctx;
 
-	run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx, run_ctx);
+	run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx,
+			       session_ctx.run_ctx);
 	return run_ctx->uprobe->cookie;
 }
 
@@ -3565,7 +3572,7 @@ static int bpf_kprobe_multi_filter(const struct bpf_prog *prog, u32 kfunc_id)
 	if (!btf_id_set8_contains(&kprobe_multi_kfunc_set_ids, kfunc_id))
 		return 0;
 
-	if (!is_kprobe_session(prog))
+	if (!is_kprobe_session(prog) && !is_uprobe_session(prog))
 		return -EACCES;
 
 	return 0;
-- 
2.45.2

[PATCHv2 bpf-next 4/9] libbpf: Add support for uprobe multi session attach

[Jiri Olsa]

Adding support to attach program in uprobe session mode
with bpf_program__attach_uprobe_multi function.

Adding session bool to bpf_uprobe_multi_opts struct that allows
to load and attach the bpf program via uprobe session.
the attachment to create uprobe multi session.

Also adding new program loader section that allows:
  SEC("uprobe.session/bpf_fentry_test*")

and loads/attaches uprobe program as uprobe session.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 tools/lib/bpf/bpf.c    |  1 +
 tools/lib/bpf/libbpf.c | 50 ++++++++++++++++++++++++++++++++++++++++--
 tools/lib/bpf/libbpf.h |  4 +++-
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/tools/lib/bpf/bpf.c b/tools/lib/bpf/bpf.c
index 2a4c71501a17..becdfa701c75 100644
--- a/tools/lib/bpf/bpf.c
+++ b/tools/lib/bpf/bpf.c
@@ -776,6 +776,7 @@ int bpf_link_create(int prog_fd, int target_fd,
 			return libbpf_err(-EINVAL);
 		break;
 	case BPF_TRACE_UPROBE_MULTI:
+	case BPF_TRACE_UPROBE_SESSION:
 		attr.link_create.uprobe_multi.flags = OPTS_GET(opts, uprobe_multi.flags, 0);
 		attr.link_create.uprobe_multi.cnt = OPTS_GET(opts, uprobe_multi.cnt, 0);
 		attr.link_create.uprobe_multi.path = ptr_to_u64(OPTS_GET(opts, uprobe_multi.path, 0));
diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
index 4a28fac4908a..492a8eb4d047 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -9344,6 +9344,7 @@ static int attach_trace(const struct bpf_program *prog, long cookie, struct bpf_
 static int attach_kprobe_multi(const struct bpf_program *prog, long cookie, struct bpf_link **link);
 static int attach_kprobe_session(const struct bpf_program *prog, long cookie, struct bpf_link **link);
 static int attach_uprobe_multi(const struct bpf_program *prog, long cookie, struct bpf_link **link);
+static int attach_uprobe_session(const struct bpf_program *prog, long cookie, struct bpf_link **link);
 static int attach_lsm(const struct bpf_program *prog, long cookie, struct bpf_link **link);
 static int attach_iter(const struct bpf_program *prog, long cookie, struct bpf_link **link);
 
@@ -9362,6 +9363,7 @@ static const struct bpf_sec_def section_defs[] = {
 	SEC_DEF("kprobe.session+",	KPROBE,	BPF_TRACE_KPROBE_SESSION, SEC_NONE, attach_kprobe_session),
 	SEC_DEF("uprobe.multi+",	KPROBE,	BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
 	SEC_DEF("uretprobe.multi+",	KPROBE,	BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
+	SEC_DEF("uprobe.session+",	KPROBE,	BPF_TRACE_UPROBE_SESSION, SEC_NONE, attach_uprobe_session),
 	SEC_DEF("uprobe.multi.s+",	KPROBE,	BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
 	SEC_DEF("uretprobe.multi.s+",	KPROBE,	BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
 	SEC_DEF("ksyscall+",		KPROBE,	0, SEC_NONE, attach_ksyscall),
@@ -11698,6 +11700,40 @@ static int attach_uprobe_multi(const struct bpf_program *prog, long cookie, stru
 	return ret;
 }
 
+static int attach_uprobe_session(const struct bpf_program *prog, long cookie, struct bpf_link **link)
+{
+	char *binary_path = NULL, *func_name = NULL;
+	LIBBPF_OPTS(bpf_uprobe_multi_opts, opts,
+		.session = true,
+	);
+	int n, ret = -EINVAL;
+	const char *spec;
+
+	*link = NULL;
+
+	spec = prog->sec_name + sizeof("uprobe.session/") - 1;
+	n = sscanf(spec, "%m[^:]:%m[^\n]",
+		   &binary_path, &func_name);
+
+	switch (n) {
+	case 1:
+		/* but auto-attach is impossible. */
+		ret = 0;
+		break;
+	case 2:
+		*link = bpf_program__attach_uprobe_multi(prog, -1, binary_path, func_name, &opts);
+		ret = *link ? 0 : -errno;
+		break;
+	default:
+		pr_warn("prog '%s': invalid format of section definition '%s'\n", prog->name,
+			prog->sec_name);
+		break;
+	}
+	free(binary_path);
+	free(func_name);
+	return ret;
+}
+
 static void gen_uprobe_legacy_event_name(char *buf, size_t buf_sz,
 					 const char *binary_path, uint64_t offset)
 {
@@ -11932,10 +11968,12 @@ bpf_program__attach_uprobe_multi(const struct bpf_program *prog,
 	const unsigned long *ref_ctr_offsets = NULL, *offsets = NULL;
 	LIBBPF_OPTS(bpf_link_create_opts, lopts);
 	unsigned long *resolved_offsets = NULL;
+	enum bpf_attach_type attach_type;
 	int err = 0, link_fd, prog_fd;
 	struct bpf_link *link = NULL;
 	char errmsg[STRERR_BUFSIZE];
 	char full_path[PATH_MAX];
+	bool retprobe, session;
 	const __u64 *cookies;
 	const char **syms;
 	size_t cnt;
@@ -12006,12 +12044,20 @@ bpf_program__attach_uprobe_multi(const struct bpf_program *prog,
 		offsets = resolved_offsets;
 	}
 
+	retprobe = OPTS_GET(opts, retprobe, false);
+	session  = OPTS_GET(opts, session, false);
+
+	if (retprobe && session)
+		return libbpf_err_ptr(-EINVAL);
+
+	attach_type = session ? BPF_TRACE_UPROBE_SESSION : BPF_TRACE_UPROBE_MULTI;
+
 	lopts.uprobe_multi.path = path;
 	lopts.uprobe_multi.offsets = offsets;
 	lopts.uprobe_multi.ref_ctr_offsets = ref_ctr_offsets;
 	lopts.uprobe_multi.cookies = cookies;
 	lopts.uprobe_multi.cnt = cnt;
-	lopts.uprobe_multi.flags = OPTS_GET(opts, retprobe, false) ? BPF_F_UPROBE_MULTI_RETURN : 0;
+	lopts.uprobe_multi.flags = retprobe ? BPF_F_UPROBE_MULTI_RETURN : 0;
 
 	if (pid == 0)
 		pid = getpid();
@@ -12025,7 +12071,7 @@ bpf_program__attach_uprobe_multi(const struct bpf_program *prog,
 	}
 	link->detach = &bpf_link__detach_fd;
 
-	link_fd = bpf_link_create(prog_fd, 0, BPF_TRACE_UPROBE_MULTI, &lopts);
+	link_fd = bpf_link_create(prog_fd, 0, attach_type, &lopts);
 	if (link_fd < 0) {
 		err = -errno;
 		pr_warn("prog '%s': failed to attach multi-uprobe: %s\n",
diff --git a/tools/lib/bpf/libbpf.h b/tools/lib/bpf/libbpf.h
index 64a6a3d323e3..f6a7835dc519 100644
--- a/tools/lib/bpf/libbpf.h
+++ b/tools/lib/bpf/libbpf.h
@@ -569,10 +569,12 @@ struct bpf_uprobe_multi_opts {
 	size_t cnt;
 	/* create return uprobes */
 	bool retprobe;
+	/* create session kprobes */
+	bool session;
 	size_t :0;
 };
 
-#define bpf_uprobe_multi_opts__last_field retprobe
+#define bpf_uprobe_multi_opts__last_field session
 
 /**
  * @brief **bpf_program__attach_uprobe_multi()** attaches a BPF program
-- 
2.45.2

[PATCHv2 bpf-next 5/9] libbpf: Add uprobe session attach type names to attach_type_name

[Jiri Olsa]

Adding uprobe session attach type name to attach_type_name,
so libbpf_bpf_attach_type_str returns proper string name for
BPF_TRACE_UPROBE_SESSION attach type.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 tools/lib/bpf/libbpf.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
index 492a8eb4d047..e69a54264580 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -133,6 +133,7 @@ static const char * const attach_type_name[] = {
 	[BPF_NETKIT_PRIMARY]		= "netkit_primary",
 	[BPF_NETKIT_PEER]		= "netkit_peer",
 	[BPF_TRACE_KPROBE_SESSION]	= "trace_kprobe_session",
+	[BPF_TRACE_UPROBE_SESSION]	= "trace_uprobe_session",
 };
 
 static const char * const link_type_name[] = {
-- 
2.45.2

[PATCHv2 bpf-next 6/9] selftests/bpf: Add uprobe session test

[Jiri Olsa]

Adding uprobe session test and testing that the entry program
return value controls execution of the return probe program.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 42 +++++++++++++++
 .../bpf/progs/uprobe_multi_session.c          | 53 +++++++++++++++++++
 2 files changed, 95 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session.c

diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
index bf6ca8e3eb13..cd9581f46c73 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -6,6 +6,7 @@
 #include "uprobe_multi.skel.h"
 #include "uprobe_multi_bench.skel.h"
 #include "uprobe_multi_usdt.skel.h"
+#include "uprobe_multi_session.skel.h"
 #include "bpf/libbpf_internal.h"
 #include "testing_helpers.h"
 #include "../sdt.h"
@@ -615,6 +616,45 @@ static void test_link_api(void)
 	__test_link_api(child);
 }
 
+static void test_session_skel_api(void)
+{
+	struct uprobe_multi_session *skel = NULL;
+	LIBBPF_OPTS(bpf_kprobe_multi_opts, opts);
+	struct bpf_link *link = NULL;
+	int err;
+
+	skel = uprobe_multi_session__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "fentry_raw_skel_load"))
+		goto cleanup;
+
+	skel->bss->pid = getpid();
+
+	err = uprobe_multi_session__attach(skel);
+	if (!ASSERT_OK(err, " uprobe_multi_session__attach"))
+		goto cleanup;
+
+	/* trigger all probes */
+	skel->bss->uprobe_multi_func_1_addr = (__u64) uprobe_multi_func_1;
+	skel->bss->uprobe_multi_func_2_addr = (__u64) uprobe_multi_func_2;
+	skel->bss->uprobe_multi_func_3_addr = (__u64) uprobe_multi_func_3;
+
+	uprobe_multi_func_1();
+	uprobe_multi_func_2();
+	uprobe_multi_func_3();
+
+	/*
+	 * We expect 2 for uprobe_multi_func_2 because it runs both entry/return probe,
+	 * uprobe_multi_func_[13] run just the entry probe.
+	 */
+	ASSERT_EQ(skel->bss->uprobe_session_result[0], 1, "uprobe_multi_func_1_result");
+	ASSERT_EQ(skel->bss->uprobe_session_result[1], 2, "uprobe_multi_func_2_result");
+	ASSERT_EQ(skel->bss->uprobe_session_result[2], 1, "uprobe_multi_func_3_result");
+
+cleanup:
+	bpf_link__destroy(link);
+	uprobe_multi_session__destroy(skel);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -703,4 +743,6 @@ void test_uprobe_multi_test(void)
 		test_bench_attach_usdt();
 	if (test__start_subtest("attach_api_fails"))
 		test_attach_api_fails();
+	if (test__start_subtest("session"))
+		test_session_skel_api();
 }
diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session.c
new file mode 100644
index 000000000000..72c00ae68372
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session.c
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+__u64 uprobe_multi_func_1_addr = 0;
+__u64 uprobe_multi_func_2_addr = 0;
+__u64 uprobe_multi_func_3_addr = 0;
+
+__u64 uprobe_session_result[3];
+
+int pid = 0;
+
+static int uprobe_multi_check(void *ctx, bool is_return)
+{
+	const __u64 funcs[] = {
+		uprobe_multi_func_1_addr,
+		uprobe_multi_func_2_addr,
+		uprobe_multi_func_3_addr,
+	};
+	unsigned int i;
+	__u64 addr;
+
+	if (bpf_get_current_pid_tgid() >> 32 != pid)
+		return 1;
+
+	addr = bpf_get_func_ip(ctx);
+
+	for (i = 0; i < ARRAY_SIZE(funcs); i++) {
+		if (funcs[i] == addr) {
+			uprobe_session_result[i]++;
+			break;
+		}
+	}
+
+	/* only uprobe_multi_func_2 executes return probe */
+	if ((addr == uprobe_multi_func_1_addr) ||
+	    (addr == uprobe_multi_func_3_addr))
+		return 1;
+
+	return 0;
+}
+
+SEC("uprobe.session//proc/self/exe:uprobe_multi_func_*")
+int uprobe(struct pt_regs *ctx)
+{
+	return uprobe_multi_check(ctx, bpf_session_is_return());
+}
-- 
2.45.2

[PATCHv2 bpf-next 7/9] selftests/bpf: Add uprobe session cookie test

[Jiri Olsa]

Adding uprobe session test that verifies the cookie value
get properly propagated from entry to return program.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 31 ++++++++++++
 .../bpf/progs/uprobe_multi_session_cookie.c   | 48 +++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c

diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
index cd9581f46c73..d5f78fc61013 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -7,6 +7,7 @@
 #include "uprobe_multi_bench.skel.h"
 #include "uprobe_multi_usdt.skel.h"
 #include "uprobe_multi_session.skel.h"
+#include "uprobe_multi_session_cookie.skel.h"
 #include "bpf/libbpf_internal.h"
 #include "testing_helpers.h"
 #include "../sdt.h"
@@ -655,6 +656,34 @@ static void test_session_skel_api(void)
 	uprobe_multi_session__destroy(skel);
 }
 
+static void test_session_cookie_skel_api(void)
+{
+	struct uprobe_multi_session_cookie *skel = NULL;
+	int err;
+
+	skel = uprobe_multi_session_cookie__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "fentry_raw_skel_load"))
+		goto cleanup;
+
+	skel->bss->pid = getpid();
+
+	err = uprobe_multi_session_cookie__attach(skel);
+	if (!ASSERT_OK(err, " kprobe_multi_session__attach"))
+		goto cleanup;
+
+	/* trigger all probes */
+	uprobe_multi_func_1();
+	uprobe_multi_func_2();
+	uprobe_multi_func_3();
+
+	ASSERT_EQ(skel->bss->test_uprobe_1_result, 1, "test_uprobe_1_result");
+	ASSERT_EQ(skel->bss->test_uprobe_2_result, 2, "test_uprobe_2_result");
+	ASSERT_EQ(skel->bss->test_uprobe_3_result, 3, "test_uprobe_3_result");
+
+cleanup:
+	uprobe_multi_session_cookie__destroy(skel);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -745,4 +774,6 @@ void test_uprobe_multi_test(void)
 		test_attach_api_fails();
 	if (test__start_subtest("session"))
 		test_session_skel_api();
+	if (test__start_subtest("session_cookie"))
+		test_session_cookie_skel_api();
 }
diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
new file mode 100644
index 000000000000..5befdf944dc6
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
+
+char _license[] SEC("license") = "GPL";
+
+int pid = 0;
+
+__u64 test_uprobe_1_result = 0;
+__u64 test_uprobe_2_result = 0;
+__u64 test_uprobe_3_result = 0;
+
+static int check_cookie(__u64 val, __u64 *result)
+{
+	__u64 *cookie;
+
+	if (bpf_get_current_pid_tgid() >> 32 != pid)
+		return 1;
+
+	cookie = bpf_session_cookie();
+
+	if (bpf_session_is_return())
+		*result = *cookie == val ? val : 0;
+	else
+		*cookie = val;
+	return 0;
+}
+
+SEC("uprobe.session//proc/self/exe:uprobe_multi_func_1")
+int uprobe_1(struct pt_regs *ctx)
+{
+	return check_cookie(1, &test_uprobe_1_result);
+}
+
+SEC("uprobe.session//proc/self/exe:uprobe_multi_func_2")
+int uprobe_2(struct pt_regs *ctx)
+{
+	return check_cookie(2, &test_uprobe_2_result);
+}
+
+SEC("uprobe.session//proc/self/exe:uprobe_multi_func_3")
+int uprobe_3(struct pt_regs *ctx)
+{
+	return check_cookie(3, &test_uprobe_3_result);
+}
-- 
2.45.2

[PATCHv2 bpf-next 8/9] selftests/bpf: Add uprobe session recursive test

[Jiri Olsa]

Adding uprobe session test that verifies the cookie value is stored
properly when single uprobe-ed function is executed recursively.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 57 +++++++++++++++++++
 .../progs/uprobe_multi_session_recursive.c    | 44 ++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c

diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
index d5f78fc61013..b521590fdbb9 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -8,6 +8,7 @@
 #include "uprobe_multi_usdt.skel.h"
 #include "uprobe_multi_session.skel.h"
 #include "uprobe_multi_session_cookie.skel.h"
+#include "uprobe_multi_session_recursive.skel.h"
 #include "bpf/libbpf_internal.h"
 #include "testing_helpers.h"
 #include "../sdt.h"
@@ -34,6 +35,12 @@ noinline void usdt_trigger(void)
 	STAP_PROBE(test, pid_filter_usdt);
 }
 
+noinline void uprobe_session_recursive(int i)
+{
+	if (i)
+		uprobe_session_recursive(i - 1);
+}
+
 struct child {
 	int go[2];
 	int c2p[2]; /* child -> parent channel */
@@ -684,6 +691,54 @@ static void test_session_cookie_skel_api(void)
 	uprobe_multi_session_cookie__destroy(skel);
 }
 
+static void test_session_recursive_skel_api(void)
+{
+	struct uprobe_multi_session_recursive *skel = NULL;
+	int i, err;
+
+	skel = uprobe_multi_session_recursive__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_recursive__open_and_load"))
+		goto cleanup;
+
+	skel->bss->pid = getpid();
+
+	err = uprobe_multi_session_recursive__attach(skel);
+	if (!ASSERT_OK(err, "uprobe_multi_session_recursive__attach"))
+		goto cleanup;
+
+	for (i = 0; i < ARRAY_SIZE(skel->bss->test_uprobe_cookie_entry); i++)
+		skel->bss->test_uprobe_cookie_entry[i] = i + 1;
+
+	uprobe_session_recursive(5);
+
+	/*
+	 *                                         entry uprobe:
+	 * uprobe_session_recursive(5) {             *cookie = 1, return 0
+	 *   uprobe_session_recursive(4) {           *cookie = 2, return 1
+	 *     uprobe_session_recursive(3) {         *cookie = 3, return 0
+	 *       uprobe_session_recursive(2) {       *cookie = 4, return 1
+	 *         uprobe_session_recursive(1) {     *cookie = 5, return 0
+	 *           uprobe_session_recursive(0) {   *cookie = 6, return 1
+	 *                                          return uprobe:
+	 *           } i = 0                          not executed
+	 *         } i = 1                            test_uprobe_cookie_return[0] = 5
+	 *       } i = 2                              not executed
+	 *     } i = 3                                test_uprobe_cookie_return[1] = 3
+	 *   } i = 4                                  not executed
+	 * } i = 5                                    test_uprobe_cookie_return[2] = 1
+	 */
+
+	ASSERT_EQ(skel->bss->idx_entry, 6, "idx_entry");
+	ASSERT_EQ(skel->bss->idx_return, 3, "idx_return");
+
+	ASSERT_EQ(skel->bss->test_uprobe_cookie_return[0], 5, "test_uprobe_cookie_return[0]");
+	ASSERT_EQ(skel->bss->test_uprobe_cookie_return[1], 3, "test_uprobe_cookie_return[1]");
+	ASSERT_EQ(skel->bss->test_uprobe_cookie_return[2], 1, "test_uprobe_cookie_return[2]");
+
+cleanup:
+	uprobe_multi_session_recursive__destroy(skel);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -776,4 +831,6 @@ void test_uprobe_multi_test(void)
 		test_session_skel_api();
 	if (test__start_subtest("session_cookie"))
 		test_session_cookie_skel_api();
+	if (test__start_subtest("session_cookie_recursive"))
+		test_session_recursive_skel_api();
 }
diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c
new file mode 100644
index 000000000000..8fbcd69fae22
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c
@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+int pid = 0;
+
+int idx_entry = 0;
+int idx_return = 0;
+
+__u64 test_uprobe_cookie_entry[6];
+__u64 test_uprobe_cookie_return[3];
+
+static int check_cookie(void)
+{
+	__u64 *cookie = bpf_session_cookie();
+
+	if (bpf_session_is_return()) {
+		if (idx_return >= ARRAY_SIZE(test_uprobe_cookie_return))
+			return 1;
+		test_uprobe_cookie_return[idx_return++] = *cookie;
+		return 0;
+	}
+
+	if (idx_entry >= ARRAY_SIZE(test_uprobe_cookie_entry))
+		return 1;
+	*cookie = test_uprobe_cookie_entry[idx_entry];
+	return idx_entry++ % 2;
+}
+
+
+SEC("uprobe.session//proc/self/exe:uprobe_session_recursive")
+int uprobe_recursive(struct pt_regs *ctx)
+{
+	if (bpf_get_current_pid_tgid() >> 32 != pid)
+		return 1;
+
+	return check_cookie();
+}
-- 
2.45.2

[PATCHv2 bpf-next 9/9] selftests/bpf: Add uprobe session consumers test

[Jiri Olsa]

Adding test that attached/detaches multiple consumers on
single uprobe and verifies all were hit as expected.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 203 ++++++++++++++++++
 .../progs/uprobe_multi_session_consumers.c    |  53 +++++
 2 files changed, 256 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c

diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
index b521590fdbb9..83eac954cf00 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -9,6 +9,7 @@
 #include "uprobe_multi_session.skel.h"
 #include "uprobe_multi_session_cookie.skel.h"
 #include "uprobe_multi_session_recursive.skel.h"
+#include "uprobe_multi_session_consumers.skel.h"
 #include "bpf/libbpf_internal.h"
 #include "testing_helpers.h"
 #include "../sdt.h"
@@ -739,6 +740,206 @@ static void test_session_recursive_skel_api(void)
 	uprobe_multi_session_recursive__destroy(skel);
 }
 
+static int uprobe_attach(struct uprobe_multi_session_consumers *skel, int bit)
+{
+	struct bpf_program **prog = &skel->progs.uprobe_0 + bit;
+	struct bpf_link **link = &skel->links.uprobe_0 + bit;
+	LIBBPF_OPTS(bpf_uprobe_multi_opts, opts);
+
+	/*
+	 * bit: 0,1 uprobe session
+	 * bit: 2,3 uprobe entry
+	 * bit: 4,5 uprobe return
+	 */
+	opts.session = bit < 2;
+	opts.retprobe = bit == 4 || bit == 5;
+
+	*link = bpf_program__attach_uprobe_multi(*prog, 0, "/proc/self/exe",
+						 "uprobe_session_consumer_test",
+						 &opts);
+	if (!ASSERT_OK_PTR(*link, "bpf_program__attach_uprobe_multi"))
+		return -1;
+	return 0;
+}
+
+static void uprobe_detach(struct uprobe_multi_session_consumers *skel, int bit)
+{
+	struct bpf_link **link = &skel->links.uprobe_0 + bit;
+
+	bpf_link__destroy(*link);
+	*link = NULL;
+}
+
+static bool test_bit(int bit, unsigned long val)
+{
+	return val & (1 << bit);
+}
+
+noinline int
+uprobe_session_consumer_test(struct uprobe_multi_session_consumers *skel,
+			     unsigned long before, unsigned long after)
+{
+	int bit;
+
+	/* detach uprobe for each unset bit in 'before' state ... */
+	for (bit = 0; bit < 6; bit++) {
+		if (test_bit(bit, before) && !test_bit(bit, after))
+			uprobe_detach(skel, bit);
+	}
+
+	/* ... and attach all new bits in 'after' state */
+	for (bit = 0; bit < 6; bit++) {
+		if (!test_bit(bit, before) && test_bit(bit, after)) {
+			if (!ASSERT_OK(uprobe_attach(skel, bit), "uprobe_attach_after"))
+				return -1;
+		}
+	}
+	return 0;
+}
+
+static void session_consumer_test(struct uprobe_multi_session_consumers *skel,
+				  unsigned long before, unsigned long after)
+{
+	int err, bit;
+
+	/* 'before' is each, we attach uprobe for every set bit */
+	for (bit = 0; bit < 6; bit++) {
+		if (test_bit(bit, before)) {
+			if (!ASSERT_OK(uprobe_attach(skel, bit), "uprobe_attach_before"))
+				goto cleanup;
+		}
+	}
+
+	err = uprobe_session_consumer_test(skel, before, after);
+	if (!ASSERT_EQ(err, 0, "uprobe_session_consumer_test"))
+		goto cleanup;
+
+	for (bit = 0; bit < 6; bit++) {
+		const char *fmt = "BUG";
+		__u64 val = 0;
+
+		if (bit == 0) {
+			/*
+			 * session with return
+			 *  +1 if defined in 'before'
+			 *  +1 if defined in 'after'
+			 */
+			if (test_bit(bit, before)) {
+				val++;
+				if (test_bit(bit, after))
+					val++;
+			}
+			fmt = "bit 0  : session with return";
+		} else if (bit == 1) {
+			/*
+			 * session without return
+			 *   +1 if defined in 'before'
+			 */
+			if (test_bit(bit, before))
+				val++;
+			fmt = "bit 1  : session with NO return";
+		} else if (bit < 4) {
+			/*
+			 * uprobe entry
+			 *   +1 if define in 'before'
+			 */
+			if (test_bit(bit, before))
+				val++;
+			fmt = "bit 3/4: uprobe";
+		} else {
+			/* uprobe return is tricky ;-)
+			 *
+			 * to trigger uretprobe consumer, the uretprobe needs to be installed,
+			 * which means one of the 'return' uprobes was alive when probe was hit:
+			 *
+			 *   bits: 0 (session with return) 4/5 uprobe return in 'installed' mask
+			 *
+			 * in addition if 'after' state removes everything that was installed in
+			 * 'before' state, then uprobe kernel object goes away and return uprobe
+			 * is not installed and we won't hit it even if it's in 'after' state.
+			 */
+			unsigned long installed = before & 0b110001; // is uretprobe installed
+			unsigned long exists    = before & after;    // did uprobe go away
+
+			if (installed && exists && test_bit(bit, after))
+				val++;
+			fmt = "bit 5/6: uretprobe";
+		}
+
+		ASSERT_EQ(skel->bss->uprobe_result[bit], val, fmt);
+		skel->bss->uprobe_result[bit] = 0;
+	}
+
+cleanup:
+	for (bit = 0; bit < 6; bit++) {
+		struct bpf_link **link = &skel->links.uprobe_0 + bit;
+
+		if (*link)
+			uprobe_detach(skel, bit);
+	}
+}
+
+static void test_session_consumers(void)
+{
+	struct uprobe_multi_session_consumers *skel;
+	int before, after;
+
+	skel = uprobe_multi_session_consumers__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_consumers__open_and_load"))
+		return;
+
+	/*
+	 * The idea of this test is to try all possible combinations of
+	 * uprobes consumers attached on single function.
+	 *
+	 *  - 1 uprobe session with return handler called
+	 *  - 1 uprobe session without return handler called
+	 *  - 2 uprobe entry consumer
+	 *  - 2 uprobe exit consumers
+	 *
+	 * The test uses 6 uprobes attached on single function, but that
+	 * translates into single uprobe with 6 consumers in kernel.
+	 *
+	 * The before/after values present the state of attached consumers
+	 * before and after the probed function:
+	 *
+	 *  bit 0   : uprobe session with return
+	 *  bit 1   : uprobe session with no return
+	 *  bit 2,3 : uprobe entry
+	 *  bit 4,5 : uprobe return
+	 *
+	 * For example for:
+	 *
+	 *   before = 0b10101
+	 *   after  = 0b00110
+	 *
+	 * it means that before we call 'uprobe_session_consumer_test' we
+	 * attach uprobes defined in 'before' value:
+	 *
+	 *   - bit 0: uprobe session with return
+	 *   - bit 2: uprobe entry
+	 *   - bit 4: uprobe return
+	 *
+	 * uprobe_session_consumer_test is called and inside it we attach
+	 * and detach * uprobes based on 'after' value:
+	 *
+	 *   - bit 0: uprobe session with return is detached
+	 *   - bit 1: uprobe session without return is attached
+	 *   - bit 2: stays untouched
+	 *   - bit 4: uprobe return is detached
+	 *
+	 * uprobe_session_consumer_test returs and we check counters values
+	 * increased by bpf programs on each uprobe to match the expected
+	 * count based on before/after bits.
+	 */
+	for (before = 0; before < 64; before++) {
+		for (after = 0; after < 64; after++)
+			session_consumer_test(skel, before, after);
+	}
+
+	uprobe_multi_session_consumers__destroy(skel);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -833,4 +1034,6 @@ void test_uprobe_multi_test(void)
 		test_session_cookie_skel_api();
 	if (test__start_subtest("session_cookie_recursive"))
 		test_session_recursive_skel_api();
+	if (test__start_subtest("session/consumers"))
+		test_session_consumers();
 }
diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
new file mode 100644
index 000000000000..035d31a0a7f8
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+__u64 uprobe_result[6];
+
+SEC("uprobe.session")
+int uprobe_0(struct pt_regs *ctx)
+{
+	uprobe_result[0]++;
+	return 0;
+}
+
+SEC("uprobe.session")
+int uprobe_1(struct pt_regs *ctx)
+{
+	uprobe_result[1]++;
+	return 1;
+}
+
+SEC("uprobe.multi")
+int uprobe_2(struct pt_regs *ctx)
+{
+	uprobe_result[2]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_3(struct pt_regs *ctx)
+{
+	uprobe_result[3]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_4(struct pt_regs *ctx)
+{
+	uprobe_result[4]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_5(struct pt_regs *ctx)
+{
+	uprobe_result[5]++;
+	return 0;
+}
-- 
2.45.2

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Peter Zijlstra]

On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:

> +static void
> +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> +{
> +	static unsigned int session_id;
> +
> +	if (uc->session) {
> +		uprobe->sessions_cnt++;
> +		uc->session_id = ++session_id ?: ++session_id;
> +	}
> +}

The way I understand this code, you create a consumer every time you do
uprobe_register() and unregister makes it go away.

Now, register one, then 4g-1 times register+unregister, then register
again.

The above seems to then result in two consumers with the same
session_id, which leads to trouble.

Hmm?

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Tue, Jul 02, 2024 at 03:04:08PM +0200, Peter Zijlstra wrote:
> On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:
> 
> > +static void
> > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > +{
> > +	static unsigned int session_id;
> > +
> > +	if (uc->session) {
> > +		uprobe->sessions_cnt++;
> > +		uc->session_id = ++session_id ?: ++session_id;
> > +	}
> > +}
> 
> The way I understand this code, you create a consumer every time you do
> uprobe_register() and unregister makes it go away.
> 
> Now, register one, then 4g-1 times register+unregister, then register
> again.
> 
> The above seems to then result in two consumers with the same
> session_id, which leads to trouble.
> 
> Hmm?

ugh true.. will make it u64 :)

I think we could store uprobe_consumer pointer+ref in session_consumer,
and that would make the unregister path more interesting.. will check

thanks,
jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:41 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding support for uprobe consumer to be defined as session and have
> new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
>
> The session means that 'handler' and 'ret_handler' callbacks are
> connected in a way that allows to:
>
>   - control execution of 'ret_handler' from 'handler' callback
>   - share data between 'handler' and 'ret_handler' callbacks
>
> The session is enabled by setting new 'session' bool field to true
> in uprobe_consumer object.
>
> We keep count of session consumers for uprobe and allocate session_consumer
> object for each in return_instance object. This allows us to store
> return values of 'handler' callbacks and data pointers of shared
> data between both handlers.
>
> The session concept fits to our common use case where we do filtering
> on entry uprobe and based on the result we decide to run the return
> uprobe (or not).
>
> It's also convenient to share the data between session callbacks.
>
> The control of 'ret_handler' callback execution is done via return
> value of the 'handler' callback. If it's 0 we install and execute
> return uprobe, if it's 1 we do not.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/linux/uprobes.h     |  16 ++++-
>  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
>  kernel/trace/bpf_trace.c    |   6 +-
>  kernel/trace/trace_uprobe.c |  12 ++--
>  4 files changed, 144 insertions(+), 19 deletions(-)
>
> diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> index f46e0ca0169c..903a860a8d01 100644
> --- a/include/linux/uprobes.h
> +++ b/include/linux/uprobes.h
> @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
>  };
>
>  struct uprobe_consumer {
> -       int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> +       int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
>         int (*ret_handler)(struct uprobe_consumer *self,
>                                 unsigned long func,
> -                               struct pt_regs *regs);
> +                               struct pt_regs *regs, __u64 *data);
>         bool (*filter)(struct uprobe_consumer *self,
>                                 enum uprobe_filter_ctx ctx,
>                                 struct mm_struct *mm);
>
>         struct uprobe_consumer *next;
> +
> +       bool                    session;        /* marks uprobe session consumer */
> +       unsigned int            session_id;     /* set when uprobe_consumer is registered */
>  };
>
>  #ifdef CONFIG_UPROBES
> @@ -80,6 +83,12 @@ struct uprobe_task {
>         unsigned int                    depth;
>  };
>
> +struct session_consumer {
> +       __u64           cookie;
> +       unsigned int    id;
> +       int             rc;

you'll be using u64 for ID, right? so this struct will be 24 bytes.
Maybe we can just use topmost bit of ID to store whether uretprobe
should run or not? It's trivial to mask out during ID comparisons

> +};
> +
>  struct return_instance {
>         struct uprobe           *uprobe;
>         unsigned long           func;
> @@ -88,6 +97,9 @@ struct return_instance {
>         bool                    chained;        /* true, if instance is nested */
>
>         struct return_instance  *next;          /* keep as stack */
> +
> +       int                     sessions_cnt;

there is 7 byte gap before next field, let's put sessions_cnt there

> +       struct session_consumer sessions[];
>  };
>
>  enum rp_check {
> diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> index 2c83ba776fc7..4da410460f2a 100644
> --- a/kernel/events/uprobes.c
> +++ b/kernel/events/uprobes.c
> @@ -63,6 +63,8 @@ struct uprobe {
>         loff_t                  ref_ctr_offset;
>         unsigned long           flags;
>
> +       unsigned int            sessions_cnt;
> +
>         /*
>          * The generic code assumes that it has two members of unknown type
>          * owned by the arch-specific code:
> @@ -750,11 +752,30 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
>         return uprobe;
>  }
>
> +static void
> +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> +{
> +       static unsigned int session_id;

(besides what Peter mentioned about wrap around of 32-bit counter)
let's use atomic here to not rely on any particular locking
(unnecessarily), this might make my life easier in the future, thanks.
This is registration time, low frequency, extra atomic won't hurt.

It might be already broken, actually, for two independently registering uprobes.

> +
> +       if (uc->session) {
> +               uprobe->sessions_cnt++;
> +               uc->session_id = ++session_id ?: ++session_id;
> +       }
> +}
> +
> +static void
> +uprobe_consumer_unaccount(struct uprobe *uprobe, struct uprobe_consumer *uc)

this fits in 100 characters, keep it single line, please. Same for
account function

> +{
> +       if (uc->session)
> +               uprobe->sessions_cnt--;
> +}
> +
>  static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
>  {
>         down_write(&uprobe->consumer_rwsem);
>         uc->next = uprobe->consumers;
>         uprobe->consumers = uc;
> +       uprobe_consumer_account(uprobe, uc);
>         up_write(&uprobe->consumer_rwsem);
>  }
>
> @@ -773,6 +794,7 @@ static bool consumer_del(struct uprobe *uprobe, struct uprobe_consumer *uc)
>                 if (*con == uc) {
>                         *con = uc->next;
>                         ret = true;
> +                       uprobe_consumer_unaccount(uprobe, uc);
>                         break;
>                 }
>         }
> @@ -1744,6 +1766,23 @@ static struct uprobe_task *get_utask(void)
>         return current->utask;
>  }
>
> +static size_t ri_size(int sessions_cnt)
> +{
> +       struct return_instance *ri __maybe_unused;
> +
> +       return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);

just use struct_size()?

> +}
> +
> +static struct return_instance *alloc_return_instance(int sessions_cnt)
> +{
> +       struct return_instance *ri;
> +
> +       ri = kzalloc(ri_size(sessions_cnt), GFP_KERNEL);
> +       if (ri)
> +               ri->sessions_cnt = sessions_cnt;
> +       return ri;
> +}
> +
>  static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
>  {
>         struct uprobe_task *n_utask;
> @@ -1756,11 +1795,11 @@ static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
>
>         p = &n_utask->return_instances;
>         for (o = o_utask->return_instances; o; o = o->next) {
> -               n = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> +               n = alloc_return_instance(o->sessions_cnt);
>                 if (!n)
>                         return -ENOMEM;
>
> -               *n = *o;
> +               memcpy(n, o, ri_size(o->sessions_cnt));
>                 get_uprobe(n->uprobe);
>                 n->next = NULL;
>
> @@ -1853,9 +1892,9 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
>         utask->return_instances = ri;
>  }
>
> -static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
> +static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
> +                             struct return_instance *ri)
>  {
> -       struct return_instance *ri;
>         struct uprobe_task *utask;
>         unsigned long orig_ret_vaddr, trampoline_vaddr;
>         bool chained;
> @@ -1874,9 +1913,11 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
>                 return;
>         }
>
> -       ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> -       if (!ri)
> -               return;
> +       if (!ri) {
> +               ri = alloc_return_instance(0);
> +               if (!ri)
> +                       return;
> +       }
>
>         trampoline_vaddr = get_trampoline_vaddr();
>         orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
> @@ -2065,35 +2106,85 @@ static struct uprobe *find_active_uprobe(unsigned long bp_vaddr, int *is_swbp)
>         return uprobe;
>  }
>
> +static struct session_consumer *
> +session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
> +                     int session_id)
> +{
> +       struct session_consumer *next;
> +
> +       next = sc ? sc + 1 : &ri->sessions[0];
> +       next->id = session_id;

it's kind of unexpected that "session_consumer_next" would actually
set an ID... Maybe drop int session_id as input argument and fill it
out outside of this function, this function being just a simple
iterator?

> +       return next;
> +}
> +
> +static struct session_consumer *
> +session_consumer_find(struct return_instance *ri, int *iter, int session_id)
> +{
> +       struct session_consumer *sc;
> +       int idx = *iter;
> +
> +       for (sc = &ri->sessions[idx]; idx < ri->sessions_cnt; idx++, sc++) {
> +               if (sc->id == session_id) {
> +                       *iter = idx + 1;
> +                       return sc;
> +               }
> +       }
> +       return NULL;
> +}
> +
>  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
>  {
>         struct uprobe_consumer *uc;
>         int remove = UPROBE_HANDLER_REMOVE;
> +       struct session_consumer *sc = NULL;
> +       struct return_instance *ri = NULL;
>         bool need_prep = false; /* prepare return uprobe, when needed */
>
>         down_read(&uprobe->register_rwsem);
> +       if (uprobe->sessions_cnt) {
> +               ri = alloc_return_instance(uprobe->sessions_cnt);
> +               if (!ri)
> +                       goto out;
> +       }
> +
>         for (uc = uprobe->consumers; uc; uc = uc->next) {
> +               __u64 *cookie = NULL;
>                 int rc = 0;
>
> +               if (uc->session) {
> +                       sc = session_consumer_next(ri, sc, uc->session_id);
> +                       cookie = &sc->cookie;
> +               }
> +
>                 if (uc->handler) {
> -                       rc = uc->handler(uc, regs);
> +                       rc = uc->handler(uc, regs, cookie);
>                         WARN(rc & ~UPROBE_HANDLER_MASK,
>                                 "bad rc=0x%x from %ps()\n", rc, uc->handler);
>                 }
>
> -               if (uc->ret_handler)
> +               if (uc->session) {
> +                       sc->rc = rc;
> +                       need_prep |= !rc;

nit:

if (rc == 0)
    need_prep = true;

and then it's *extremely obvious* what happens and under which conditions

> +               } else if (uc->ret_handler) {
>                         need_prep = true;
> +               }
>
>                 remove &= rc;
>         }
>
> +       /* no removal if there's at least one session consumer */
> +       remove &= !uprobe->sessions_cnt;

this is counter (not error, not pointer), let's stick to ` == 0`, please

is this

if (uprobe->sessions_cnt != 0)
   remove = 0;

? I can't tell (honestly), without spending ridiculous amounts of
mental resources (for the underlying simplicity of the condition).

> +
>         if (need_prep && !remove)
> -               prepare_uretprobe(uprobe, regs); /* put bp at return */
> +               prepare_uretprobe(uprobe, regs, ri); /* put bp at return */
> +       else
> +               kfree(ri);
>
>         if (remove && uprobe->consumers) {
>                 WARN_ON(!uprobe_is_active(uprobe));
>                 unapply_uprobe(uprobe, current->mm);
>         }
> + out:
>         up_read(&uprobe->register_rwsem);
>  }
>

[...]

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Tue, Jul 2, 2024 at 9:11 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Tue, Jul 02, 2024 at 03:04:08PM +0200, Peter Zijlstra wrote:
> > On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:
> >
> > > +static void
> > > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > > +{
> > > +   static unsigned int session_id;
> > > +
> > > +   if (uc->session) {
> > > +           uprobe->sessions_cnt++;
> > > +           uc->session_id = ++session_id ?: ++session_id;
> > > +   }
> > > +}
> >
> > The way I understand this code, you create a consumer every time you do
> > uprobe_register() and unregister makes it go away.
> >
> > Now, register one, then 4g-1 times register+unregister, then register
> > again.
> >
> > The above seems to then result in two consumers with the same
> > session_id, which leads to trouble.
> >
> > Hmm?
>
> ugh true.. will make it u64 :)
>
> I think we could store uprobe_consumer pointer+ref in session_consumer,
> and that would make the unregister path more interesting.. will check

More interesting how? It's actually a great idea, uprobe_consumer
pointer itself is a unique ID and 64-bit. We can still use lowest bit
for RC (see my other reply).

>
> thanks,
> jirka

Re: [PATCHv2 bpf-next 2/9] bpf: Add support for uprobe multi session attach

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:42 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding support to attach bpf program for entry and return probe
> of the same function. This is common use case which at the moment
> requires to create two uprobe multi links.
>
> Adding new BPF_TRACE_UPROBE_SESSION attach type that instructs
> kernel to attach single 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 the bpf program on return
> probe respectively.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/uapi/linux/bpf.h       |  1 +
>  kernel/bpf/syscall.c           |  9 +++++++--
>  kernel/trace/bpf_trace.c       | 25 +++++++++++++++++++------
>  tools/include/uapi/linux/bpf.h |  1 +
>  4 files changed, 28 insertions(+), 8 deletions(-)
>

LGTM

Acked-by: Andrii Nakryiko <andrii@kernel.org>

> diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> index 35bcf52dbc65..1d93cb014884 100644
> --- a/include/uapi/linux/bpf.h
> +++ b/include/uapi/linux/bpf.h
> @@ -1116,6 +1116,7 @@ enum bpf_attach_type {
>         BPF_NETKIT_PRIMARY,
>         BPF_NETKIT_PEER,
>         BPF_TRACE_KPROBE_SESSION,
> +       BPF_TRACE_UPROBE_SESSION,
>         __MAX_BPF_ATTACH_TYPE
>  };
>
> diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
> index 869265852d51..2a63a528fa3c 100644
> --- a/kernel/bpf/syscall.c
> +++ b/kernel/bpf/syscall.c
> @@ -4049,10 +4049,14 @@ static int bpf_prog_attach_check_attach_type(const struct bpf_prog *prog,
>                 if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI &&
>                     attach_type != BPF_TRACE_UPROBE_MULTI)
>                         return -EINVAL;
> +               if (prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION &&
> +                   attach_type != BPF_TRACE_UPROBE_SESSION)
> +                       return -EINVAL;
>                 if (attach_type != BPF_PERF_EVENT &&
>                     attach_type != BPF_TRACE_KPROBE_MULTI &&
>                     attach_type != BPF_TRACE_KPROBE_SESSION &&
> -                   attach_type != BPF_TRACE_UPROBE_MULTI)
> +                   attach_type != BPF_TRACE_UPROBE_MULTI &&
> +                   attach_type != BPF_TRACE_UPROBE_SESSION)
>                         return -EINVAL;
>                 return 0;
>         case BPF_PROG_TYPE_SCHED_CLS:
> @@ -5315,7 +5319,8 @@ static int link_create(union bpf_attr *attr, bpfptr_t uattr)
>                 else if (attr->link_create.attach_type == BPF_TRACE_KPROBE_MULTI ||
>                          attr->link_create.attach_type == BPF_TRACE_KPROBE_SESSION)
>                         ret = bpf_kprobe_multi_link_attach(attr, prog);
> -               else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI)
> +               else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI ||
> +                        attr->link_create.attach_type == BPF_TRACE_UPROBE_SESSION)
>                         ret = bpf_uprobe_multi_link_attach(attr, prog);
>                 break;
>         default:
> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index 02d052639dfe..1b19c1cdb5e1 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -1645,6 +1645,17 @@ static inline bool is_kprobe_session(const struct bpf_prog *prog)
>         return prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION;
>  }
>
> +static inline bool is_uprobe_multi(const struct bpf_prog *prog)
> +{
> +       return prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI ||
> +              prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
> +}
> +
> +static inline bool is_uprobe_session(const struct bpf_prog *prog)
> +{
> +       return prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
> +}
> +
>  static const struct bpf_func_proto *
>  kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
>  {
> @@ -1662,13 +1673,13 @@ kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
>         case BPF_FUNC_get_func_ip:
>                 if (is_kprobe_multi(prog))
>                         return &bpf_get_func_ip_proto_kprobe_multi;
> -               if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
> +               if (is_uprobe_multi(prog))
>                         return &bpf_get_func_ip_proto_uprobe_multi;
>                 return &bpf_get_func_ip_proto_kprobe;
>         case BPF_FUNC_get_attach_cookie:
>                 if (is_kprobe_multi(prog))
>                         return &bpf_get_attach_cookie_proto_kmulti;
> -               if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
> +               if (is_uprobe_multi(prog))
>                         return &bpf_get_attach_cookie_proto_umulti;
>                 return &bpf_get_attach_cookie_proto_trace;
>         default:
> @@ -3387,7 +3398,7 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
>         if (sizeof(u64) != sizeof(void *))
>                 return -EOPNOTSUPP;
>
> -       if (prog->expected_attach_type != BPF_TRACE_UPROBE_MULTI)
> +       if (!is_uprobe_multi(prog))
>                 return -EINVAL;
>
>         flags = attr->link_create.uprobe_multi.flags;
> @@ -3463,10 +3474,12 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
>
>                 uprobes[i].link = link;
>
> -               if (flags & BPF_F_UPROBE_MULTI_RETURN)
> -                       uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
> -               else
> +               if (!(flags & BPF_F_UPROBE_MULTI_RETURN))
>                         uprobes[i].consumer.handler = uprobe_multi_link_handler;
> +               if (flags & BPF_F_UPROBE_MULTI_RETURN || is_uprobe_session(prog))
> +                       uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
> +               if (is_uprobe_session(prog))
> +                       uprobes[i].consumer.session = true;
>
>                 if (pid)
>                         uprobes[i].consumer.filter = uprobe_multi_link_filter;
> diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
> index 35bcf52dbc65..1d93cb014884 100644
> --- a/tools/include/uapi/linux/bpf.h
> +++ b/tools/include/uapi/linux/bpf.h
> @@ -1116,6 +1116,7 @@ enum bpf_attach_type {
>         BPF_NETKIT_PRIMARY,
>         BPF_NETKIT_PEER,
>         BPF_TRACE_KPROBE_SESSION,
> +       BPF_TRACE_UPROBE_SESSION,
>         __MAX_BPF_ATTACH_TYPE
>  };
>
> --
> 2.45.2
>

Re: [PATCHv2 bpf-next 3/9] bpf: Add support for uprobe multi session context

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:42 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Placing bpf_session_run_ctx layer in between bpf_run_ctx and
> bpf_uprobe_multi_run_ctx, so the session data can be retrieved
> from uprobe_multi link.
>
> Plus granting session kfuncs access to uprobe session programs.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  kernel/trace/bpf_trace.c | 23 +++++++++++++++--------
>  1 file changed, 15 insertions(+), 8 deletions(-)
>

Acked-by: Andrii Nakryiko <andrii@kernel.org>


> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index 1b19c1cdb5e1..d431b880ca11 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -3184,7 +3184,7 @@ struct bpf_uprobe_multi_link {
>  };
>
>  struct bpf_uprobe_multi_run_ctx {
> -       struct bpf_run_ctx run_ctx;
> +       struct bpf_session_run_ctx session_ctx;
>         unsigned long entry_ip;
>         struct bpf_uprobe *uprobe;
>  };
> @@ -3297,10 +3297,15 @@ static const struct bpf_link_ops bpf_uprobe_multi_link_lops = {
>
>  static int uprobe_prog_run(struct bpf_uprobe *uprobe,
>                            unsigned long entry_ip,
> -                          struct pt_regs *regs)
> +                          struct pt_regs *regs,
> +                          bool is_return, void *data)
>  {
>         struct bpf_uprobe_multi_link *link = uprobe->link;
>         struct bpf_uprobe_multi_run_ctx run_ctx = {
> +               .session_ctx = {
> +                       .is_return = is_return,
> +                       .data = data,
> +               },
>                 .entry_ip = entry_ip,
>                 .uprobe = uprobe,
>         };
> @@ -3319,7 +3324,7 @@ static int uprobe_prog_run(struct bpf_uprobe *uprobe,
>
>         migrate_disable();
>
> -       old_run_ctx = bpf_set_run_ctx(&run_ctx.run_ctx);
> +       old_run_ctx = bpf_set_run_ctx(&run_ctx.session_ctx.run_ctx);
>         err = bpf_prog_run(link->link.prog, regs);
>         bpf_reset_run_ctx(old_run_ctx);
>
> @@ -3349,7 +3354,7 @@ uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs,
>         struct bpf_uprobe *uprobe;
>
>         uprobe = container_of(con, struct bpf_uprobe, consumer);
> -       return uprobe_prog_run(uprobe, instruction_pointer(regs), regs);
> +       return uprobe_prog_run(uprobe, instruction_pointer(regs), regs, false, data);
>  }
>
>  static int
> @@ -3359,14 +3364,15 @@ uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, s
>         struct bpf_uprobe *uprobe;
>
>         uprobe = container_of(con, struct bpf_uprobe, consumer);
> -       return uprobe_prog_run(uprobe, func, regs);
> +       return uprobe_prog_run(uprobe, func, regs, true, data);
>  }
>
>  static u64 bpf_uprobe_multi_entry_ip(struct bpf_run_ctx *ctx)
>  {
>         struct bpf_uprobe_multi_run_ctx *run_ctx;
>
> -       run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx, run_ctx);
> +       run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx,
> +                              session_ctx.run_ctx);
>         return run_ctx->entry_ip;
>  }
>
> @@ -3374,7 +3380,8 @@ static u64 bpf_uprobe_multi_cookie(struct bpf_run_ctx *ctx)
>  {
>         struct bpf_uprobe_multi_run_ctx *run_ctx;
>
> -       run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx, run_ctx);
> +       run_ctx = container_of(current->bpf_ctx, struct bpf_uprobe_multi_run_ctx,
> +                              session_ctx.run_ctx);
>         return run_ctx->uprobe->cookie;
>  }
>
> @@ -3565,7 +3572,7 @@ static int bpf_kprobe_multi_filter(const struct bpf_prog *prog, u32 kfunc_id)
>         if (!btf_id_set8_contains(&kprobe_multi_kfunc_set_ids, kfunc_id))
>                 return 0;
>
> -       if (!is_kprobe_session(prog))
> +       if (!is_kprobe_session(prog) && !is_uprobe_session(prog))
>                 return -EACCES;
>
>         return 0;
> --
> 2.45.2
>
>

Re: [PATCHv2 bpf-next 4/9] libbpf: Add support for uprobe multi session attach

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:42 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding support to attach program in uprobe session mode
> with bpf_program__attach_uprobe_multi function.
>
> Adding session bool to bpf_uprobe_multi_opts struct that allows
> to load and attach the bpf program via uprobe session.
> the attachment to create uprobe multi session.
>
> Also adding new program loader section that allows:
>   SEC("uprobe.session/bpf_fentry_test*")
>
> and loads/attaches uprobe program as uprobe session.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  tools/lib/bpf/bpf.c    |  1 +
>  tools/lib/bpf/libbpf.c | 50 ++++++++++++++++++++++++++++++++++++++++--
>  tools/lib/bpf/libbpf.h |  4 +++-
>  3 files changed, 52 insertions(+), 3 deletions(-)
>

[...]

> @@ -9362,6 +9363,7 @@ static const struct bpf_sec_def section_defs[] = {
>         SEC_DEF("kprobe.session+",      KPROBE, BPF_TRACE_KPROBE_SESSION, SEC_NONE, attach_kprobe_session),
>         SEC_DEF("uprobe.multi+",        KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
>         SEC_DEF("uretprobe.multi+",     KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
> +       SEC_DEF("uprobe.session+",      KPROBE, BPF_TRACE_UPROBE_SESSION, SEC_NONE, attach_uprobe_session),

sleepable ones as well?

>         SEC_DEF("uprobe.multi.s+",      KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
>         SEC_DEF("uretprobe.multi.s+",   KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
>         SEC_DEF("ksyscall+",            KPROBE, 0, SEC_NONE, attach_ksyscall),
> @@ -11698,6 +11700,40 @@ static int attach_uprobe_multi(const struct bpf_program *prog, long cookie, stru
>         return ret;
>  }
>
> +static int attach_uprobe_session(const struct bpf_program *prog, long cookie, struct bpf_link **link)
> +{
> +       char *binary_path = NULL, *func_name = NULL;
> +       LIBBPF_OPTS(bpf_uprobe_multi_opts, opts,
> +               .session = true,
> +       );

nit: keep a single line?

> +       int n, ret = -EINVAL;
> +       const char *spec;
> +
> +       *link = NULL;
> +
> +       spec = prog->sec_name + sizeof("uprobe.session/") - 1;
> +       n = sscanf(spec, "%m[^:]:%m[^\n]",
> +                  &binary_path, &func_name);

single line, wrapping lines is a necessary evil, please

[...]

Re: [PATCHv2 bpf-next 5/9] libbpf: Add uprobe session attach type names to attach_type_name

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding uprobe session attach type name to attach_type_name,
> so libbpf_bpf_attach_type_str returns proper string name for
> BPF_TRACE_UPROBE_SESSION attach type.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  tools/lib/bpf/libbpf.c | 1 +
>  1 file changed, 1 insertion(+)
>

Can you merge this into a patch that adds BPF_TRACE_UPROBE_SESSION to
keep bisectability of BPF selftests? It's a trivial patch, so
shouldn't be a big deal.

> diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
> index 492a8eb4d047..e69a54264580 100644
> --- a/tools/lib/bpf/libbpf.c
> +++ b/tools/lib/bpf/libbpf.c
> @@ -133,6 +133,7 @@ static const char * const attach_type_name[] = {
>         [BPF_NETKIT_PRIMARY]            = "netkit_primary",
>         [BPF_NETKIT_PEER]               = "netkit_peer",
>         [BPF_TRACE_KPROBE_SESSION]      = "trace_kprobe_session",
> +       [BPF_TRACE_UPROBE_SESSION]      = "trace_uprobe_session",
>  };
>
>  static const char * const link_type_name[] = {
> --
> 2.45.2
>

Re: [PATCHv2 bpf-next 6/9] selftests/bpf: Add uprobe session test

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding uprobe session test and testing that the entry program
> return value controls execution of the return probe program.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  .../bpf/prog_tests/uprobe_multi_test.c        | 42 +++++++++++++++
>  .../bpf/progs/uprobe_multi_session.c          | 53 +++++++++++++++++++
>  2 files changed, 95 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session.c
>

LGTM.
Acked-by: Andrii Nakryiko <andrii@kernel.org>

> diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> index bf6ca8e3eb13..cd9581f46c73 100644
> --- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> +++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> @@ -6,6 +6,7 @@
>  #include "uprobe_multi.skel.h"
>  #include "uprobe_multi_bench.skel.h"
>  #include "uprobe_multi_usdt.skel.h"
> +#include "uprobe_multi_session.skel.h"
>  #include "bpf/libbpf_internal.h"
>  #include "testing_helpers.h"
>  #include "../sdt.h"
> @@ -615,6 +616,45 @@ static void test_link_api(void)
>         __test_link_api(child);
>  }
>
> +static void test_session_skel_api(void)
> +{
> +       struct uprobe_multi_session *skel = NULL;
> +       LIBBPF_OPTS(bpf_kprobe_multi_opts, opts);
> +       struct bpf_link *link = NULL;
> +       int err;
> +
> +       skel = uprobe_multi_session__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "fentry_raw_skel_load"))
> +               goto cleanup;
> +
> +       skel->bss->pid = getpid();
> +
> +       err = uprobe_multi_session__attach(skel);
> +       if (!ASSERT_OK(err, " uprobe_multi_session__attach"))
> +               goto cleanup;
> +
> +       /* trigger all probes */
> +       skel->bss->uprobe_multi_func_1_addr = (__u64) uprobe_multi_func_1;
> +       skel->bss->uprobe_multi_func_2_addr = (__u64) uprobe_multi_func_2;
> +       skel->bss->uprobe_multi_func_3_addr = (__u64) uprobe_multi_func_3;
> +
> +       uprobe_multi_func_1();
> +       uprobe_multi_func_2();
> +       uprobe_multi_func_3();
> +
> +       /*
> +        * We expect 2 for uprobe_multi_func_2 because it runs both entry/return probe,
> +        * uprobe_multi_func_[13] run just the entry probe.
> +        */
> +       ASSERT_EQ(skel->bss->uprobe_session_result[0], 1, "uprobe_multi_func_1_result");
> +       ASSERT_EQ(skel->bss->uprobe_session_result[1], 2, "uprobe_multi_func_2_result");
> +       ASSERT_EQ(skel->bss->uprobe_session_result[2], 1, "uprobe_multi_func_3_result");
> +
> +cleanup:
> +       bpf_link__destroy(link);
> +       uprobe_multi_session__destroy(skel);
> +}
> +
>  static void test_bench_attach_uprobe(void)
>  {
>         long attach_start_ns = 0, attach_end_ns = 0;
> @@ -703,4 +743,6 @@ void test_uprobe_multi_test(void)
>                 test_bench_attach_usdt();
>         if (test__start_subtest("attach_api_fails"))
>                 test_attach_api_fails();
> +       if (test__start_subtest("session"))
> +               test_session_skel_api();
>  }
> diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session.c
> new file mode 100644
> index 000000000000..72c00ae68372
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session.c
> @@ -0,0 +1,53 @@
> +// SPDX-License-Identifier: GPL-2.0
> +#include <linux/bpf.h>
> +#include <bpf/bpf_helpers.h>
> +#include <bpf/bpf_tracing.h>
> +#include <stdbool.h>
> +#include "bpf_kfuncs.h"
> +#include "bpf_misc.h"
> +
> +char _license[] SEC("license") = "GPL";
> +
> +__u64 uprobe_multi_func_1_addr = 0;
> +__u64 uprobe_multi_func_2_addr = 0;
> +__u64 uprobe_multi_func_3_addr = 0;
> +
> +__u64 uprobe_session_result[3];
> +
> +int pid = 0;
> +
> +static int uprobe_multi_check(void *ctx, bool is_return)
> +{
> +       const __u64 funcs[] = {
> +               uprobe_multi_func_1_addr,
> +               uprobe_multi_func_2_addr,
> +               uprobe_multi_func_3_addr,
> +       };
> +       unsigned int i;
> +       __u64 addr;
> +
> +       if (bpf_get_current_pid_tgid() >> 32 != pid)
> +               return 1;
> +
> +       addr = bpf_get_func_ip(ctx);
> +
> +       for (i = 0; i < ARRAY_SIZE(funcs); i++) {
> +               if (funcs[i] == addr) {
> +                       uprobe_session_result[i]++;
> +                       break;
> +               }
> +       }
> +
> +       /* only uprobe_multi_func_2 executes return probe */
> +       if ((addr == uprobe_multi_func_1_addr) ||
> +           (addr == uprobe_multi_func_3_addr))
> +               return 1;
> +
> +       return 0;
> +}
> +
> +SEC("uprobe.session//proc/self/exe:uprobe_multi_func_*")
> +int uprobe(struct pt_regs *ctx)
> +{
> +       return uprobe_multi_check(ctx, bpf_session_is_return());
> +}
> --
> 2.45.2
>

Re: [PATCHv2 bpf-next 7/9] selftests/bpf: Add uprobe session cookie test

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding uprobe session test that verifies the cookie value
> get properly propagated from entry to return program.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  .../bpf/prog_tests/uprobe_multi_test.c        | 31 ++++++++++++
>  .../bpf/progs/uprobe_multi_session_cookie.c   | 48 +++++++++++++++++++
>  2 files changed, 79 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
>

LGTM

Acked-by: Andrii Nakryiko <andrii@kernel.org>


> diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> index cd9581f46c73..d5f78fc61013 100644
> --- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> +++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> @@ -7,6 +7,7 @@
>  #include "uprobe_multi_bench.skel.h"
>  #include "uprobe_multi_usdt.skel.h"
>  #include "uprobe_multi_session.skel.h"
> +#include "uprobe_multi_session_cookie.skel.h"
>  #include "bpf/libbpf_internal.h"
>  #include "testing_helpers.h"
>  #include "../sdt.h"
> @@ -655,6 +656,34 @@ static void test_session_skel_api(void)
>         uprobe_multi_session__destroy(skel);
>  }
>
> +static void test_session_cookie_skel_api(void)
> +{
> +       struct uprobe_multi_session_cookie *skel = NULL;
> +       int err;
> +
> +       skel = uprobe_multi_session_cookie__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "fentry_raw_skel_load"))
> +               goto cleanup;
> +
> +       skel->bss->pid = getpid();
> +
> +       err = uprobe_multi_session_cookie__attach(skel);
> +       if (!ASSERT_OK(err, " kprobe_multi_session__attach"))
> +               goto cleanup;
> +
> +       /* trigger all probes */
> +       uprobe_multi_func_1();
> +       uprobe_multi_func_2();
> +       uprobe_multi_func_3();
> +
> +       ASSERT_EQ(skel->bss->test_uprobe_1_result, 1, "test_uprobe_1_result");
> +       ASSERT_EQ(skel->bss->test_uprobe_2_result, 2, "test_uprobe_2_result");
> +       ASSERT_EQ(skel->bss->test_uprobe_3_result, 3, "test_uprobe_3_result");
> +
> +cleanup:
> +       uprobe_multi_session_cookie__destroy(skel);
> +}
> +
>  static void test_bench_attach_uprobe(void)
>  {
>         long attach_start_ns = 0, attach_end_ns = 0;
> @@ -745,4 +774,6 @@ void test_uprobe_multi_test(void)
>                 test_attach_api_fails();
>         if (test__start_subtest("session"))
>                 test_session_skel_api();
> +       if (test__start_subtest("session_cookie"))
> +               test_session_cookie_skel_api();
>  }
> diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
> new file mode 100644
> index 000000000000..5befdf944dc6
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
> @@ -0,0 +1,48 @@
> +// SPDX-License-Identifier: GPL-2.0
> +#include <linux/bpf.h>
> +#include <bpf/bpf_helpers.h>
> +#include <bpf/bpf_tracing.h>
> +#include <stdbool.h>
> +#include "bpf_kfuncs.h"
> +
> +char _license[] SEC("license") = "GPL";
> +
> +int pid = 0;
> +
> +__u64 test_uprobe_1_result = 0;
> +__u64 test_uprobe_2_result = 0;
> +__u64 test_uprobe_3_result = 0;
> +
> +static int check_cookie(__u64 val, __u64 *result)
> +{
> +       __u64 *cookie;
> +
> +       if (bpf_get_current_pid_tgid() >> 32 != pid)
> +               return 1;
> +
> +       cookie = bpf_session_cookie();
> +
> +       if (bpf_session_is_return())
> +               *result = *cookie == val ? val : 0;
> +       else
> +               *cookie = val;
> +       return 0;
> +}
> +
> +SEC("uprobe.session//proc/self/exe:uprobe_multi_func_1")
> +int uprobe_1(struct pt_regs *ctx)
> +{
> +       return check_cookie(1, &test_uprobe_1_result);
> +}
> +
> +SEC("uprobe.session//proc/self/exe:uprobe_multi_func_2")
> +int uprobe_2(struct pt_regs *ctx)
> +{
> +       return check_cookie(2, &test_uprobe_2_result);
> +}
> +
> +SEC("uprobe.session//proc/self/exe:uprobe_multi_func_3")
> +int uprobe_3(struct pt_regs *ctx)
> +{
> +       return check_cookie(3, &test_uprobe_3_result);
> +}
> --
> 2.45.2
>

Re: [PATCHv2 bpf-next 8/9] selftests/bpf: Add uprobe session recursive test

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding uprobe session test that verifies the cookie value is stored
> properly when single uprobe-ed function is executed recursively.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  .../bpf/prog_tests/uprobe_multi_test.c        | 57 +++++++++++++++++++
>  .../progs/uprobe_multi_session_recursive.c    | 44 ++++++++++++++
>  2 files changed, 101 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c
>

Nice!

Acked-by: Andrii Nakryiko <andrii@kernel.org>


[...]

> +static void test_session_recursive_skel_api(void)
> +{
> +       struct uprobe_multi_session_recursive *skel = NULL;
> +       int i, err;
> +
> +       skel = uprobe_multi_session_recursive__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_recursive__open_and_load"))
> +               goto cleanup;
> +
> +       skel->bss->pid = getpid();
> +
> +       err = uprobe_multi_session_recursive__attach(skel);
> +       if (!ASSERT_OK(err, "uprobe_multi_session_recursive__attach"))
> +               goto cleanup;
> +
> +       for (i = 0; i < ARRAY_SIZE(skel->bss->test_uprobe_cookie_entry); i++)
> +               skel->bss->test_uprobe_cookie_entry[i] = i + 1;
> +
> +       uprobe_session_recursive(5);
> +
> +       /*

nit: unnecessary empty comment line

> +        *                                         entry uprobe:
> +        * uprobe_session_recursive(5) {             *cookie = 1, return 0
> +        *   uprobe_session_recursive(4) {           *cookie = 2, return 1
> +        *     uprobe_session_recursive(3) {         *cookie = 3, return 0
> +        *       uprobe_session_recursive(2) {       *cookie = 4, return 1
> +        *         uprobe_session_recursive(1) {     *cookie = 5, return 0
> +        *           uprobe_session_recursive(0) {   *cookie = 6, return 1
> +        *                                          return uprobe:
> +        *           } i = 0                          not executed
> +        *         } i = 1                            test_uprobe_cookie_return[0] = 5
> +        *       } i = 2                              not executed
> +        *     } i = 3                                test_uprobe_cookie_return[1] = 3
> +        *   } i = 4                                  not executed
> +        * } i = 5                                    test_uprobe_cookie_return[2] = 1
> +        */
> +

[...]

Re: [PATCHv2 bpf-next 9/9] selftests/bpf: Add uprobe session consumers test

[Andrii Nakryiko]

On Mon, Jul 1, 2024 at 9:44 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding test that attached/detaches multiple consumers on
> single uprobe and verifies all were hit as expected.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  .../bpf/prog_tests/uprobe_multi_test.c        | 203 ++++++++++++++++++
>  .../progs/uprobe_multi_session_consumers.c    |  53 +++++
>  2 files changed, 256 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
>

This is clever, though bit notation obscures the meaning of the code a
bit. But thanks for the long comment explaining the overall idea.

> diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> index b521590fdbb9..83eac954cf00 100644
> --- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> +++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> @@ -9,6 +9,7 @@
>  #include "uprobe_multi_session.skel.h"
>  #include "uprobe_multi_session_cookie.skel.h"
>  #include "uprobe_multi_session_recursive.skel.h"
> +#include "uprobe_multi_session_consumers.skel.h"
>  #include "bpf/libbpf_internal.h"
>  #include "testing_helpers.h"
>  #include "../sdt.h"
> @@ -739,6 +740,206 @@ static void test_session_recursive_skel_api(void)
>         uprobe_multi_session_recursive__destroy(skel);
>  }
>
> +static int uprobe_attach(struct uprobe_multi_session_consumers *skel, int bit)
> +{
> +       struct bpf_program **prog = &skel->progs.uprobe_0 + bit;
> +       struct bpf_link **link = &skel->links.uprobe_0 + bit;
> +       LIBBPF_OPTS(bpf_uprobe_multi_opts, opts);
> +
> +       /*
> +        * bit: 0,1 uprobe session
> +        * bit: 2,3 uprobe entry
> +        * bit: 4,5 uprobe return
> +        */
> +       opts.session = bit < 2;
> +       opts.retprobe = bit == 4 || bit == 5;
> +
> +       *link = bpf_program__attach_uprobe_multi(*prog, 0, "/proc/self/exe",
> +                                                "uprobe_session_consumer_test",
> +                                                &opts);
> +       if (!ASSERT_OK_PTR(*link, "bpf_program__attach_uprobe_multi"))
> +               return -1;
> +       return 0;
> +}
> +
> +static void uprobe_detach(struct uprobe_multi_session_consumers *skel, int bit)
> +{
> +       struct bpf_link **link = &skel->links.uprobe_0 + bit;

ok, this is nasty, no one guarantees this should keep working,
explicit switch would be preferable

> +
> +       bpf_link__destroy(*link);
> +       *link = NULL;
> +}
> +
> +static bool test_bit(int bit, unsigned long val)
> +{
> +       return val & (1 << bit);
> +}
> +
> +noinline int
> +uprobe_session_consumer_test(struct uprobe_multi_session_consumers *skel,
> +                            unsigned long before, unsigned long after)
> +{
> +       int bit;
> +
> +       /* detach uprobe for each unset bit in 'before' state ... */
> +       for (bit = 0; bit < 6; bit++) {

Does "bit" correspond to the uprobe_X program? Maybe call it an uprobe
index or something, if that's the case? bits are just representations,
but semantically meaningful is identifier of an uprobe program, right?

> +               if (test_bit(bit, before) && !test_bit(bit, after))
> +                       uprobe_detach(skel, bit);
> +       }
> +
> +       /* ... and attach all new bits in 'after' state */
> +       for (bit = 0; bit < 6; bit++) {
> +               if (!test_bit(bit, before) && test_bit(bit, after)) {
> +                       if (!ASSERT_OK(uprobe_attach(skel, bit), "uprobe_attach_after"))
> +                               return -1;
> +               }
> +       }
> +       return 0;
> +}
> +

[...]

> +
> +static void test_session_consumers(void)
> +{
> +       struct uprobe_multi_session_consumers *skel;
> +       int before, after;
> +
> +       skel = uprobe_multi_session_consumers__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_consumers__open_and_load"))
> +               return;
> +
> +       /*
> +        * The idea of this test is to try all possible combinations of
> +        * uprobes consumers attached on single function.
> +        *
> +        *  - 1 uprobe session with return handler called
> +        *  - 1 uprobe session without return handler called
> +        *  - 2 uprobe entry consumer
> +        *  - 2 uprobe exit consumers
> +        *
> +        * The test uses 6 uprobes attached on single function, but that
> +        * translates into single uprobe with 6 consumers in kernel.
> +        *
> +        * The before/after values present the state of attached consumers
> +        * before and after the probed function:
> +        *
> +        *  bit 0   : uprobe session with return
> +        *  bit 1   : uprobe session with no return
> +        *  bit 2,3 : uprobe entry
> +        *  bit 4,5 : uprobe return
> +        *
> +        * For example for:
> +        *
> +        *   before = 0b10101
> +        *   after  = 0b00110
> +        *
> +        * it means that before we call 'uprobe_session_consumer_test' we
> +        * attach uprobes defined in 'before' value:
> +        *
> +        *   - bit 0: uprobe session with return
> +        *   - bit 2: uprobe entry
> +        *   - bit 4: uprobe return
> +        *
> +        * uprobe_session_consumer_test is called and inside it we attach
> +        * and detach * uprobes based on 'after' value:
> +        *
> +        *   - bit 0: uprobe session with return is detached
> +        *   - bit 1: uprobe session without return is attached
> +        *   - bit 2: stays untouched
> +        *   - bit 4: uprobe return is detached
> +        *
> +        * uprobe_session_consumer_test returs and we check counters values
> +        * increased by bpf programs on each uprobe to match the expected
> +        * count based on before/after bits.
> +        */
> +       for (before = 0; before < 64; before++) {
> +               for (after = 0; after < 64; after++)
> +                       session_consumer_test(skel, before, after);
> +       }
> +
> +       uprobe_multi_session_consumers__destroy(skel);
> +}
> +

[...]

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Masami Hiramatsu]

Hi Jiri,

On Mon,  1 Jul 2024 18:41:07 +0200
Jiri Olsa <jolsa@kernel.org> wrote:

> Adding support for uprobe consumer to be defined as session and have
> new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
> 
> The session means that 'handler' and 'ret_handler' callbacks are
> connected in a way that allows to:
> 
>   - control execution of 'ret_handler' from 'handler' callback
>   - share data between 'handler' and 'ret_handler' callbacks
> 
> The session is enabled by setting new 'session' bool field to true
> in uprobe_consumer object.
> 
> We keep count of session consumers for uprobe and allocate session_consumer
> object for each in return_instance object. This allows us to store
> return values of 'handler' callbacks and data pointers of shared
> data between both handlers.
> 
> The session concept fits to our common use case where we do filtering
> on entry uprobe and based on the result we decide to run the return
> uprobe (or not).
> 
> It's also convenient to share the data between session callbacks.
> 
> The control of 'ret_handler' callback execution is done via return
> value of the 'handler' callback. If it's 0 we install and execute
> return uprobe, if it's 1 we do not.
> 
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/linux/uprobes.h     |  16 ++++-
>  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
>  kernel/trace/bpf_trace.c    |   6 +-
>  kernel/trace/trace_uprobe.c |  12 ++--
>  4 files changed, 144 insertions(+), 19 deletions(-)
> 
> diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> index f46e0ca0169c..903a860a8d01 100644
> --- a/include/linux/uprobes.h
> +++ b/include/linux/uprobes.h
> @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
>  };
>  
>  struct uprobe_consumer {
> -	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> +	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
>  	int (*ret_handler)(struct uprobe_consumer *self,
>  				unsigned long func,
> -				struct pt_regs *regs);
> +				struct pt_regs *regs, __u64 *data);
>  	bool (*filter)(struct uprobe_consumer *self,
>  				enum uprobe_filter_ctx ctx,
>  				struct mm_struct *mm);
>  
>  	struct uprobe_consumer *next;
> +
> +	bool			session;	/* marks uprobe session consumer */
> +	unsigned int		session_id;	/* set when uprobe_consumer is registered */

Hmm, why this has both session and session_id?
I also think we can use the address of uprobe_consumer itself as a unique id.

Also, if we can set session enabled by default, and skip ret_handler by handler's
return value, it is more simpler. (If handler returns a specific value, skip ret_handler)

>  };
>  
>  #ifdef CONFIG_UPROBES
> @@ -80,6 +83,12 @@ struct uprobe_task {
>  	unsigned int			depth;
>  };
>  
> +struct session_consumer {
> +	__u64		cookie;

And this cookie looks not scalable. If we can pass a data to handler, I would like to
reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.

	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);

uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
at handler_chain().

> +	unsigned int	id;
> +	int		rc;
> +};
> +
>  struct return_instance {
>  	struct uprobe		*uprobe;
>  	unsigned long		func;
> @@ -88,6 +97,9 @@ struct return_instance {
>  	bool			chained;	/* true, if instance is nested */
>  
>  	struct return_instance	*next;		/* keep as stack */
> +
> +	int			sessions_cnt;
> +	struct session_consumer	sessions[];

In that case, we don't have this array, but 

	char data[];

And decode data array, which is a slice of variable length structure;

struct session_consumer {
	struct uprobe_consumer *uc;
	char data[];
};

The size of session_consumer is uc->session_data_size + sizeof(uc).

What would you think?

Thank you,

>  };
>  
>  enum rp_check {
> diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> index 2c83ba776fc7..4da410460f2a 100644
> --- a/kernel/events/uprobes.c
> +++ b/kernel/events/uprobes.c
> @@ -63,6 +63,8 @@ struct uprobe {
>  	loff_t			ref_ctr_offset;
>  	unsigned long		flags;
>  
> +	unsigned int		sessions_cnt;
> +
>  	/*
>  	 * The generic code assumes that it has two members of unknown type
>  	 * owned by the arch-specific code:
> @@ -750,11 +752,30 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
>  	return uprobe;
>  }
>  
> +static void
> +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> +{
> +	static unsigned int session_id;
> +
> +	if (uc->session) {
> +		uprobe->sessions_cnt++;
> +		uc->session_id = ++session_id ?: ++session_id;
> +	}
> +}
> +
> +static void
> +uprobe_consumer_unaccount(struct uprobe *uprobe, struct uprobe_consumer *uc)
> +{
> +	if (uc->session)
> +		uprobe->sessions_cnt--;
> +}
> +
>  static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
>  {
>  	down_write(&uprobe->consumer_rwsem);
>  	uc->next = uprobe->consumers;
>  	uprobe->consumers = uc;
> +	uprobe_consumer_account(uprobe, uc);
>  	up_write(&uprobe->consumer_rwsem);
>  }
>  
> @@ -773,6 +794,7 @@ static bool consumer_del(struct uprobe *uprobe, struct uprobe_consumer *uc)
>  		if (*con == uc) {
>  			*con = uc->next;
>  			ret = true;
> +			uprobe_consumer_unaccount(uprobe, uc);
>  			break;
>  		}
>  	}
> @@ -1744,6 +1766,23 @@ static struct uprobe_task *get_utask(void)
>  	return current->utask;
>  }
>  
> +static size_t ri_size(int sessions_cnt)
> +{
> +	struct return_instance *ri __maybe_unused;
> +
> +	return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
> +}
> +
> +static struct return_instance *alloc_return_instance(int sessions_cnt)
> +{
> +	struct return_instance *ri;
> +
> +	ri = kzalloc(ri_size(sessions_cnt), GFP_KERNEL);
> +	if (ri)
> +		ri->sessions_cnt = sessions_cnt;
> +	return ri;
> +}
> +
>  static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
>  {
>  	struct uprobe_task *n_utask;
> @@ -1756,11 +1795,11 @@ static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
>  
>  	p = &n_utask->return_instances;
>  	for (o = o_utask->return_instances; o; o = o->next) {
> -		n = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> +		n = alloc_return_instance(o->sessions_cnt);
>  		if (!n)
>  			return -ENOMEM;
>  
> -		*n = *o;
> +		memcpy(n, o, ri_size(o->sessions_cnt));
>  		get_uprobe(n->uprobe);
>  		n->next = NULL;
>  
> @@ -1853,9 +1892,9 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
>  	utask->return_instances = ri;
>  }
>  
> -static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
> +static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
> +			      struct return_instance *ri)
>  {
> -	struct return_instance *ri;
>  	struct uprobe_task *utask;
>  	unsigned long orig_ret_vaddr, trampoline_vaddr;
>  	bool chained;
> @@ -1874,9 +1913,11 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
>  		return;
>  	}
>  
> -	ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> -	if (!ri)
> -		return;
> +	if (!ri) {
> +		ri = alloc_return_instance(0);
> +		if (!ri)
> +			return;
> +	}
>  
>  	trampoline_vaddr = get_trampoline_vaddr();
>  	orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
> @@ -2065,35 +2106,85 @@ static struct uprobe *find_active_uprobe(unsigned long bp_vaddr, int *is_swbp)
>  	return uprobe;
>  }
>  
> +static struct session_consumer *
> +session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
> +		      int session_id)
> +{
> +	struct session_consumer *next;
> +
> +	next = sc ? sc + 1 : &ri->sessions[0];
> +	next->id = session_id;
> +	return next;
> +}
> +
> +static struct session_consumer *
> +session_consumer_find(struct return_instance *ri, int *iter, int session_id)
> +{
> +	struct session_consumer *sc;
> +	int idx = *iter;
> +
> +	for (sc = &ri->sessions[idx]; idx < ri->sessions_cnt; idx++, sc++) {
> +		if (sc->id == session_id) {
> +			*iter = idx + 1;
> +			return sc;
> +		}
> +	}
> +	return NULL;
> +}
> +
>  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
>  {
>  	struct uprobe_consumer *uc;
>  	int remove = UPROBE_HANDLER_REMOVE;
> +	struct session_consumer *sc = NULL;
> +	struct return_instance *ri = NULL;
>  	bool need_prep = false; /* prepare return uprobe, when needed */
>  
>  	down_read(&uprobe->register_rwsem);
> +	if (uprobe->sessions_cnt) {
> +		ri = alloc_return_instance(uprobe->sessions_cnt);
> +		if (!ri)
> +			goto out;
> +	}
> +
>  	for (uc = uprobe->consumers; uc; uc = uc->next) {
> +		__u64 *cookie = NULL;
>  		int rc = 0;
>  
> +		if (uc->session) {
> +			sc = session_consumer_next(ri, sc, uc->session_id);
> +			cookie = &sc->cookie;
> +		}
> +
>  		if (uc->handler) {
> -			rc = uc->handler(uc, regs);
> +			rc = uc->handler(uc, regs, cookie);
>  			WARN(rc & ~UPROBE_HANDLER_MASK,
>  				"bad rc=0x%x from %ps()\n", rc, uc->handler);
>  		}
>  
> -		if (uc->ret_handler)
> +		if (uc->session) {
> +			sc->rc = rc;
> +			need_prep |= !rc;
> +		} else if (uc->ret_handler) {
>  			need_prep = true;
> +		}
>  
>  		remove &= rc;
>  	}
>  
> +	/* no removal if there's at least one session consumer */
> +	remove &= !uprobe->sessions_cnt;
> +
>  	if (need_prep && !remove)
> -		prepare_uretprobe(uprobe, regs); /* put bp at return */
> +		prepare_uretprobe(uprobe, regs, ri); /* put bp at return */
> +	else
> +		kfree(ri);
>  
>  	if (remove && uprobe->consumers) {
>  		WARN_ON(!uprobe_is_active(uprobe));
>  		unapply_uprobe(uprobe, current->mm);
>  	}
> + out:
>  	up_read(&uprobe->register_rwsem);
>  }
>  
> @@ -2101,12 +2192,28 @@ static void
>  handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
>  {
>  	struct uprobe *uprobe = ri->uprobe;
> +	struct session_consumer *sc;
>  	struct uprobe_consumer *uc;
> +	int session_iter = 0;
>  
>  	down_read(&uprobe->register_rwsem);
>  	for (uc = uprobe->consumers; uc; uc = uc->next) {
> +		__u64 *cookie = NULL;
> +
> +		if (uc->session) {
> +			/*
> +			 * Consumers could be added and removed, but they will not
> +			 * change position, so we can iterate sessions just once
> +			 * and keep the last found session as base for next search.
> +			 */
> +			sc = session_consumer_find(ri, &session_iter, uc->session_id);
> +			if (!sc || sc->rc)
> +				continue;
> +			cookie = &sc->cookie;
> +		}
> +
>  		if (uc->ret_handler)
> -			uc->ret_handler(uc, ri->func, regs);
> +			uc->ret_handler(uc, ri->func, regs, cookie);
>  	}
>  	up_read(&uprobe->register_rwsem);
>  }
> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index cd098846e251..02d052639dfe 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -3332,7 +3332,8 @@ uprobe_multi_link_filter(struct uprobe_consumer *con, enum uprobe_filter_ctx ctx
>  }
>  
>  static int
> -uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs)
> +uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs,
> +			  __u64 *data)
>  {
>  	struct bpf_uprobe *uprobe;
>  
> @@ -3341,7 +3342,8 @@ uprobe_multi_link_handler(struct uprobe_consumer *con, struct pt_regs *regs)
>  }
>  
>  static int
> -uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, struct pt_regs *regs)
> +uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, struct pt_regs *regs,
> +			      __u64 *data)
>  {
>  	struct bpf_uprobe *uprobe;
>  
> diff --git a/kernel/trace/trace_uprobe.c b/kernel/trace/trace_uprobe.c
> index c98e3b3386ba..7068c279a244 100644
> --- a/kernel/trace/trace_uprobe.c
> +++ b/kernel/trace/trace_uprobe.c
> @@ -88,9 +88,11 @@ static struct trace_uprobe *to_trace_uprobe(struct dyn_event *ev)
>  static int register_uprobe_event(struct trace_uprobe *tu);
>  static int unregister_uprobe_event(struct trace_uprobe *tu);
>  
> -static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs);
> +static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs,
> +			     __u64 *data);
>  static int uretprobe_dispatcher(struct uprobe_consumer *con,
> -				unsigned long func, struct pt_regs *regs);
> +				unsigned long func, struct pt_regs *regs,
> +				__u64 *data);
>  
>  #ifdef CONFIG_STACK_GROWSUP
>  static unsigned long adjust_stack_addr(unsigned long addr, unsigned int n)
> @@ -1504,7 +1506,8 @@ trace_uprobe_register(struct trace_event_call *event, enum trace_reg type,
>  	}
>  }
>  
> -static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs)
> +static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs,
> +			     __u64 *data)
>  {
>  	struct trace_uprobe *tu;
>  	struct uprobe_dispatch_data udd;
> @@ -1534,7 +1537,8 @@ static int uprobe_dispatcher(struct uprobe_consumer *con, struct pt_regs *regs)
>  }
>  
>  static int uretprobe_dispatcher(struct uprobe_consumer *con,
> -				unsigned long func, struct pt_regs *regs)
> +				unsigned long func, struct pt_regs *regs,
> +				__u64 *data)
>  {
>  	struct trace_uprobe *tu;
>  	struct uprobe_dispatch_data udd;
> -- 
> 2.45.2
> 


-- 
Masami Hiramatsu (Google) <mhiramat@kernel.org>

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Tue, Jul 2, 2024 at 4:55 PM Masami Hiramatsu <mhiramat@kernel.org> wrote:
>
> Hi Jiri,
>
> On Mon,  1 Jul 2024 18:41:07 +0200
> Jiri Olsa <jolsa@kernel.org> wrote:
>
> > Adding support for uprobe consumer to be defined as session and have
> > new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
> >
> > The session means that 'handler' and 'ret_handler' callbacks are
> > connected in a way that allows to:
> >
> >   - control execution of 'ret_handler' from 'handler' callback
> >   - share data between 'handler' and 'ret_handler' callbacks
> >
> > The session is enabled by setting new 'session' bool field to true
> > in uprobe_consumer object.
> >
> > We keep count of session consumers for uprobe and allocate session_consumer
> > object for each in return_instance object. This allows us to store
> > return values of 'handler' callbacks and data pointers of shared
> > data between both handlers.
> >
> > The session concept fits to our common use case where we do filtering
> > on entry uprobe and based on the result we decide to run the return
> > uprobe (or not).
> >
> > It's also convenient to share the data between session callbacks.
> >
> > The control of 'ret_handler' callback execution is done via return
> > value of the 'handler' callback. If it's 0 we install and execute
> > return uprobe, if it's 1 we do not.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  include/linux/uprobes.h     |  16 ++++-
> >  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
> >  kernel/trace/bpf_trace.c    |   6 +-
> >  kernel/trace/trace_uprobe.c |  12 ++--
> >  4 files changed, 144 insertions(+), 19 deletions(-)
> >
> > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > index f46e0ca0169c..903a860a8d01 100644
> > --- a/include/linux/uprobes.h
> > +++ b/include/linux/uprobes.h
> > @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
> >  };
> >
> >  struct uprobe_consumer {
> > -     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> > +     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
> >       int (*ret_handler)(struct uprobe_consumer *self,
> >                               unsigned long func,
> > -                             struct pt_regs *regs);
> > +                             struct pt_regs *regs, __u64 *data);
> >       bool (*filter)(struct uprobe_consumer *self,
> >                               enum uprobe_filter_ctx ctx,
> >                               struct mm_struct *mm);
> >
> >       struct uprobe_consumer *next;
> > +
> > +     bool                    session;        /* marks uprobe session consumer */
> > +     unsigned int            session_id;     /* set when uprobe_consumer is registered */
>
> Hmm, why this has both session and session_id?

session is caller's request to establish session semantics. Jiri, I
think it's better to move it higher next to
handler/ret_handler/filter, that's the part of uprobe_consumer struct
which has read-only caller-provided data (I'm adding offset and
ref_ctr_offset there as well).

> I also think we can use the address of uprobe_consumer itself as a unique id.

+1

>
> Also, if we can set session enabled by default, and skip ret_handler by handler's
> return value, it is more simpler. (If handler returns a specific value, skip ret_handler)

you mean derive if it's a session or not by both handler and
ret_handler being set? I guess this works fine for BPF side, because
there we never had them both set. If this doesn't regress others, I
think it's OK. We just need to make sure we don't unnecessarily
allocate session state for consumers that don't set both handler and
ret_handler. That would be a waste.

>
> >  };
> >
> >  #ifdef CONFIG_UPROBES
> > @@ -80,6 +83,12 @@ struct uprobe_task {
> >       unsigned int                    depth;
> >  };
> >
> > +struct session_consumer {
> > +     __u64           cookie;
>
> And this cookie looks not scalable. If we can pass a data to handler, I would like to
> reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.
>
>         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);
>
> uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
> at handler_chain().

The goal here is to keep this simple and fast. I'd prefer to keep it
small and fixed size, if possible. I'm thinking about caching and
reusing return_instance as one of the future optimizations, so if we
can keep this more or less fixed (assuming there is typically not more
than 1 or 2 consumers per uprobe, which seems realistic), this will
provide a way to avoid excessive memory allocations.

>
> > +     unsigned int    id;
> > +     int             rc;
> > +};
> > +
> >  struct return_instance {
> >       struct uprobe           *uprobe;
> >       unsigned long           func;
> > @@ -88,6 +97,9 @@ struct return_instance {
> >       bool                    chained;        /* true, if instance is nested */
> >
> >       struct return_instance  *next;          /* keep as stack */
> > +
> > +     int                     sessions_cnt;
> > +     struct session_consumer sessions[];
>
> In that case, we don't have this array, but
>
>         char data[];
>
> And decode data array, which is a slice of variable length structure;
>
> struct session_consumer {
>         struct uprobe_consumer *uc;
>         char data[];
> };
>
> The size of session_consumer is uc->session_data_size + sizeof(uc).
>
> What would you think?
>
> Thank you,
>
> >  };
> >
> >  enum rp_check {
> > diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> > index 2c83ba776fc7..4da410460f2a 100644
> > --- a/kernel/events/uprobes.c
> > +++ b/kernel/events/uprobes.c
> > @@ -63,6 +63,8 @@ struct uprobe {
> >       loff_t                  ref_ctr_offset;
> >       unsigned long           flags;
> >
> > +     unsigned int            sessions_cnt;

[...]

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Peter Zijlstra]

On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:
> > +static size_t ri_size(int sessions_cnt)
> > +{
> > +       struct return_instance *ri __maybe_unused;
> > +
> > +       return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
> 
> just use struct_size()?

Yeah, lets not. This is readable, struct_size() is not.

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Tue, Jul 02, 2024 at 01:52:38PM -0700, Andrii Nakryiko wrote:
> On Tue, Jul 2, 2024 at 9:11 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> >
> > On Tue, Jul 02, 2024 at 03:04:08PM +0200, Peter Zijlstra wrote:
> > > On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:
> > >
> > > > +static void
> > > > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > > > +{
> > > > +   static unsigned int session_id;
> > > > +
> > > > +   if (uc->session) {
> > > > +           uprobe->sessions_cnt++;
> > > > +           uc->session_id = ++session_id ?: ++session_id;
> > > > +   }
> > > > +}
> > >
> > > The way I understand this code, you create a consumer every time you do
> > > uprobe_register() and unregister makes it go away.
> > >
> > > Now, register one, then 4g-1 times register+unregister, then register
> > > again.
> > >
> > > The above seems to then result in two consumers with the same
> > > session_id, which leads to trouble.
> > >
> > > Hmm?
> >
> > ugh true.. will make it u64 :)
> >
> > I think we could store uprobe_consumer pointer+ref in session_consumer,
> > and that would make the unregister path more interesting.. will check
> 
> More interesting how? It's actually a great idea, uprobe_consumer

nah, got confused ;-)

> pointer itself is a unique ID and 64-bit. We can still use lowest bit
> for RC (see my other reply).

I used pointers in the previous version, but then I thought what if the
consumer gets free-ed and new one created (with same address.. maybe not
likely but possible, right?) before the return probe is hit

jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Tue, Jul 02, 2024 at 05:13:38PM -0700, Andrii Nakryiko wrote:
> On Tue, Jul 2, 2024 at 4:55 PM Masami Hiramatsu <mhiramat@kernel.org> wrote:
> >
> > Hi Jiri,
> >
> > On Mon,  1 Jul 2024 18:41:07 +0200
> > Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > > Adding support for uprobe consumer to be defined as session and have
> > > new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
> > >
> > > The session means that 'handler' and 'ret_handler' callbacks are
> > > connected in a way that allows to:
> > >
> > >   - control execution of 'ret_handler' from 'handler' callback
> > >   - share data between 'handler' and 'ret_handler' callbacks
> > >
> > > The session is enabled by setting new 'session' bool field to true
> > > in uprobe_consumer object.
> > >
> > > We keep count of session consumers for uprobe and allocate session_consumer
> > > object for each in return_instance object. This allows us to store
> > > return values of 'handler' callbacks and data pointers of shared
> > > data between both handlers.
> > >
> > > The session concept fits to our common use case where we do filtering
> > > on entry uprobe and based on the result we decide to run the return
> > > uprobe (or not).
> > >
> > > It's also convenient to share the data between session callbacks.
> > >
> > > The control of 'ret_handler' callback execution is done via return
> > > value of the 'handler' callback. If it's 0 we install and execute
> > > return uprobe, if it's 1 we do not.
> > >
> > > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > > ---
> > >  include/linux/uprobes.h     |  16 ++++-
> > >  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
> > >  kernel/trace/bpf_trace.c    |   6 +-
> > >  kernel/trace/trace_uprobe.c |  12 ++--
> > >  4 files changed, 144 insertions(+), 19 deletions(-)
> > >
> > > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > > index f46e0ca0169c..903a860a8d01 100644
> > > --- a/include/linux/uprobes.h
> > > +++ b/include/linux/uprobes.h
> > > @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
> > >  };
> > >
> > >  struct uprobe_consumer {
> > > -     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> > > +     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
> > >       int (*ret_handler)(struct uprobe_consumer *self,
> > >                               unsigned long func,
> > > -                             struct pt_regs *regs);
> > > +                             struct pt_regs *regs, __u64 *data);
> > >       bool (*filter)(struct uprobe_consumer *self,
> > >                               enum uprobe_filter_ctx ctx,
> > >                               struct mm_struct *mm);
> > >
> > >       struct uprobe_consumer *next;
> > > +
> > > +     bool                    session;        /* marks uprobe session consumer */
> > > +     unsigned int            session_id;     /* set when uprobe_consumer is registered */
> >
> > Hmm, why this has both session and session_id?
> 
> session is caller's request to establish session semantics. Jiri, I

and session_id is set when uprobe is registered and used when
return uprobe is executed to find matching uprobe_consumer,
plz check handle_uretprobe_chain/session_consumer_find

> think it's better to move it higher next to
> handler/ret_handler/filter, that's the part of uprobe_consumer struct
> which has read-only caller-provided data (I'm adding offset and
> ref_ctr_offset there as well).

ok, makes sense

> 
> > I also think we can use the address of uprobe_consumer itself as a unique id.
> 
> +1
> 
> >
> > Also, if we can set session enabled by default, and skip ret_handler by handler's
> > return value, it is more simpler. (If handler returns a specific value, skip ret_handler)
> 
> you mean derive if it's a session or not by both handler and
> ret_handler being set? I guess this works fine for BPF side, because
> there we never had them both set. If this doesn't regress others, I
> think it's OK. We just need to make sure we don't unnecessarily
> allocate session state for consumers that don't set both handler and
> ret_handler. That would be a waste.

hum.. so the current code installs return uprobe if there's ret_handler
defined in consumer and also the entry 'handler' needs to return 0

if entry 'handler' returns 1 the uprobe is unregistered

we could define new return value from 'handler' to 'not execute the
'ret_handler' and have 'handler' return values:

  0 - execute 'ret_handler' if defined
  1 - remove the uprobe
  2 - do NOT execute 'ret_handler'  // this current triggers WARN

we could delay the allocation of 'return_instance' until the first
consumer returns 0, so there's no perf regression

that way we could treat all consumers the same and we wouldn't need
the session flag..

ok looks like good idea ;-) will try that

thanks,
jirka

> 
> >
> > >  };
> > >
> > >  #ifdef CONFIG_UPROBES
> > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > >       unsigned int                    depth;
> > >  };
> > >
> > > +struct session_consumer {
> > > +     __u64           cookie;
> >
> > And this cookie looks not scalable. If we can pass a data to handler, I would like to
> > reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.
> >
> >         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);
> >
> > uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
> > at handler_chain().
> 
> The goal here is to keep this simple and fast. I'd prefer to keep it
> small and fixed size, if possible. I'm thinking about caching and
> reusing return_instance as one of the future optimizations, so if we
> can keep this more or less fixed (assuming there is typically not more
> than 1 or 2 consumers per uprobe, which seems realistic), this will
> provide a way to avoid excessive memory allocations.
> 
> >
> > > +     unsigned int    id;
> > > +     int             rc;
> > > +};
> > > +
> > >  struct return_instance {
> > >       struct uprobe           *uprobe;
> > >       unsigned long           func;
> > > @@ -88,6 +97,9 @@ struct return_instance {
> > >       bool                    chained;        /* true, if instance is nested */
> > >
> > >       struct return_instance  *next;          /* keep as stack */
> > > +
> > > +     int                     sessions_cnt;
> > > +     struct session_consumer sessions[];
> >
> > In that case, we don't have this array, but
> >
> >         char data[];
> >
> > And decode data array, which is a slice of variable length structure;
> >
> > struct session_consumer {
> >         struct uprobe_consumer *uc;
> >         char data[];
> > };
> >
> > The size of session_consumer is uc->session_data_size + sizeof(uc).
> >
> > What would you think?
> >
> > Thank you,
> >
> > >  };
> > >
> > >  enum rp_check {
> > > diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> > > index 2c83ba776fc7..4da410460f2a 100644
> > > --- a/kernel/events/uprobes.c
> > > +++ b/kernel/events/uprobes.c
> > > @@ -63,6 +63,8 @@ struct uprobe {
> > >       loff_t                  ref_ctr_offset;
> > >       unsigned long           flags;
> > >
> > > +     unsigned int            sessions_cnt;
> 
> [...]

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Wed, Jul 03, 2024 at 05:31:32PM +0200, Jiri Olsa wrote:
> On Tue, Jul 02, 2024 at 01:52:38PM -0700, Andrii Nakryiko wrote:
> > On Tue, Jul 2, 2024 at 9:11 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> > >
> > > On Tue, Jul 02, 2024 at 03:04:08PM +0200, Peter Zijlstra wrote:
> > > > On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:
> > > >
> > > > > +static void
> > > > > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > > > > +{
> > > > > +   static unsigned int session_id;
> > > > > +
> > > > > +   if (uc->session) {
> > > > > +           uprobe->sessions_cnt++;
> > > > > +           uc->session_id = ++session_id ?: ++session_id;
> > > > > +   }
> > > > > +}
> > > >
> > > > The way I understand this code, you create a consumer every time you do
> > > > uprobe_register() and unregister makes it go away.
> > > >
> > > > Now, register one, then 4g-1 times register+unregister, then register
> > > > again.
> > > >
> > > > The above seems to then result in two consumers with the same
> > > > session_id, which leads to trouble.
> > > >
> > > > Hmm?
> > >
> > > ugh true.. will make it u64 :)
> > >
> > > I think we could store uprobe_consumer pointer+ref in session_consumer,
> > > and that would make the unregister path more interesting.. will check
> > 
> > More interesting how? It's actually a great idea, uprobe_consumer
> 
> nah, got confused ;-)

actually.. the idea was to store uprobe_consumer pointers in 'return_instance'
and iterate them on return probe (not uprobe->consumers).. but that would
require the unregister code to somehow remove them (replace with null)

but we wouldn't need to do the search for matching consumer with session_id

also it probably regress current code, because we would execute only uprobe
return consumers that were registered at the time the function entry was hit,
whereas in current code the return uprobe executes all registered return
consumers

jirka

> 
> > pointer itself is a unique ID and 64-bit. We can still use lowest bit
> > for RC (see my other reply).
> 
> I used pointers in the previous version, but then I thought what if the
> consumer gets free-ed and new one created (with same address.. maybe not
> likely but possible, right?) before the return probe is hit
> 
> jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:

SNIP

> >  #ifdef CONFIG_UPROBES
> > @@ -80,6 +83,12 @@ struct uprobe_task {
> >         unsigned int                    depth;
> >  };
> >
> > +struct session_consumer {
> > +       __u64           cookie;
> > +       unsigned int    id;
> > +       int             rc;
> 
> you'll be using u64 for ID, right? so this struct will be 24 bytes.

yes

> Maybe we can just use topmost bit of ID to store whether uretprobe
> should run or not? It's trivial to mask out during ID comparisons

actually.. I think we could store just consumers that need to be
executed in return probe so there will be no need for 'rc' value

> 
> > +};
> > +
> >  struct return_instance {
> >         struct uprobe           *uprobe;
> >         unsigned long           func;
> > @@ -88,6 +97,9 @@ struct return_instance {
> >         bool                    chained;        /* true, if instance is nested */
> >
> >         struct return_instance  *next;          /* keep as stack */
> > +
> > +       int                     sessions_cnt;
> 
> there is 7 byte gap before next field, let's put sessions_cnt there

ok

> 
> > +       struct session_consumer sessions[];
> >  };
> >
> >  enum rp_check {
> > diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> > index 2c83ba776fc7..4da410460f2a 100644
> > --- a/kernel/events/uprobes.c
> > +++ b/kernel/events/uprobes.c
> > @@ -63,6 +63,8 @@ struct uprobe {
> >         loff_t                  ref_ctr_offset;
> >         unsigned long           flags;
> >
> > +       unsigned int            sessions_cnt;
> > +
> >         /*
> >          * The generic code assumes that it has two members of unknown type
> >          * owned by the arch-specific code:
> > @@ -750,11 +752,30 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
> >         return uprobe;
> >  }
> >
> > +static void
> > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > +{
> > +       static unsigned int session_id;
> 
> (besides what Peter mentioned about wrap around of 32-bit counter)
> let's use atomic here to not rely on any particular locking
> (unnecessarily), this might make my life easier in the future, thanks.
> This is registration time, low frequency, extra atomic won't hurt.
> 
> It might be already broken, actually, for two independently registering uprobes.

ok, will try

> 
> > +
> > +       if (uc->session) {
> > +               uprobe->sessions_cnt++;
> > +               uc->session_id = ++session_id ?: ++session_id;
> > +       }
> > +}
> > +
> > +static void
> > +uprobe_consumer_unaccount(struct uprobe *uprobe, struct uprobe_consumer *uc)
> 
> this fits in 100 characters, keep it single line, please. Same for
> account function

ok

> 
> > +{
> > +       if (uc->session)
> > +               uprobe->sessions_cnt--;
> > +}
> > +
> >  static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
> >  {
> >         down_write(&uprobe->consumer_rwsem);
> >         uc->next = uprobe->consumers;
> >         uprobe->consumers = uc;
> > +       uprobe_consumer_account(uprobe, uc);
> >         up_write(&uprobe->consumer_rwsem);
> >  }
> >
> > @@ -773,6 +794,7 @@ static bool consumer_del(struct uprobe *uprobe, struct uprobe_consumer *uc)
> >                 if (*con == uc) {
> >                         *con = uc->next;
> >                         ret = true;
> > +                       uprobe_consumer_unaccount(uprobe, uc);
> >                         break;
> >                 }
> >         }
> > @@ -1744,6 +1766,23 @@ static struct uprobe_task *get_utask(void)
> >         return current->utask;
> >  }
> >
> > +static size_t ri_size(int sessions_cnt)
> > +{
> > +       struct return_instance *ri __maybe_unused;
> > +
> > +       return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
> 
> just use struct_size()?
> 
> > +}
> > +
> > +static struct return_instance *alloc_return_instance(int sessions_cnt)
> > +{
> > +       struct return_instance *ri;
> > +
> > +       ri = kzalloc(ri_size(sessions_cnt), GFP_KERNEL);
> > +       if (ri)
> > +               ri->sessions_cnt = sessions_cnt;
> > +       return ri;
> > +}
> > +
> >  static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
> >  {
> >         struct uprobe_task *n_utask;
> > @@ -1756,11 +1795,11 @@ static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
> >
> >         p = &n_utask->return_instances;
> >         for (o = o_utask->return_instances; o; o = o->next) {
> > -               n = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> > +               n = alloc_return_instance(o->sessions_cnt);
> >                 if (!n)
> >                         return -ENOMEM;
> >
> > -               *n = *o;
> > +               memcpy(n, o, ri_size(o->sessions_cnt));
> >                 get_uprobe(n->uprobe);
> >                 n->next = NULL;
> >
> > @@ -1853,9 +1892,9 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
> >         utask->return_instances = ri;
> >  }
> >
> > -static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
> > +static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
> > +                             struct return_instance *ri)
> >  {
> > -       struct return_instance *ri;
> >         struct uprobe_task *utask;
> >         unsigned long orig_ret_vaddr, trampoline_vaddr;
> >         bool chained;
> > @@ -1874,9 +1913,11 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
> >                 return;
> >         }
> >
> > -       ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
> > -       if (!ri)
> > -               return;
> > +       if (!ri) {
> > +               ri = alloc_return_instance(0);
> > +               if (!ri)
> > +                       return;
> > +       }
> >
> >         trampoline_vaddr = get_trampoline_vaddr();
> >         orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
> > @@ -2065,35 +2106,85 @@ static struct uprobe *find_active_uprobe(unsigned long bp_vaddr, int *is_swbp)
> >         return uprobe;
> >  }
> >
> > +static struct session_consumer *
> > +session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
> > +                     int session_id)
> > +{
> > +       struct session_consumer *next;
> > +
> > +       next = sc ? sc + 1 : &ri->sessions[0];
> > +       next->id = session_id;
> 
> it's kind of unexpected that "session_consumer_next" would actually
> set an ID... Maybe drop int session_id as input argument and fill it
> out outside of this function, this function being just a simple
> iterator?

yea, I was going back and forth on what to have in that function
or not, to keep the change minimal, but makes sense, will move

> 
> > +       return next;
> > +}
> > +
> > +static struct session_consumer *
> > +session_consumer_find(struct return_instance *ri, int *iter, int session_id)
> > +{
> > +       struct session_consumer *sc;
> > +       int idx = *iter;
> > +
> > +       for (sc = &ri->sessions[idx]; idx < ri->sessions_cnt; idx++, sc++) {
> > +               if (sc->id == session_id) {
> > +                       *iter = idx + 1;
> > +                       return sc;
> > +               }
> > +       }
> > +       return NULL;
> > +}
> > +
> >  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
> >  {
> >         struct uprobe_consumer *uc;
> >         int remove = UPROBE_HANDLER_REMOVE;
> > +       struct session_consumer *sc = NULL;
> > +       struct return_instance *ri = NULL;
> >         bool need_prep = false; /* prepare return uprobe, when needed */
> >
> >         down_read(&uprobe->register_rwsem);
> > +       if (uprobe->sessions_cnt) {
> > +               ri = alloc_return_instance(uprobe->sessions_cnt);
> > +               if (!ri)
> > +                       goto out;
> > +       }
> > +
> >         for (uc = uprobe->consumers; uc; uc = uc->next) {
> > +               __u64 *cookie = NULL;
> >                 int rc = 0;
> >
> > +               if (uc->session) {
> > +                       sc = session_consumer_next(ri, sc, uc->session_id);
> > +                       cookie = &sc->cookie;
> > +               }
> > +
> >                 if (uc->handler) {
> > -                       rc = uc->handler(uc, regs);
> > +                       rc = uc->handler(uc, regs, cookie);
> >                         WARN(rc & ~UPROBE_HANDLER_MASK,
> >                                 "bad rc=0x%x from %ps()\n", rc, uc->handler);
> >                 }
> >
> > -               if (uc->ret_handler)
> > +               if (uc->session) {
> > +                       sc->rc = rc;
> > +                       need_prep |= !rc;
> 
> nit:
> 
> if (rc == 0)
>     need_prep = true;
> 
> and then it's *extremely obvious* what happens and under which conditions

ok

> 
> > +               } else if (uc->ret_handler) {
> >                         need_prep = true;
> > +               }
> >
> >                 remove &= rc;
> >         }
> >
> > +       /* no removal if there's at least one session consumer */
> > +       remove &= !uprobe->sessions_cnt;
> 
> this is counter (not error, not pointer), let's stick to ` == 0`, please
> 
> is this
> 
> if (uprobe->sessions_cnt != 0)
>    remove = 0;

yes ;-) will change

jirka

> 
> ? I can't tell (honestly), without spending ridiculous amounts of
> mental resources (for the underlying simplicity of the condition).

SNIP

Re: [PATCHv2 bpf-next 4/9] libbpf: Add support for uprobe multi session attach

[Jiri Olsa]

On Tue, Jul 02, 2024 at 02:34:34PM -0700, Andrii Nakryiko wrote:
> On Mon, Jul 1, 2024 at 9:42 AM Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > Adding support to attach program in uprobe session mode
> > with bpf_program__attach_uprobe_multi function.
> >
> > Adding session bool to bpf_uprobe_multi_opts struct that allows
> > to load and attach the bpf program via uprobe session.
> > the attachment to create uprobe multi session.
> >
> > Also adding new program loader section that allows:
> >   SEC("uprobe.session/bpf_fentry_test*")
> >
> > and loads/attaches uprobe program as uprobe session.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  tools/lib/bpf/bpf.c    |  1 +
> >  tools/lib/bpf/libbpf.c | 50 ++++++++++++++++++++++++++++++++++++++++--
> >  tools/lib/bpf/libbpf.h |  4 +++-
> >  3 files changed, 52 insertions(+), 3 deletions(-)
> >
> 
> [...]
> 
> > @@ -9362,6 +9363,7 @@ static const struct bpf_sec_def section_defs[] = {
> >         SEC_DEF("kprobe.session+",      KPROBE, BPF_TRACE_KPROBE_SESSION, SEC_NONE, attach_kprobe_session),
> >         SEC_DEF("uprobe.multi+",        KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
> >         SEC_DEF("uretprobe.multi+",     KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_NONE, attach_uprobe_multi),
> > +       SEC_DEF("uprobe.session+",      KPROBE, BPF_TRACE_UPROBE_SESSION, SEC_NONE, attach_uprobe_session),
> 
> sleepable ones as well?

ah right, forgot.. will add

> 
> >         SEC_DEF("uprobe.multi.s+",      KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
> >         SEC_DEF("uretprobe.multi.s+",   KPROBE, BPF_TRACE_UPROBE_MULTI, SEC_SLEEPABLE, attach_uprobe_multi),
> >         SEC_DEF("ksyscall+",            KPROBE, 0, SEC_NONE, attach_ksyscall),
> > @@ -11698,6 +11700,40 @@ static int attach_uprobe_multi(const struct bpf_program *prog, long cookie, stru
> >         return ret;
> >  }
> >
> > +static int attach_uprobe_session(const struct bpf_program *prog, long cookie, struct bpf_link **link)
> > +{
> > +       char *binary_path = NULL, *func_name = NULL;
> > +       LIBBPF_OPTS(bpf_uprobe_multi_opts, opts,
> > +               .session = true,
> > +       );
> 
> nit: keep a single line?

ok

> 
> > +       int n, ret = -EINVAL;
> > +       const char *spec;
> > +
> > +       *link = NULL;
> > +
> > +       spec = prog->sec_name + sizeof("uprobe.session/") - 1;
> > +       n = sscanf(spec, "%m[^:]:%m[^\n]",
> > +                  &binary_path, &func_name);
> 
> single line, wrapping lines is a necessary evil, please

ok

thanks,
jirka

Re: [PATCHv2 bpf-next 5/9] libbpf: Add uprobe session attach type names to attach_type_name

[Jiri Olsa]

On Tue, Jul 02, 2024 at 02:56:34PM -0700, Andrii Nakryiko wrote:
> On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > Adding uprobe session attach type name to attach_type_name,
> > so libbpf_bpf_attach_type_str returns proper string name for
> > BPF_TRACE_UPROBE_SESSION attach type.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  tools/lib/bpf/libbpf.c | 1 +
> >  1 file changed, 1 insertion(+)
> >
> 
> Can you merge this into a patch that adds BPF_TRACE_UPROBE_SESSION to
> keep bisectability of BPF selftests? It's a trivial patch, so
> shouldn't be a big deal.

ok

jirka

> 
> > diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
> > index 492a8eb4d047..e69a54264580 100644
> > --- a/tools/lib/bpf/libbpf.c
> > +++ b/tools/lib/bpf/libbpf.c
> > @@ -133,6 +133,7 @@ static const char * const attach_type_name[] = {
> >         [BPF_NETKIT_PRIMARY]            = "netkit_primary",
> >         [BPF_NETKIT_PEER]               = "netkit_peer",
> >         [BPF_TRACE_KPROBE_SESSION]      = "trace_kprobe_session",
> > +       [BPF_TRACE_UPROBE_SESSION]      = "trace_uprobe_session",
> >  };
> >
> >  static const char * const link_type_name[] = {
> > --
> > 2.45.2
> >

Re: [PATCHv2 bpf-next 8/9] selftests/bpf: Add uprobe session recursive test

[Jiri Olsa]

On Tue, Jul 02, 2024 at 03:01:35PM -0700, Andrii Nakryiko wrote:
> On Mon, Jul 1, 2024 at 9:43 AM Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > Adding uprobe session test that verifies the cookie value is stored
> > properly when single uprobe-ed function is executed recursively.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  .../bpf/prog_tests/uprobe_multi_test.c        | 57 +++++++++++++++++++
> >  .../progs/uprobe_multi_session_recursive.c    | 44 ++++++++++++++
> >  2 files changed, 101 insertions(+)
> >  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c
> >
> 
> Nice!
> 
> Acked-by: Andrii Nakryiko <andrii@kernel.org>
> 
> 
> [...]
> 
> > +static void test_session_recursive_skel_api(void)
> > +{
> > +       struct uprobe_multi_session_recursive *skel = NULL;
> > +       int i, err;
> > +
> > +       skel = uprobe_multi_session_recursive__open_and_load();
> > +       if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_recursive__open_and_load"))
> > +               goto cleanup;
> > +
> > +       skel->bss->pid = getpid();
> > +
> > +       err = uprobe_multi_session_recursive__attach(skel);
> > +       if (!ASSERT_OK(err, "uprobe_multi_session_recursive__attach"))
> > +               goto cleanup;
> > +
> > +       for (i = 0; i < ARRAY_SIZE(skel->bss->test_uprobe_cookie_entry); i++)
> > +               skel->bss->test_uprobe_cookie_entry[i] = i + 1;
> > +
> > +       uprobe_session_recursive(5);
> > +
> > +       /*
> 
> nit: unnecessary empty comment line

ok

jirka

> 
> > +        *                                         entry uprobe:
> > +        * uprobe_session_recursive(5) {             *cookie = 1, return 0
> > +        *   uprobe_session_recursive(4) {           *cookie = 2, return 1
> > +        *     uprobe_session_recursive(3) {         *cookie = 3, return 0
> > +        *       uprobe_session_recursive(2) {       *cookie = 4, return 1
> > +        *         uprobe_session_recursive(1) {     *cookie = 5, return 0
> > +        *           uprobe_session_recursive(0) {   *cookie = 6, return 1
> > +        *                                          return uprobe:
> > +        *           } i = 0                          not executed
> > +        *         } i = 1                            test_uprobe_cookie_return[0] = 5
> > +        *       } i = 2                              not executed
> > +        *     } i = 3                                test_uprobe_cookie_return[1] = 3
> > +        *   } i = 4                                  not executed
> > +        * } i = 5                                    test_uprobe_cookie_return[2] = 1
> > +        */
> > +
> 
> [...]

Re: [PATCHv2 bpf-next 9/9] selftests/bpf: Add uprobe session consumers test

[Jiri Olsa]

On Tue, Jul 02, 2024 at 03:10:55PM -0700, Andrii Nakryiko wrote:
> On Mon, Jul 1, 2024 at 9:44 AM Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > Adding test that attached/detaches multiple consumers on
> > single uprobe and verifies all were hit as expected.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  .../bpf/prog_tests/uprobe_multi_test.c        | 203 ++++++++++++++++++
> >  .../progs/uprobe_multi_session_consumers.c    |  53 +++++
> >  2 files changed, 256 insertions(+)
> >  create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
> >
> 
> This is clever, though bit notation obscures the meaning of the code a
> bit. But thanks for the long comment explaining the overall idea.
> 
> > diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> > index b521590fdbb9..83eac954cf00 100644
> > --- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> > +++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
> > @@ -9,6 +9,7 @@
> >  #include "uprobe_multi_session.skel.h"
> >  #include "uprobe_multi_session_cookie.skel.h"
> >  #include "uprobe_multi_session_recursive.skel.h"
> > +#include "uprobe_multi_session_consumers.skel.h"
> >  #include "bpf/libbpf_internal.h"
> >  #include "testing_helpers.h"
> >  #include "../sdt.h"
> > @@ -739,6 +740,206 @@ static void test_session_recursive_skel_api(void)
> >         uprobe_multi_session_recursive__destroy(skel);
> >  }
> >
> > +static int uprobe_attach(struct uprobe_multi_session_consumers *skel, int bit)
> > +{
> > +       struct bpf_program **prog = &skel->progs.uprobe_0 + bit;
> > +       struct bpf_link **link = &skel->links.uprobe_0 + bit;
> > +       LIBBPF_OPTS(bpf_uprobe_multi_opts, opts);
> > +
> > +       /*
> > +        * bit: 0,1 uprobe session
> > +        * bit: 2,3 uprobe entry
> > +        * bit: 4,5 uprobe return
> > +        */
> > +       opts.session = bit < 2;
> > +       opts.retprobe = bit == 4 || bit == 5;
> > +
> > +       *link = bpf_program__attach_uprobe_multi(*prog, 0, "/proc/self/exe",
> > +                                                "uprobe_session_consumer_test",
> > +                                                &opts);
> > +       if (!ASSERT_OK_PTR(*link, "bpf_program__attach_uprobe_multi"))
> > +               return -1;
> > +       return 0;
> > +}
> > +
> > +static void uprobe_detach(struct uprobe_multi_session_consumers *skel, int bit)
> > +{
> > +       struct bpf_link **link = &skel->links.uprobe_0 + bit;
> 
> ok, this is nasty, no one guarantees this should keep working,
> explicit switch would be preferable

I see, ok, will replace that with a switch

> 
> > +
> > +       bpf_link__destroy(*link);
> > +       *link = NULL;
> > +}
> > +
> > +static bool test_bit(int bit, unsigned long val)
> > +{
> > +       return val & (1 << bit);
> > +}
> > +
> > +noinline int
> > +uprobe_session_consumer_test(struct uprobe_multi_session_consumers *skel,
> > +                            unsigned long before, unsigned long after)
> > +{
> > +       int bit;
> > +
> > +       /* detach uprobe for each unset bit in 'before' state ... */
> > +       for (bit = 0; bit < 6; bit++) {
> 
> Does "bit" correspond to the uprobe_X program? Maybe call it an uprobe
> index or something, if that's the case? bits are just representations,
> but semantically meaningful is identifier of an uprobe program, right?

right.. so it corresponds to program 'uprobe_<bit>' so maybe 'idx' is better

thanks,
jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Wed, Jul 3, 2024 at 1:10 AM Peter Zijlstra <peterz@infradead.org> wrote:
>
> On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:
> > > +static size_t ri_size(int sessions_cnt)
> > > +{
> > > +       struct return_instance *ri __maybe_unused;
> > > +
> > > +       return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
> >
> > just use struct_size()?
>
> Yeah, lets not. This is readable, struct_size() is not.

This hack with __maybe_unused is more readable than the standard
struct_size() helper that was added specifically for cases like this,
really?

I wonder if Kees agrees and whether there are any downsides to using
struct_size()

struct_size(struct return_instance, sessions, sessions_cnt) seems
readable enough to me, in any case.

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Kees Cook]

On Wed, Jul 03, 2024 at 11:31:11AM -0700, Andrii Nakryiko wrote:
> On Wed, Jul 3, 2024 at 1:10 AM Peter Zijlstra <peterz@infradead.org> wrote:
> >
> > On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:
> > > > +static size_t ri_size(int sessions_cnt)
> > > > +{
> > > > +       struct return_instance *ri __maybe_unused;
> > > > +
> > > > +       return sizeof(*ri) + sessions_cnt * sizeof(ri->sessions[0]);
> > >
> > > just use struct_size()?
> >
> > Yeah, lets not. This is readable, struct_size() is not.
> 
> This hack with __maybe_unused is more readable than the standard
> struct_size() helper that was added specifically for cases like this,
> really?
> 
> I wonder if Kees agrees and whether there are any downsides to using
> struct_size()
> 
> struct_size(struct return_instance, sessions, sessions_cnt) seems
> readable enough to me, in any case.

Yes, please use struct_size_t(). This is exactly what it was designed for.

Though with only 2 instances of ri_size(), maybe just use struct_size()
directly?

Also, please annotate struct return_instance with __counted_by:

+	int			sessions_cnt;
+	struct session_consumer	sessions[] __counted_by(sessions_cnt);


-- 
Kees Cook

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Wed, Jul 3, 2024 at 8:31 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Tue, Jul 02, 2024 at 01:52:38PM -0700, Andrii Nakryiko wrote:
> > On Tue, Jul 2, 2024 at 9:11 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> > >
> > > On Tue, Jul 02, 2024 at 03:04:08PM +0200, Peter Zijlstra wrote:
> > > > On Mon, Jul 01, 2024 at 06:41:07PM +0200, Jiri Olsa wrote:
> > > >
> > > > > +static void
> > > > > +uprobe_consumer_account(struct uprobe *uprobe, struct uprobe_consumer *uc)
> > > > > +{
> > > > > +   static unsigned int session_id;
> > > > > +
> > > > > +   if (uc->session) {
> > > > > +           uprobe->sessions_cnt++;
> > > > > +           uc->session_id = ++session_id ?: ++session_id;
> > > > > +   }
> > > > > +}
> > > >
> > > > The way I understand this code, you create a consumer every time you do
> > > > uprobe_register() and unregister makes it go away.
> > > >
> > > > Now, register one, then 4g-1 times register+unregister, then register
> > > > again.
> > > >
> > > > The above seems to then result in two consumers with the same
> > > > session_id, which leads to trouble.
> > > >
> > > > Hmm?
> > >
> > > ugh true.. will make it u64 :)
> > >
> > > I think we could store uprobe_consumer pointer+ref in session_consumer,
> > > and that would make the unregister path more interesting.. will check
> >
> > More interesting how? It's actually a great idea, uprobe_consumer
>
> nah, got confused ;-)
>
> > pointer itself is a unique ID and 64-bit. We can still use lowest bit
> > for RC (see my other reply).
>
> I used pointers in the previous version, but then I thought what if the
> consumer gets free-ed and new one created (with same address.. maybe not
> likely but possible, right?) before the return probe is hit

I think no matter what we do, uprobe_unregister() API has to guarantee
that when it returns consumer won't be hit (i.e., we removed consumer
from uprobe->consumers list, waited for RCU grace period(s), etc). So
I don't think this should be a problem. And that's one of the reasons
for the need for batched unregister, because we'll have to do sync_rcu
call there for this.

>
> jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Wed, Jul 3, 2024 at 9:09 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Tue, Jul 02, 2024 at 05:13:38PM -0700, Andrii Nakryiko wrote:
> > On Tue, Jul 2, 2024 at 4:55 PM Masami Hiramatsu <mhiramat@kernel.org> wrote:
> > >
> > > Hi Jiri,
> > >
> > > On Mon,  1 Jul 2024 18:41:07 +0200
> > > Jiri Olsa <jolsa@kernel.org> wrote:
> > >
> > > > Adding support for uprobe consumer to be defined as session and have
> > > > new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
> > > >
> > > > The session means that 'handler' and 'ret_handler' callbacks are
> > > > connected in a way that allows to:
> > > >
> > > >   - control execution of 'ret_handler' from 'handler' callback
> > > >   - share data between 'handler' and 'ret_handler' callbacks
> > > >
> > > > The session is enabled by setting new 'session' bool field to true
> > > > in uprobe_consumer object.
> > > >
> > > > We keep count of session consumers for uprobe and allocate session_consumer
> > > > object for each in return_instance object. This allows us to store
> > > > return values of 'handler' callbacks and data pointers of shared
> > > > data between both handlers.
> > > >
> > > > The session concept fits to our common use case where we do filtering
> > > > on entry uprobe and based on the result we decide to run the return
> > > > uprobe (or not).
> > > >
> > > > It's also convenient to share the data between session callbacks.
> > > >
> > > > The control of 'ret_handler' callback execution is done via return
> > > > value of the 'handler' callback. If it's 0 we install and execute
> > > > return uprobe, if it's 1 we do not.
> > > >
> > > > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > > > ---
> > > >  include/linux/uprobes.h     |  16 ++++-
> > > >  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
> > > >  kernel/trace/bpf_trace.c    |   6 +-
> > > >  kernel/trace/trace_uprobe.c |  12 ++--
> > > >  4 files changed, 144 insertions(+), 19 deletions(-)
> > > >
> > > > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > > > index f46e0ca0169c..903a860a8d01 100644
> > > > --- a/include/linux/uprobes.h
> > > > +++ b/include/linux/uprobes.h
> > > > @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
> > > >  };
> > > >
> > > >  struct uprobe_consumer {
> > > > -     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> > > > +     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
> > > >       int (*ret_handler)(struct uprobe_consumer *self,
> > > >                               unsigned long func,
> > > > -                             struct pt_regs *regs);
> > > > +                             struct pt_regs *regs, __u64 *data);
> > > >       bool (*filter)(struct uprobe_consumer *self,
> > > >                               enum uprobe_filter_ctx ctx,
> > > >                               struct mm_struct *mm);
> > > >
> > > >       struct uprobe_consumer *next;
> > > > +
> > > > +     bool                    session;        /* marks uprobe session consumer */
> > > > +     unsigned int            session_id;     /* set when uprobe_consumer is registered */
> > >
> > > Hmm, why this has both session and session_id?
> >
> > session is caller's request to establish session semantics. Jiri, I
>
> and session_id is set when uprobe is registered and used when
> return uprobe is executed to find matching uprobe_consumer,
> plz check handle_uretprobe_chain/session_consumer_find
>
> > think it's better to move it higher next to
> > handler/ret_handler/filter, that's the part of uprobe_consumer struct
> > which has read-only caller-provided data (I'm adding offset and
> > ref_ctr_offset there as well).
>
> ok, makes sense
>
> >
> > > I also think we can use the address of uprobe_consumer itself as a unique id.
> >
> > +1
> >
> > >
> > > Also, if we can set session enabled by default, and skip ret_handler by handler's
> > > return value, it is more simpler. (If handler returns a specific value, skip ret_handler)
> >
> > you mean derive if it's a session or not by both handler and
> > ret_handler being set? I guess this works fine for BPF side, because
> > there we never had them both set. If this doesn't regress others, I
> > think it's OK. We just need to make sure we don't unnecessarily
> > allocate session state for consumers that don't set both handler and
> > ret_handler. That would be a waste.
>
> hum.. so the current code installs return uprobe if there's ret_handler
> defined in consumer and also the entry 'handler' needs to return 0
>
> if entry 'handler' returns 1 the uprobe is unregistered
>
> we could define new return value from 'handler' to 'not execute the
> 'ret_handler' and have 'handler' return values:
>
>   0 - execute 'ret_handler' if defined
>   1 - remove the uprobe
>   2 - do NOT execute 'ret_handler'  // this current triggers WARN
>
> we could delay the allocation of 'return_instance' until the first
> consumer returns 0, so there's no perf regression
>
> that way we could treat all consumers the same and we wouldn't need
> the session flag..
>
> ok looks like good idea ;-) will try that

Just please double check that we don't pass through 1 or 2 as a return
result for BPF uprobes/multi-uprobes, so that we don't have any
accidental changes of behavior.

>
> thanks,
> jirka
>
> >
> > >
> > > >  };
> > > >
> > > >  #ifdef CONFIG_UPROBES
> > > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > > >       unsigned int                    depth;
> > > >  };
> > > >
> > > > +struct session_consumer {
> > > > +     __u64           cookie;
> > >
> > > And this cookie looks not scalable. If we can pass a data to handler, I would like to
> > > reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.
> > >
> > >         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);
> > >
> > > uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
> > > at handler_chain().
> >
> > The goal here is to keep this simple and fast. I'd prefer to keep it
> > small and fixed size, if possible. I'm thinking about caching and
> > reusing return_instance as one of the future optimizations, so if we
> > can keep this more or less fixed (assuming there is typically not more
> > than 1 or 2 consumers per uprobe, which seems realistic), this will
> > provide a way to avoid excessive memory allocations.
> >
> > >
> > > > +     unsigned int    id;
> > > > +     int             rc;
> > > > +};
> > > > +
> > > >  struct return_instance {
> > > >       struct uprobe           *uprobe;
> > > >       unsigned long           func;
> > > > @@ -88,6 +97,9 @@ struct return_instance {
> > > >       bool                    chained;        /* true, if instance is nested */
> > > >
> > > >       struct return_instance  *next;          /* keep as stack */
> > > > +
> > > > +     int                     sessions_cnt;
> > > > +     struct session_consumer sessions[];
> > >
> > > In that case, we don't have this array, but
> > >
> > >         char data[];
> > >
> > > And decode data array, which is a slice of variable length structure;
> > >
> > > struct session_consumer {
> > >         struct uprobe_consumer *uc;
> > >         char data[];
> > > };
> > >
> > > The size of session_consumer is uc->session_data_size + sizeof(uc).
> > >
> > > What would you think?
> > >
> > > Thank you,
> > >
> > > >  };
> > > >
> > > >  enum rp_check {
> > > > diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> > > > index 2c83ba776fc7..4da410460f2a 100644
> > > > --- a/kernel/events/uprobes.c
> > > > +++ b/kernel/events/uprobes.c
> > > > @@ -63,6 +63,8 @@ struct uprobe {
> > > >       loff_t                  ref_ctr_offset;
> > > >       unsigned long           flags;
> > > >
> > > > +     unsigned int            sessions_cnt;
> >
> > [...]

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Andrii Nakryiko]

On Wed, Jul 3, 2024 at 10:13 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:
>
> SNIP
>
> > >  #ifdef CONFIG_UPROBES
> > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > >         unsigned int                    depth;
> > >  };
> > >
> > > +struct session_consumer {
> > > +       __u64           cookie;
> > > +       unsigned int    id;
> > > +       int             rc;
> >
> > you'll be using u64 for ID, right? so this struct will be 24 bytes.
>
> yes
>
> > Maybe we can just use topmost bit of ID to store whether uretprobe
> > should run or not? It's trivial to mask out during ID comparisons
>
> actually.. I think we could store just consumers that need to be
> executed in return probe so there will be no need for 'rc' value

ah, nice idea. NULL would mean we have session uprobe, but for this
particular run we "disabled" uretprobe part of it. Great. And for
non-session uprobes we just won't have session_consumer at all, right?

[...]

> > > +static struct session_consumer *
> > > +session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
> > > +                     int session_id)
> > > +{
> > > +       struct session_consumer *next;
> > > +
> > > +       next = sc ? sc + 1 : &ri->sessions[0];
> > > +       next->id = session_id;
> >
> > it's kind of unexpected that "session_consumer_next" would actually
> > set an ID... Maybe drop int session_id as input argument and fill it
> > out outside of this function, this function being just a simple
> > iterator?
>
> yea, I was going back and forth on what to have in that function
> or not, to keep the change minimal, but makes sense, will move
>

great, thanks

> >
> > > +       return next;
> > > +}
> > > +

[...]

> >
> > > +               } else if (uc->ret_handler) {
> > >                         need_prep = true;
> > > +               }
> > >
> > >                 remove &= rc;
> > >         }
> > >
> > > +       /* no removal if there's at least one session consumer */
> > > +       remove &= !uprobe->sessions_cnt;
> >
> > this is counter (not error, not pointer), let's stick to ` == 0`, please
> >
> > is this
> >
> > if (uprobe->sessions_cnt != 0)
> >    remove = 0;
>
> yes ;-) will change
>

Thanks, I feel bad for being the only one to call this out, but I find
all these '!<some_integer_variable>` constructs extremely unintuitive
and hard to reason about quickly. It's only pointers and error cases
that are more or less intuitive. Everything else, including
!strcmp(...) is just mind bending and exhausting... Perhaps I'm just
not a kernel engineer enough :)

> jirka
>
> >
> > ? I can't tell (honestly), without spending ridiculous amounts of
> > mental resources (for the underlying simplicity of the condition).
>
> SNIP

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Peter Zijlstra]

On Wed, Jul 03, 2024 at 01:36:19PM -0700, Kees Cook wrote:

> Yes, please use struct_size_t(). This is exactly what it was designed for.

Kees, please, just let up, not going to happen. I'm getting really fed
up having to endlessly repeat what a piece of shite struct_size() is.

Put your time and effort into doing a proper language extension so we
can go and delete all that __builtin_*_overflow() based garbage.

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Masami Hiramatsu]

On Wed, 3 Jul 2024 14:43:27 -0700
Andrii Nakryiko <andrii.nakryiko@gmail.com> wrote:

> On Wed, Jul 3, 2024 at 9:09 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> >
> > On Tue, Jul 02, 2024 at 05:13:38PM -0700, Andrii Nakryiko wrote:
> > > On Tue, Jul 2, 2024 at 4:55 PM Masami Hiramatsu <mhiramat@kernel.org> wrote:
> > > >
> > > > Hi Jiri,
> > > >
> > > > On Mon,  1 Jul 2024 18:41:07 +0200
> > > > Jiri Olsa <jolsa@kernel.org> wrote:
> > > >
> > > > > Adding support for uprobe consumer to be defined as session and have
> > > > > new behaviour for consumer's 'handler' and 'ret_handler' callbacks.
> > > > >
> > > > > The session means that 'handler' and 'ret_handler' callbacks are
> > > > > connected in a way that allows to:
> > > > >
> > > > >   - control execution of 'ret_handler' from 'handler' callback
> > > > >   - share data between 'handler' and 'ret_handler' callbacks
> > > > >
> > > > > The session is enabled by setting new 'session' bool field to true
> > > > > in uprobe_consumer object.
> > > > >
> > > > > We keep count of session consumers for uprobe and allocate session_consumer
> > > > > object for each in return_instance object. This allows us to store
> > > > > return values of 'handler' callbacks and data pointers of shared
> > > > > data between both handlers.
> > > > >
> > > > > The session concept fits to our common use case where we do filtering
> > > > > on entry uprobe and based on the result we decide to run the return
> > > > > uprobe (or not).
> > > > >
> > > > > It's also convenient to share the data between session callbacks.
> > > > >
> > > > > The control of 'ret_handler' callback execution is done via return
> > > > > value of the 'handler' callback. If it's 0 we install and execute
> > > > > return uprobe, if it's 1 we do not.
> > > > >
> > > > > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > > > > ---
> > > > >  include/linux/uprobes.h     |  16 ++++-
> > > > >  kernel/events/uprobes.c     | 129 +++++++++++++++++++++++++++++++++---
> > > > >  kernel/trace/bpf_trace.c    |   6 +-
> > > > >  kernel/trace/trace_uprobe.c |  12 ++--
> > > > >  4 files changed, 144 insertions(+), 19 deletions(-)
> > > > >
> > > > > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > > > > index f46e0ca0169c..903a860a8d01 100644
> > > > > --- a/include/linux/uprobes.h
> > > > > +++ b/include/linux/uprobes.h
> > > > > @@ -34,15 +34,18 @@ enum uprobe_filter_ctx {
> > > > >  };
> > > > >
> > > > >  struct uprobe_consumer {
> > > > > -     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> > > > > +     int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, __u64 *data);
> > > > >       int (*ret_handler)(struct uprobe_consumer *self,
> > > > >                               unsigned long func,
> > > > > -                             struct pt_regs *regs);
> > > > > +                             struct pt_regs *regs, __u64 *data);
> > > > >       bool (*filter)(struct uprobe_consumer *self,
> > > > >                               enum uprobe_filter_ctx ctx,
> > > > >                               struct mm_struct *mm);
> > > > >
> > > > >       struct uprobe_consumer *next;
> > > > > +
> > > > > +     bool                    session;        /* marks uprobe session consumer */
> > > > > +     unsigned int            session_id;     /* set when uprobe_consumer is registered */
> > > >
> > > > Hmm, why this has both session and session_id?
> > >
> > > session is caller's request to establish session semantics. Jiri, I
> >
> > and session_id is set when uprobe is registered and used when
> > return uprobe is executed to find matching uprobe_consumer,
> > plz check handle_uretprobe_chain/session_consumer_find
> >
> > > think it's better to move it higher next to
> > > handler/ret_handler/filter, that's the part of uprobe_consumer struct
> > > which has read-only caller-provided data (I'm adding offset and
> > > ref_ctr_offset there as well).
> >
> > ok, makes sense
> >
> > >
> > > > I also think we can use the address of uprobe_consumer itself as a unique id.
> > >
> > > +1
> > >
> > > >
> > > > Also, if we can set session enabled by default, and skip ret_handler by handler's
> > > > return value, it is more simpler. (If handler returns a specific value, skip ret_handler)
> > >
> > > you mean derive if it's a session or not by both handler and
> > > ret_handler being set? I guess this works fine for BPF side, because
> > > there we never had them both set. If this doesn't regress others, I
> > > think it's OK. We just need to make sure we don't unnecessarily
> > > allocate session state for consumers that don't set both handler and
> > > ret_handler. That would be a waste.
> >
> > hum.. so the current code installs return uprobe if there's ret_handler
> > defined in consumer and also the entry 'handler' needs to return 0
> >
> > if entry 'handler' returns 1 the uprobe is unregistered
> >
> > we could define new return value from 'handler' to 'not execute the
> > 'ret_handler' and have 'handler' return values:
> >
> >   0 - execute 'ret_handler' if defined
> >   1 - remove the uprobe
> >   2 - do NOT execute 'ret_handler'  // this current triggers WARN
> >
> > we could delay the allocation of 'return_instance' until the first
> > consumer returns 0, so there's no perf regression
> >
> > that way we could treat all consumers the same and we wouldn't need
> > the session flag..
> >
> > ok looks like good idea ;-) will try that
> 
> Just please double check that we don't pass through 1 or 2 as a return
> result for BPF uprobes/multi-uprobes, so that we don't have any
> accidental changes of behavior.

Agreed. BTW, even if the uprobe is removed, the ret_handler should be called?
I think both 1 and 2 case, we should skip ret_handler.

> > > > >
> > > > >  #ifdef CONFIG_UPROBES
> > > > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > > > >       unsigned int                    depth;
> > > > >  };
> > > > >
> > > > > +struct session_consumer {
> > > > > +     __u64           cookie;
> > > >
> > > > And this cookie looks not scalable. If we can pass a data to handler, I would like to
> > > > reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.
> > > >
> > > >         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);
> > > >
> > > > uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
> > > > at handler_chain().
> > >
> > > The goal here is to keep this simple and fast. I'd prefer to keep it
> > > small and fixed size, if possible. I'm thinking about caching and
> > > reusing return_instance as one of the future optimizations, so if we
> > > can keep this more or less fixed (assuming there is typically not more
> > > than 1 or 2 consumers per uprobe, which seems realistic), this will
> > > provide a way to avoid excessive memory allocations.

Hmm, so you mean user will allocate another "data map" and use cookie as
a key to access the data? That is possible but sounds a bit redundant.
If such "data map" allocation is also provided, it is more useful.


Thank you,




-- 
Masami Hiramatsu (Google) <mhiramat@kernel.org>

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Wed, Jul 03, 2024 at 02:48:28PM -0700, Andrii Nakryiko wrote:
> On Wed, Jul 3, 2024 at 10:13 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> >
> > On Tue, Jul 02, 2024 at 01:51:28PM -0700, Andrii Nakryiko wrote:
> >
> > SNIP
> >
> > > >  #ifdef CONFIG_UPROBES
> > > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > > >         unsigned int                    depth;
> > > >  };
> > > >
> > > > +struct session_consumer {
> > > > +       __u64           cookie;
> > > > +       unsigned int    id;
> > > > +       int             rc;
> > >
> > > you'll be using u64 for ID, right? so this struct will be 24 bytes.
> >
> > yes
> >
> > > Maybe we can just use topmost bit of ID to store whether uretprobe
> > > should run or not? It's trivial to mask out during ID comparisons
> >
> > actually.. I think we could store just consumers that need to be
> > executed in return probe so there will be no need for 'rc' value
> 
> ah, nice idea. NULL would mean we have session uprobe, but for this
> particular run we "disabled" uretprobe part of it. Great. And for
> non-session uprobes we just won't have session_consumer at all, right?

hm, I think we don't need to add both session or non-session consumer
if it's not supposed to run.. let's see

> 
> [...]
> 
> > > > +static struct session_consumer *
> > > > +session_consumer_next(struct return_instance *ri, struct session_consumer *sc,
> > > > +                     int session_id)
> > > > +{
> > > > +       struct session_consumer *next;
> > > > +
> > > > +       next = sc ? sc + 1 : &ri->sessions[0];
> > > > +       next->id = session_id;
> > >
> > > it's kind of unexpected that "session_consumer_next" would actually
> > > set an ID... Maybe drop int session_id as input argument and fill it
> > > out outside of this function, this function being just a simple
> > > iterator?
> >
> > yea, I was going back and forth on what to have in that function
> > or not, to keep the change minimal, but makes sense, will move
> >
> 
> great, thanks
> 
> > >
> > > > +       return next;
> > > > +}
> > > > +
> 
> [...]
> 
> > >
> > > > +               } else if (uc->ret_handler) {
> > > >                         need_prep = true;
> > > > +               }
> > > >
> > > >                 remove &= rc;
> > > >         }
> > > >
> > > > +       /* no removal if there's at least one session consumer */
> > > > +       remove &= !uprobe->sessions_cnt;
> > >
> > > this is counter (not error, not pointer), let's stick to ` == 0`, please
> > >
> > > is this
> > >
> > > if (uprobe->sessions_cnt != 0)
> > >    remove = 0;
> >
> > yes ;-) will change
> >
> 
> Thanks, I feel bad for being the only one to call this out, but I find
> all these '!<some_integer_variable>` constructs extremely unintuitive
> and hard to reason about quickly. It's only pointers and error cases
> that are more or less intuitive. Everything else, including
> !strcmp(...) is just mind bending and exhausting... Perhaps I'm just
> not a kernel engineer enough :)

heh I was going for minimal change.. but it's intrusive enough already,
so let's keep it at least readable

jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Jiri Olsa]

On Fri, Jul 05, 2024 at 05:35:44PM +0900, Masami Hiramatsu wrote:

SNIP

> > > > > Also, if we can set session enabled by default, and skip ret_handler by handler's
> > > > > return value, it is more simpler. (If handler returns a specific value, skip ret_handler)
> > > >
> > > > you mean derive if it's a session or not by both handler and
> > > > ret_handler being set? I guess this works fine for BPF side, because
> > > > there we never had them both set. If this doesn't regress others, I
> > > > think it's OK. We just need to make sure we don't unnecessarily
> > > > allocate session state for consumers that don't set both handler and
> > > > ret_handler. That would be a waste.
> > >
> > > hum.. so the current code installs return uprobe if there's ret_handler
> > > defined in consumer and also the entry 'handler' needs to return 0
> > >
> > > if entry 'handler' returns 1 the uprobe is unregistered
> > >
> > > we could define new return value from 'handler' to 'not execute the
> > > 'ret_handler' and have 'handler' return values:
> > >
> > >   0 - execute 'ret_handler' if defined
> > >   1 - remove the uprobe
> > >   2 - do NOT execute 'ret_handler'  // this current triggers WARN
> > >
> > > we could delay the allocation of 'return_instance' until the first
> > > consumer returns 0, so there's no perf regression
> > >
> > > that way we could treat all consumers the same and we wouldn't need
> > > the session flag..
> > >
> > > ok looks like good idea ;-) will try that
> > 
> > Just please double check that we don't pass through 1 or 2 as a return
> > result for BPF uprobes/multi-uprobes, so that we don't have any
> > accidental changes of behavior.
> 
> Agreed. BTW, even if the uprobe is removed, the ret_handler should be called?
> I think both 1 and 2 case, we should skip ret_handler.

do you mean what happens when the uretprobe is installed and its consumer
is unregistered before it's triggered?

I think it won't get executed, because the consumer is removed right away,
even if the uprobe object stays because the return_instance holds ref to it

> 
> > > > > >
> > > > > >  #ifdef CONFIG_UPROBES
> > > > > > @@ -80,6 +83,12 @@ struct uprobe_task {
> > > > > >       unsigned int                    depth;
> > > > > >  };
> > > > > >
> > > > > > +struct session_consumer {
> > > > > > +     __u64           cookie;
> > > > >
> > > > > And this cookie looks not scalable. If we can pass a data to handler, I would like to
> > > > > reuse it to pass the target function parameters to ret_handler as kretprobe/fprobe does.
> > > > >
> > > > >         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs, void *data);
> > > > >
> > > > > uprobes can collect its uc's required sizes and allocate the memory (shadow stack frame)
> > > > > at handler_chain().
> > > >
> > > > The goal here is to keep this simple and fast. I'd prefer to keep it
> > > > small and fixed size, if possible. I'm thinking about caching and
> > > > reusing return_instance as one of the future optimizations, so if we
> > > > can keep this more or less fixed (assuming there is typically not more
> > > > than 1 or 2 consumers per uprobe, which seems realistic), this will
> > > > provide a way to avoid excessive memory allocations.
> 
> Hmm, so you mean user will allocate another "data map" and use cookie as
> a key to access the data? That is possible but sounds a bit redundant.
> If such "data map" allocation is also provided, it is more useful.

is the argument only about the size of the shared data?

we can make it configurable.. for bpf user we will probably stick to
8 bytes to match the kprobe session and to be compatible with existing
helpers accessing the cookie

jirka

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Kees Cook]

On Fri, Jul 05, 2024 at 09:10:36AM +0200, Peter Zijlstra wrote:
> On Wed, Jul 03, 2024 at 01:36:19PM -0700, Kees Cook wrote:
> 
> > Yes, please use struct_size_t(). This is exactly what it was designed for.
> 
> Kees, please, just let up, not going to happen. I'm getting really fed
> up having to endlessly repeat what a piece of shite struct_size() is.

I mean, okay, but the wrapper in the patch is basically the same thing.
*shrug*

> Put your time and effort into doing a proper language extension so we
> can go and delete all that __builtin_*_overflow() based garbage.

We are! That's in the future. Today, we have a saturating wrapper that
provides type checking for the calculation's operands, and is in common
use through-out the kernel. These are all things that the open-coded
does not provide, so I continue to see it as an improvement over what
else is available right now.

I got asked for my opinion about whether to use struct_size() or not. In
my opinion, this is a good place for it. I know you don't agree with me,
but that wasn't the question. :)

-Kees

-- 
Kees Cook

Re: [PATCHv2 bpf-next 1/9] uprobe: Add support for session consumer

[Peter Zijlstra]

On Fri, Jul 05, 2024 at 03:38:12PM +0200, Jiri Olsa wrote:

> > Agreed. BTW, even if the uprobe is removed, the ret_handler should be called?
> > I think both 1 and 2 case, we should skip ret_handler.
> 
> do you mean what happens when the uretprobe is installed and its consumer
> is unregistered before it's triggered?
> 
> I think it won't get executed, because the consumer is removed right away,
> even if the uprobe object stays because the return_instance holds ref to it

Yep, that is my understanding too. RI keeps the uprobe object around,
but the consumers can go at any time.