[RFC bpf-next 00/10] uprobe, bpf: Add session support

[RFC bpf-next 00/10] 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.

thanks,
jirka


---
Jiri Olsa (10):
      uprobe: Add session callbacks to uprobe_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: Move ARRAY_SIZE to bpf_misc.h
      selftests/bpf: Add uprobe session test
      selftests/bpf: Add uprobe session errors test
      selftests/bpf: Add uprobe session cookie test
      selftests/bpf: Add uprobe session recursive test

 include/linux/uprobes.h                                            |  18 ++++++++++
 include/uapi/linux/bpf.h                                           |   1 +
 kernel/bpf/syscall.c                                               |   9 +++--
 kernel/events/uprobes.c                                            |  69 +++++++++++++++++++++++++++++++++-----
 kernel/trace/bpf_trace.c                                           |  72 +++++++++++++++++++++++++++++++---------
 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         | 153 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 tools/testing/selftests/bpf/progs/bpf_misc.h                       |   2 ++
 tools/testing/selftests/bpf/progs/iters.c                          |   2 --
 tools/testing/selftests/bpf/progs/kprobe_multi_session.c           |   3 +-
 tools/testing/selftests/bpf/progs/linked_list.c                    |   5 +--
 tools/testing/selftests/bpf/progs/netif_receive_skb.c              |   5 +--
 tools/testing/selftests/bpf/progs/profiler.inc.h                   |   5 +--
 tools/testing/selftests/bpf/progs/setget_sockopt.c                 |   5 +--
 tools/testing/selftests/bpf/progs/test_bpf_ma.c                    |   4 ---
 tools/testing/selftests/bpf/progs/test_sysctl_loop1.c              |   5 +--
 tools/testing/selftests/bpf/progs/test_sysctl_loop2.c              |   5 +--
 tools/testing/selftests/bpf/progs/test_sysctl_prog.c               |   5 +--
 tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.c      |   1 +
 tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.h      |   2 --
 tools/testing/selftests/bpf/progs/uprobe_multi_session.c           |  52 +++++++++++++++++++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c    |  50 ++++++++++++++++++++++++++++
 tools/testing/selftests/bpf/progs/uprobe_multi_session_recursive.c |  44 ++++++++++++++++++++++++
 tools/testing/selftests/bpf/progs/verifier_subprog_precision.c     |   2 --
 27 files changed, 507 insertions(+), 69 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session.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

[RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

Adding new set of callbacks that are triggered on entry and return
uprobe execution for the attached function.

The session means that those callbacks are 'connected' in a way
that allows to:
  - control execution of return callback from entry callback
  - share data between entry and return callbacks

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 return uprobe execution is done via return value of the
entry session callback, where 0 means to install and execute return
uprobe, 1 means to not install.

Current implementation has a restriction that allows to register only
single consumer with session callbacks for a uprobe and also restricting
standard callbacks consumers.

Which means that there can be only single user of a uprobe (inode +
offset) when session consumer is registered to it.

This is because all registered consumers are executed when uprobe or
return uprobe is hit and wihout additional layer (like fgraph's shadow
stack) that would keep the state of the return callback, we have no
way to find out which consumer should be executed.

I'm not sure how big limitation this is for people, our current use
case seems to be ok with that. Fixing this would be more complex/bigger
change to uprobes, thoughts?

Hence sending this as RFC to gather more opinions and feedback.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/uprobes.h | 18 +++++++++++
 kernel/events/uprobes.c | 69 +++++++++++++++++++++++++++++++++++------
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
index f46e0ca0169c..a2f2d5ac3cee 100644
--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h
@@ -34,6 +34,12 @@ enum uprobe_filter_ctx {
 };
 
 struct uprobe_consumer {
+	/*
+	 * The handler callback return value controls removal of the uprobe.
+	 *  0 on success, uprobe stays
+	 *  1 on failure, remove the uprobe
+	 *    console warning for anything else
+	 */
 	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
 	int (*ret_handler)(struct uprobe_consumer *self,
 				unsigned long func,
@@ -42,6 +48,17 @@ struct uprobe_consumer {
 				enum uprobe_filter_ctx ctx,
 				struct mm_struct *mm);
 
+	/* The handler_session callback return value controls execution of
+	 * the return uprobe and ret_handler_session callback.
+	 *  0 on success
+	 *  1 on failure, DO NOT install/execute the return uprobe
+	 *    console warning for anything else
+	 */
+	int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
+			       unsigned long *data);
+	int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
+				   struct pt_regs *regs, unsigned long *data);
+
 	struct uprobe_consumer *next;
 };
 
@@ -85,6 +102,7 @@ struct return_instance {
 	unsigned long		func;
 	unsigned long		stack;		/* stack pointer */
 	unsigned long		orig_ret_vaddr; /* original return address */
+	unsigned long		data;
 	bool			chained;	/* true, if instance is nested */
 
 	struct return_instance	*next;		/* keep as stack */
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index 2c83ba776fc7..17b0771272a6 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -750,12 +750,32 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
 	return uprobe;
 }
 
-static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
+/*
+ * Make sure all the uprobe consumers have only one type of entry
+ * callback registered (either handler or handler_session) due to
+ * different return value actions.
+ */
+static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
+{
+	if (!curr)
+		return 0;
+	if (curr->handler_session || uc->handler_session)
+		return -EBUSY;
+	return 0;
+}
+
+static int consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
 {
+	int err;
+
 	down_write(&uprobe->consumer_rwsem);
-	uc->next = uprobe->consumers;
-	uprobe->consumers = uc;
+	err = consumer_check(uprobe->consumers, uc);
+	if (!err) {
+		uc->next = uprobe->consumers;
+		uprobe->consumers = uc;
+	}
 	up_write(&uprobe->consumer_rwsem);
+	return err;
 }
 
 /*
@@ -1114,6 +1134,21 @@ void uprobe_unregister(struct inode *inode, loff_t offset, struct uprobe_consume
 }
 EXPORT_SYMBOL_GPL(uprobe_unregister);
 
+static int check_handler(struct uprobe_consumer *uc)
+{
+	/* Uprobe must have at least one set consumer. */
+	if (!uc->handler && !uc->ret_handler &&
+	    !uc->handler_session && !uc->ret_handler_session)
+		return -1;
+	/* Session consumer is exclusive. */
+	if (uc->handler && uc->handler_session)
+		return -1;
+	/* Session consumer must have both entry and return handler. */
+	if (!!uc->handler_session != !!uc->ret_handler_session)
+		return -1;
+	return 0;
+}
+
 /*
  * __uprobe_register - register a probe
  * @inode: the file in which the probe has to be placed.
@@ -1138,8 +1173,7 @@ static int __uprobe_register(struct inode *inode, loff_t offset,
 	struct uprobe *uprobe;
 	int ret;
 
-	/* Uprobe must have at least one set consumer */
-	if (!uc->handler && !uc->ret_handler)
+	if (check_handler(uc))
 		return -EINVAL;
 
 	/* copy_insn() uses read_mapping_page() or shmem_read_mapping_page() */
@@ -1173,11 +1207,14 @@ static int __uprobe_register(struct inode *inode, loff_t offset,
 	down_write(&uprobe->register_rwsem);
 	ret = -EAGAIN;
 	if (likely(uprobe_is_active(uprobe))) {
-		consumer_add(uprobe, uc);
+		ret = consumer_add(uprobe, uc);
+		if (ret)
+			goto fail;
 		ret = register_for_each_vma(uprobe, uc);
 		if (ret)
 			__uprobe_unregister(uprobe, uc);
 	}
+ fail:
 	up_write(&uprobe->register_rwsem);
 	put_uprobe(uprobe);
 
@@ -1853,7 +1890,7 @@ 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, unsigned long data)
 {
 	struct return_instance *ri;
 	struct uprobe_task *utask;
@@ -1909,6 +1946,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 	ri->stack = user_stack_pointer(regs);
 	ri->orig_ret_vaddr = orig_ret_vaddr;
 	ri->chained = chained;
+	ri->data = data;
 
 	utask->depth++;
 	ri->next = utask->return_instances;
@@ -2070,6 +2108,7 @@ static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
 	struct uprobe_consumer *uc;
 	int remove = UPROBE_HANDLER_REMOVE;
 	bool need_prep = false; /* prepare return uprobe, when needed */
+	unsigned long data = 0;
 
 	down_read(&uprobe->register_rwsem);
 	for (uc = uprobe->consumers; uc; uc = uc->next) {
@@ -2081,14 +2120,24 @@ static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
 				"bad rc=0x%x from %ps()\n", rc, uc->handler);
 		}
 
-		if (uc->ret_handler)
+		if (uc->handler_session) {
+			rc = uc->handler_session(uc, regs, &data);
+			WARN(rc & ~UPROBE_HANDLER_MASK,
+				"bad rc=0x%x from %ps()\n", rc, uc->handler_session);
+		}
+
+		if (uc->ret_handler || uc->ret_handler_session)
 			need_prep = true;
 
 		remove &= rc;
 	}
 
 	if (need_prep && !remove)
-		prepare_uretprobe(uprobe, regs); /* put bp at return */
+		prepare_uretprobe(uprobe, regs, data); /* put bp at return */
+
+	/* remove uprobe only for non-session consumers */
+	if (uprobe->consumers && remove)
+		remove &= !!uprobe->consumers->handler;
 
 	if (remove && uprobe->consumers) {
 		WARN_ON(!uprobe_is_active(uprobe));
@@ -2107,6 +2156,8 @@ handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
 	for (uc = uprobe->consumers; uc; uc = uc->next) {
 		if (uc->ret_handler)
 			uc->ret_handler(uc, ri->func, regs);
+		if (uc->ret_handler_session)
+			uc->ret_handler_session(uc, ri->func, regs, &ri->data);
 	}
 	up_read(&uprobe->register_rwsem);
 }
-- 
2.45.1

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

I'll try to read this code tomorrow, right now I don't really understand
what does it do and why.

However,

On 06/04, Jiri Olsa wrote:
>
>  struct uprobe_consumer {
> +	/*
> +	 * The handler callback return value controls removal of the uprobe.
> +	 *  0 on success, uprobe stays
> +	 *  1 on failure, remove the uprobe
> +	 *    console warning for anything else
> +	 */
>  	int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);

This is misleading. It is not about success/failure, it is about filtering.

consumer->handler() returns UPROBE_HANDLER_REMOVE if this consumer is not
interested in this task, so this uprobe can be removed (unless another
consumer returns 0).

> +/*
> + * Make sure all the uprobe consumers have only one type of entry
> + * callback registered (either handler or handler_session) due to
> + * different return value actions.
> + */
> +static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
> +{
> +	if (!curr)
> +		return 0;
> +	if (curr->handler_session || uc->handler_session)
> +		return -EBUSY;
> +	return 0;
> +}

Hmm, I don't understand this code, it doesn't match the comment...

The comment says "all the uprobe consumers have only one type" but
consumer_check() will always fail if the the 1st or 2nd consumer has
->handler_session != NULL ?

Perhaps you meant

	if (!!curr->handler != !!uc->handler)
		return -EBUSY;

?

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

On 06/05, Oleg Nesterov wrote:
>
> > +/*
> > + * Make sure all the uprobe consumers have only one type of entry
> > + * callback registered (either handler or handler_session) due to
> > + * different return value actions.
> > + */
> > +static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
> > +{
> > +	if (!curr)
> > +		return 0;
> > +	if (curr->handler_session || uc->handler_session)
> > +		return -EBUSY;
> > +	return 0;
> > +}
>
> Hmm, I don't understand this code, it doesn't match the comment...
>
> The comment says "all the uprobe consumers have only one type" but
> consumer_check() will always fail if the the 1st or 2nd consumer has
> ->handler_session != NULL ?
>
> Perhaps you meant
>
> 	if (!!curr->handler != !!uc->handler)
> 		return -EBUSY;
>
> ?

OK, the changelog says

	Which means that there can be only single user of a uprobe (inode +
	offset) when session consumer is registered to it.

so the code is correct. But I still think the comment is misleading.

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

On 06/05, Oleg Nesterov wrote:
>
> On 06/05, Oleg Nesterov wrote:
> >
> > > +/*
> > > + * Make sure all the uprobe consumers have only one type of entry
> > > + * callback registered (either handler or handler_session) due to
> > > + * different return value actions.
> > > + */
> > > +static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
> > > +{
> > > +	if (!curr)
> > > +		return 0;
> > > +	if (curr->handler_session || uc->handler_session)
> > > +		return -EBUSY;
> > > +	return 0;
> > > +}
> >
> > Hmm, I don't understand this code, it doesn't match the comment...
> >
> > The comment says "all the uprobe consumers have only one type" but
> > consumer_check() will always fail if the the 1st or 2nd consumer has
> > ->handler_session != NULL ?
> >
> > Perhaps you meant
> >
> > 	if (!!curr->handler != !!uc->handler)
> > 		return -EBUSY;
> >
> > ?
>
> OK, the changelog says
>
> 	Which means that there can be only single user of a uprobe (inode +
> 	offset) when session consumer is registered to it.
>
> so the code is correct. But I still think the comment is misleading.

Cough... perhaps it is correct but I am still confused even we forget about
the comment ;)

OK, uprobe can have a single consumer with ->handler_session != NULL. I guess
this is because return_instance->data is "global".

So uprobe can have multiple handler_session == NULL consumers before
handler_session != NULL, but not after ?

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Wed, Jun 05, 2024 at 06:36:25PM +0200, Oleg Nesterov wrote:
> On 06/05, Oleg Nesterov wrote:
> >
> > On 06/05, Oleg Nesterov wrote:
> > >
> > > > +/*
> > > > + * Make sure all the uprobe consumers have only one type of entry
> > > > + * callback registered (either handler or handler_session) due to
> > > > + * different return value actions.
> > > > + */
> > > > +static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
> > > > +{
> > > > +	if (!curr)
> > > > +		return 0;
> > > > +	if (curr->handler_session || uc->handler_session)
> > > > +		return -EBUSY;
> > > > +	return 0;
> > > > +}
> > >
> > > Hmm, I don't understand this code, it doesn't match the comment...
> > >
> > > The comment says "all the uprobe consumers have only one type" but
> > > consumer_check() will always fail if the the 1st or 2nd consumer has
> > > ->handler_session != NULL ?
> > >
> > > Perhaps you meant
> > >
> > > 	if (!!curr->handler != !!uc->handler)
> > > 		return -EBUSY;
> > >
> > > ?
> >
> > OK, the changelog says
> >
> > 	Which means that there can be only single user of a uprobe (inode +
> > 	offset) when session consumer is registered to it.
> >
> > so the code is correct. But I still think the comment is misleading.
> 
> Cough... perhaps it is correct but I am still confused even we forget about
> the comment ;)
> 
> OK, uprobe can have a single consumer with ->handler_session != NULL. I guess
> this is because return_instance->data is "global".
> 
> So uprobe can have multiple handler_session == NULL consumers before
> handler_session != NULL, but not after ?

ah yea it should have done what's in the comment, so it's missing
the check for handler.. session handlers are meant to be exclusive

thanks,
jirka

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Andrii Nakryiko]

On Tue, Jun 4, 2024 at 1:02 PM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding new set of callbacks that are triggered on entry and return
> uprobe execution for the attached function.
>
> The session means that those callbacks are 'connected' in a way
> that allows to:
>   - control execution of return callback from entry callback
>   - share data between entry and return callbacks
>
> 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 return uprobe execution is done via return value of the
> entry session callback, where 0 means to install and execute return
> uprobe, 1 means to not install.
>
> Current implementation has a restriction that allows to register only
> single consumer with session callbacks for a uprobe and also restricting
> standard callbacks consumers.
>
> Which means that there can be only single user of a uprobe (inode +
> offset) when session consumer is registered to it.
>
> This is because all registered consumers are executed when uprobe or
> return uprobe is hit and wihout additional layer (like fgraph's shadow
> stack) that would keep the state of the return callback, we have no
> way to find out which consumer should be executed.
>
> I'm not sure how big limitation this is for people, our current use
> case seems to be ok with that. Fixing this would be more complex/bigger
> change to uprobes, thoughts?

I think it's a pretty big limitation, because in production you don't
always know ahead of time all possible users of uprobe, so any such
limitations will cause problems, issue reports, investigation, etc.

As one possible solution, what if we do

struct return_instance {
    ...
    u64 session_cookies[];
};

and allocate sizeof(struct return_instance) + 8 *
<num-of-session-consumers> and then at runtime pass
&session_cookies[i] as data pointer to session-aware callbacks?

>
> Hence sending this as RFC to gather more opinions and feedback.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/linux/uprobes.h | 18 +++++++++++
>  kernel/events/uprobes.c | 69 +++++++++++++++++++++++++++++++++++------
>  2 files changed, 78 insertions(+), 9 deletions(-)
>
> diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> index f46e0ca0169c..a2f2d5ac3cee 100644
> --- a/include/linux/uprobes.h
> +++ b/include/linux/uprobes.h
> @@ -34,6 +34,12 @@ enum uprobe_filter_ctx {
>  };
>
>  struct uprobe_consumer {
> +       /*
> +        * The handler callback return value controls removal of the uprobe.
> +        *  0 on success, uprobe stays
> +        *  1 on failure, remove the uprobe
> +        *    console warning for anything else
> +        */
>         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
>         int (*ret_handler)(struct uprobe_consumer *self,
>                                 unsigned long func,
> @@ -42,6 +48,17 @@ struct uprobe_consumer {
>                                 enum uprobe_filter_ctx ctx,
>                                 struct mm_struct *mm);
>
> +       /* The handler_session callback return value controls execution of
> +        * the return uprobe and ret_handler_session callback.
> +        *  0 on success
> +        *  1 on failure, DO NOT install/execute the return uprobe
> +        *    console warning for anything else
> +        */
> +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> +                              unsigned long *data);
> +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> +                                  struct pt_regs *regs, unsigned long *data);
> +

We should try to avoid an alternative set of callbacks, IMO. Let's
extend existing ones with `unsigned long *data`, but specify that
unless consumer sets some flag on registration that it needs a session
cookie, we'll pass NULL here? Or just allocate cookie data for each
registered consumer for simplicity, don't know; given we don't expect
many consumers on exactly the same uprobe, it might be ok to keep it
simple.


>         struct uprobe_consumer *next;
>  };
>
> @@ -85,6 +102,7 @@ struct return_instance {
>         unsigned long           func;
>         unsigned long           stack;          /* stack pointer */
>         unsigned long           orig_ret_vaddr; /* original return address */
> +       unsigned long           data;
>         bool                    chained;        /* true, if instance is nested */
>
>         struct return_instance  *next;          /* keep as stack */
> diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
> index 2c83ba776fc7..17b0771272a6 100644
> --- a/kernel/events/uprobes.c
> +++ b/kernel/events/uprobes.c
> @@ -750,12 +750,32 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
>         return uprobe;
>  }
>
> -static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
> +/*
> + * Make sure all the uprobe consumers have only one type of entry
> + * callback registered (either handler or handler_session) due to
> + * different return value actions.
> + */
> +static int consumer_check(struct uprobe_consumer *curr, struct uprobe_consumer *uc)
> +{
> +       if (!curr)
> +               return 0;
> +       if (curr->handler_session || uc->handler_session)
> +               return -EBUSY;
> +       return 0;
> +}
> +
> +static int consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
>  {
> +       int err;
> +
>         down_write(&uprobe->consumer_rwsem);
> -       uc->next = uprobe->consumers;
> -       uprobe->consumers = uc;
> +       err = consumer_check(uprobe->consumers, uc);
> +       if (!err) {
> +               uc->next = uprobe->consumers;
> +               uprobe->consumers = uc;
> +       }
>         up_write(&uprobe->consumer_rwsem);
> +       return err;
>  }
>
>  /*
> @@ -1114,6 +1134,21 @@ void uprobe_unregister(struct inode *inode, loff_t offset, struct uprobe_consume
>  }
>  EXPORT_SYMBOL_GPL(uprobe_unregister);
>
> +static int check_handler(struct uprobe_consumer *uc)
> +{
> +       /* Uprobe must have at least one set consumer. */
> +       if (!uc->handler && !uc->ret_handler &&
> +           !uc->handler_session && !uc->ret_handler_session)
> +               return -1;
> +       /* Session consumer is exclusive. */
> +       if (uc->handler && uc->handler_session)
> +               return -1;
> +       /* Session consumer must have both entry and return handler. */
> +       if (!!uc->handler_session != !!uc->ret_handler_session)
> +               return -1;
> +       return 0;
> +}
> +
>  /*
>   * __uprobe_register - register a probe
>   * @inode: the file in which the probe has to be placed.
> @@ -1138,8 +1173,7 @@ static int __uprobe_register(struct inode *inode, loff_t offset,
>         struct uprobe *uprobe;
>         int ret;
>
> -       /* Uprobe must have at least one set consumer */
> -       if (!uc->handler && !uc->ret_handler)
> +       if (check_handler(uc))
>                 return -EINVAL;
>
>         /* copy_insn() uses read_mapping_page() or shmem_read_mapping_page() */
> @@ -1173,11 +1207,14 @@ static int __uprobe_register(struct inode *inode, loff_t offset,
>         down_write(&uprobe->register_rwsem);
>         ret = -EAGAIN;
>         if (likely(uprobe_is_active(uprobe))) {
> -               consumer_add(uprobe, uc);
> +               ret = consumer_add(uprobe, uc);
> +               if (ret)
> +                       goto fail;
>                 ret = register_for_each_vma(uprobe, uc);
>                 if (ret)
>                         __uprobe_unregister(uprobe, uc);
>         }
> + fail:
>         up_write(&uprobe->register_rwsem);
>         put_uprobe(uprobe);
>
> @@ -1853,7 +1890,7 @@ 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, unsigned long data)
>  {
>         struct return_instance *ri;
>         struct uprobe_task *utask;
> @@ -1909,6 +1946,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
>         ri->stack = user_stack_pointer(regs);
>         ri->orig_ret_vaddr = orig_ret_vaddr;
>         ri->chained = chained;
> +       ri->data = data;
>
>         utask->depth++;
>         ri->next = utask->return_instances;
> @@ -2070,6 +2108,7 @@ static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
>         struct uprobe_consumer *uc;
>         int remove = UPROBE_HANDLER_REMOVE;
>         bool need_prep = false; /* prepare return uprobe, when needed */
> +       unsigned long data = 0;
>
>         down_read(&uprobe->register_rwsem);
>         for (uc = uprobe->consumers; uc; uc = uc->next) {
> @@ -2081,14 +2120,24 @@ static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
>                                 "bad rc=0x%x from %ps()\n", rc, uc->handler);
>                 }
>
> -               if (uc->ret_handler)
> +               if (uc->handler_session) {
> +                       rc = uc->handler_session(uc, regs, &data);
> +                       WARN(rc & ~UPROBE_HANDLER_MASK,
> +                               "bad rc=0x%x from %ps()\n", rc, uc->handler_session);
> +               }
> +
> +               if (uc->ret_handler || uc->ret_handler_session)
>                         need_prep = true;
>
>                 remove &= rc;
>         }
>
>         if (need_prep && !remove)
> -               prepare_uretprobe(uprobe, regs); /* put bp at return */
> +               prepare_uretprobe(uprobe, regs, data); /* put bp at return */
> +
> +       /* remove uprobe only for non-session consumers */
> +       if (uprobe->consumers && remove)
> +               remove &= !!uprobe->consumers->handler;
>
>         if (remove && uprobe->consumers) {
>                 WARN_ON(!uprobe_is_active(uprobe));
> @@ -2107,6 +2156,8 @@ handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
>         for (uc = uprobe->consumers; uc; uc = uc->next) {
>                 if (uc->ret_handler)
>                         uc->ret_handler(uc, ri->func, regs);
> +               if (uc->ret_handler_session)
> +                       uc->ret_handler_session(uc, ri->func, regs, &ri->data);
>         }
>         up_read(&uprobe->register_rwsem);
>  }
> --
> 2.45.1
>

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

On 06/05, Andrii Nakryiko wrote:
>
> so any such
> limitations will cause problems, issue reports, investigation, etc.

Agreed...

> As one possible solution, what if we do
>
> struct return_instance {
>     ...
>     u64 session_cookies[];
> };
>
> and allocate sizeof(struct return_instance) + 8 *
> <num-of-session-consumers> and then at runtime pass
> &session_cookies[i] as data pointer to session-aware callbacks?

I too thought about this, but I guess it is not that simple.

Just for example. Suppose we have 2 session-consumers C1 and C2.
What if uprobe_unregister(C1) comes before the probed function
returns?

We need something like map_cookie_to_consumer().

> > +       /* The handler_session callback return value controls execution of
> > +        * the return uprobe and ret_handler_session callback.
> > +        *  0 on success
> > +        *  1 on failure, DO NOT install/execute the return uprobe
> > +        *    console warning for anything else
> > +        */
> > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > +                              unsigned long *data);
> > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > +                                  struct pt_regs *regs, unsigned long *data);
> > +
>
> We should try to avoid an alternative set of callbacks, IMO. Let's
> extend existing ones with `unsigned long *data`,

Oh yes, agreed.

And the comment about the return value looks confusing too. I mean, the
logic doesn't differ from the ret-code from ->handler().

"DO NOT install/execute the return uprobe" is not true if another
non-session-consumer returns 0.

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Andrii Nakryiko]

On Wed, Jun 5, 2024 at 10:57 AM Oleg Nesterov <oleg@redhat.com> wrote:
>
> On 06/05, Andrii Nakryiko wrote:
> >
> > so any such
> > limitations will cause problems, issue reports, investigation, etc.
>
> Agreed...
>
> > As one possible solution, what if we do
> >
> > struct return_instance {
> >     ...
> >     u64 session_cookies[];
> > };
> >
> > and allocate sizeof(struct return_instance) + 8 *
> > <num-of-session-consumers> and then at runtime pass
> > &session_cookies[i] as data pointer to session-aware callbacks?
>
> I too thought about this, but I guess it is not that simple.
>
> Just for example. Suppose we have 2 session-consumers C1 and C2.
> What if uprobe_unregister(C1) comes before the probed function
> returns?
>
> We need something like map_cookie_to_consumer().

Fair enough. The easy way to solve this is to have


struct uprobe_session_cookie {
    int consumer_id;
    u64 cookie;
};

And add id to each new consumer when it is added to struct uprobe.
Unfortunately, it's impossible to tell when a new consumer was added
to the list (as a front item, but maybe we just change it to be
appended instead of prepending) vs when the old consumer was removed,
so in some cases we'd need to do a linear search.

But the good news is that in the common case we wouldn't need to
search and the next item in session_cookies[] array would be the one
we need.

WDYT? It's still fast, and it's simpler than the shadow stack idea, IMO.

P.S. Regardless, maybe we should change the order in which we insert
consumers to uprobe? Right now uprobe consumer added later will be
executed first, which, while not wrong, is counter-intuitive. And also
it breaks a nice natural order when we need to match it up with stuff
like session_cookies[] as described above.

>
> > > +       /* The handler_session callback return value controls execution of
> > > +        * the return uprobe and ret_handler_session callback.
> > > +        *  0 on success
> > > +        *  1 on failure, DO NOT install/execute the return uprobe
> > > +        *    console warning for anything else
> > > +        */
> > > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > > +                              unsigned long *data);
> > > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > > +                                  struct pt_regs *regs, unsigned long *data);
> > > +
> >
> > We should try to avoid an alternative set of callbacks, IMO. Let's
> > extend existing ones with `unsigned long *data`,
>
> Oh yes, agreed.
>
> And the comment about the return value looks confusing too. I mean, the
> logic doesn't differ from the ret-code from ->handler().
>
> "DO NOT install/execute the return uprobe" is not true if another
> non-session-consumer returns 0.
>
> Oleg.
>

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Wed, Jun 05, 2024 at 01:47:00PM -0700, Andrii Nakryiko wrote:
> On Wed, Jun 5, 2024 at 10:57 AM Oleg Nesterov <oleg@redhat.com> wrote:
> >
> > On 06/05, Andrii Nakryiko wrote:
> > >
> > > so any such
> > > limitations will cause problems, issue reports, investigation, etc.
> >
> > Agreed...
> >
> > > As one possible solution, what if we do
> > >
> > > struct return_instance {
> > >     ...
> > >     u64 session_cookies[];
> > > };
> > >
> > > and allocate sizeof(struct return_instance) + 8 *
> > > <num-of-session-consumers> and then at runtime pass
> > > &session_cookies[i] as data pointer to session-aware callbacks?
> >
> > I too thought about this, but I guess it is not that simple.
> >
> > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > What if uprobe_unregister(C1) comes before the probed function
> > returns?
> >
> > We need something like map_cookie_to_consumer().
> 
> Fair enough. The easy way to solve this is to have
> 
> 
> struct uprobe_session_cookie {
>     int consumer_id;
>     u64 cookie;
> };
> 
> And add id to each new consumer when it is added to struct uprobe.
> Unfortunately, it's impossible to tell when a new consumer was added
> to the list (as a front item, but maybe we just change it to be
> appended instead of prepending) vs when the old consumer was removed,
> so in some cases we'd need to do a linear search.

also we probably need to add the flag if we want to execute the return
handler..  we can have multiple session handlers and if just one of them
returns 0 we need to install the return probe

and then when return probe hits, we need to execute only that consumer's
return handler

jirka

> 
> But the good news is that in the common case we wouldn't need to
> search and the next item in session_cookies[] array would be the one
> we need.
> 
> WDYT? It's still fast, and it's simpler than the shadow stack idea, IMO.
> 
> P.S. Regardless, maybe we should change the order in which we insert
> consumers to uprobe? Right now uprobe consumer added later will be
> executed first, which, while not wrong, is counter-intuitive. And also
> it breaks a nice natural order when we need to match it up with stuff
> like session_cookies[] as described above.
> 
> >
> > > > +       /* The handler_session callback return value controls execution of
> > > > +        * the return uprobe and ret_handler_session callback.
> > > > +        *  0 on success
> > > > +        *  1 on failure, DO NOT install/execute the return uprobe
> > > > +        *    console warning for anything else
> > > > +        */
> > > > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > > > +                              unsigned long *data);
> > > > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > > > +                                  struct pt_regs *regs, unsigned long *data);
> > > > +
> > >
> > > We should try to avoid an alternative set of callbacks, IMO. Let's
> > > extend existing ones with `unsigned long *data`,
> >
> > Oh yes, agreed.
> >
> > And the comment about the return value looks confusing too. I mean, the
> > logic doesn't differ from the ret-code from ->handler().
> >
> > "DO NOT install/execute the return uprobe" is not true if another
> > non-session-consumer returns 0.
> >
> > Oleg.
> >

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

On 06/05, Andrii Nakryiko wrote:
>
> WDYT? It's still fast, and it's simpler than the shadow stack idea, IMO.

Andrii. I am alredy sleeping, I'll try to read your email tomorrow.
Right now I can only say that everything is simpler than the shadow stack ;)

> P.S. Regardless, maybe we should change the order in which we insert
> consumers to uprobe? Right now uprobe consumer added later will be
> executed first, which, while not wrong, is counter-intuitive.

Agreed...

Even if currently this doesn't really matter, I guess it is supposed
that uc->handler() is "non-intrusive".

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> On 06/05, Andrii Nakryiko wrote:
> >
> > so any such
> > limitations will cause problems, issue reports, investigation, etc.
> 
> Agreed...
> 
> > As one possible solution, what if we do
> >
> > struct return_instance {
> >     ...
> >     u64 session_cookies[];
> > };
> >
> > and allocate sizeof(struct return_instance) + 8 *
> > <num-of-session-consumers> and then at runtime pass
> > &session_cookies[i] as data pointer to session-aware callbacks?
> 
> I too thought about this, but I guess it is not that simple.
> 
> Just for example. Suppose we have 2 session-consumers C1 and C2.
> What if uprobe_unregister(C1) comes before the probed function
> returns?
> 
> We need something like map_cookie_to_consumer().

I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?

return instance is freed after the consumers' return handlers are executed,
so there's no leak if some consumer gets unregistered before that

> 
> > > +       /* The handler_session callback return value controls execution of
> > > +        * the return uprobe and ret_handler_session callback.
> > > +        *  0 on success
> > > +        *  1 on failure, DO NOT install/execute the return uprobe
> > > +        *    console warning for anything else
> > > +        */
> > > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > > +                              unsigned long *data);
> > > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > > +                                  struct pt_regs *regs, unsigned long *data);
> > > +
> >
> > We should try to avoid an alternative set of callbacks, IMO. Let's
> > extend existing ones with `unsigned long *data`,
> 
> Oh yes, agreed.
> 
> And the comment about the return value looks confusing too. I mean, the
> logic doesn't differ from the ret-code from ->handler().
> 
> "DO NOT install/execute the return uprobe" is not true if another
> non-session-consumer returns 0.

well they are meant to be exclusive, so there'd be no other non-session-consumer

jirka

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Oleg Nesterov]

On 06/05, Jiri Olsa wrote:
>
> > And the comment about the return value looks confusing too. I mean, the
> > logic doesn't differ from the ret-code from ->handler().
> >
> > "DO NOT install/execute the return uprobe" is not true if another
> > non-session-consumer returns 0.
>
> well they are meant to be exclusive, so there'd be no other non-session-consumer

OK. (but may be the changelog can explain more clearly why they can't
co-exist with the non-session-consumers).

But again, this doesn't differ from the the ret-code from the
non-session-consumer->handler().

If it returns 1 == UPROBE_HANDLER_REMOVE, then without other consumers
prepare_uretprobe() won't be called.

Oleg.

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Wed, Jun 05, 2024 at 10:50:11PM +0200, Jiri Olsa wrote:
> On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> > On 06/05, Andrii Nakryiko wrote:
> > >
> > > so any such
> > > limitations will cause problems, issue reports, investigation, etc.
> > 
> > Agreed...
> > 
> > > As one possible solution, what if we do
> > >
> > > struct return_instance {
> > >     ...
> > >     u64 session_cookies[];
> > > };
> > >
> > > and allocate sizeof(struct return_instance) + 8 *
> > > <num-of-session-consumers> and then at runtime pass
> > > &session_cookies[i] as data pointer to session-aware callbacks?
> > 
> > I too thought about this, but I guess it is not that simple.
> > 
> > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > What if uprobe_unregister(C1) comes before the probed function
> > returns?
> > 
> > We need something like map_cookie_to_consumer().
> 
> I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?

ok, hash table is probably too big for this.. I guess some solution that
would iterate consumers and cookies made sure it matches would be fine

jirka

> 
> return instance is freed after the consumers' return handlers are executed,
> so there's no leak if some consumer gets unregistered before that
> 
> > 
> > > > +       /* The handler_session callback return value controls execution of
> > > > +        * the return uprobe and ret_handler_session callback.
> > > > +        *  0 on success
> > > > +        *  1 on failure, DO NOT install/execute the return uprobe
> > > > +        *    console warning for anything else
> > > > +        */
> > > > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > > > +                              unsigned long *data);
> > > > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > > > +                                  struct pt_regs *regs, unsigned long *data);
> > > > +
> > >
> > > We should try to avoid an alternative set of callbacks, IMO. Let's
> > > extend existing ones with `unsigned long *data`,
> > 
> > Oh yes, agreed.
> > 
> > And the comment about the return value looks confusing too. I mean, the
> > logic doesn't differ from the ret-code from ->handler().
> > 
> > "DO NOT install/execute the return uprobe" is not true if another
> > non-session-consumer returns 0.
> 
> well they are meant to be exclusive, so there'd be no other non-session-consumer
> 
> jirka

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Andrii Nakryiko]

On Thu, Jun 6, 2024 at 9:46 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Wed, Jun 05, 2024 at 10:50:11PM +0200, Jiri Olsa wrote:
> > On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> > > On 06/05, Andrii Nakryiko wrote:
> > > >
> > > > so any such
> > > > limitations will cause problems, issue reports, investigation, etc.
> > >
> > > Agreed...
> > >
> > > > As one possible solution, what if we do
> > > >
> > > > struct return_instance {
> > > >     ...
> > > >     u64 session_cookies[];
> > > > };
> > > >
> > > > and allocate sizeof(struct return_instance) + 8 *
> > > > <num-of-session-consumers> and then at runtime pass
> > > > &session_cookies[i] as data pointer to session-aware callbacks?
> > >
> > > I too thought about this, but I guess it is not that simple.
> > >
> > > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > > What if uprobe_unregister(C1) comes before the probed function
> > > returns?
> > >
> > > We need something like map_cookie_to_consumer().
> >
> > I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?
>
> ok, hash table is probably too big for this.. I guess some solution that
> would iterate consumers and cookies made sure it matches would be fine
>

Yes, I was hoping to avoid hash tables for this, and in the common
case have no added overhead.

> jirka
>
> >
> > return instance is freed after the consumers' return handlers are executed,
> > so there's no leak if some consumer gets unregistered before that
> >
> > >
> > > > > +       /* The handler_session callback return value controls execution of
> > > > > +        * the return uprobe and ret_handler_session callback.
> > > > > +        *  0 on success
> > > > > +        *  1 on failure, DO NOT install/execute the return uprobe
> > > > > +        *    console warning for anything else
> > > > > +        */
> > > > > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > > > > +                              unsigned long *data);
> > > > > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > > > > +                                  struct pt_regs *regs, unsigned long *data);
> > > > > +
> > > >
> > > > We should try to avoid an alternative set of callbacks, IMO. Let's
> > > > extend existing ones with `unsigned long *data`,
> > >
> > > Oh yes, agreed.
> > >
> > > And the comment about the return value looks confusing too. I mean, the
> > > logic doesn't differ from the ret-code from ->handler().
> > >
> > > "DO NOT install/execute the return uprobe" is not true if another
> > > non-session-consumer returns 0.
> >
> > well they are meant to be exclusive, so there'd be no other non-session-consumer
> >
> > jirka

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Thu, Jun 06, 2024 at 09:52:39AM -0700, Andrii Nakryiko wrote:
> On Thu, Jun 6, 2024 at 9:46 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> >
> > On Wed, Jun 05, 2024 at 10:50:11PM +0200, Jiri Olsa wrote:
> > > On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> > > > On 06/05, Andrii Nakryiko wrote:
> > > > >
> > > > > so any such
> > > > > limitations will cause problems, issue reports, investigation, etc.
> > > >
> > > > Agreed...
> > > >
> > > > > As one possible solution, what if we do
> > > > >
> > > > > struct return_instance {
> > > > >     ...
> > > > >     u64 session_cookies[];
> > > > > };
> > > > >
> > > > > and allocate sizeof(struct return_instance) + 8 *
> > > > > <num-of-session-consumers> and then at runtime pass
> > > > > &session_cookies[i] as data pointer to session-aware callbacks?
> > > >
> > > > I too thought about this, but I guess it is not that simple.
> > > >
> > > > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > > > What if uprobe_unregister(C1) comes before the probed function
> > > > returns?
> > > >
> > > > We need something like map_cookie_to_consumer().
> > >
> > > I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?
> >
> > ok, hash table is probably too big for this.. I guess some solution that
> > would iterate consumers and cookies made sure it matches would be fine
> >
> 
> Yes, I was hoping to avoid hash tables for this, and in the common
> case have no added overhead.

hi,
here's first stab on that.. the change below:
  - extends current handlers with extra argument rather than adding new
    set of handlers
  - store session consumers objects within return_instance object and
  - iterate these objects ^^^ in handle_uretprobe_chain

I guess it could be still polished, but I wonder if this could
be the right direction to do this.. thoughts? ;-)

thanks,
jirka


---
diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
index f46e0ca0169c..4e40e8352eac 100644
--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h
@@ -34,15 +34,19 @@ 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,
+			unsigned long *data);
 	int (*ret_handler)(struct uprobe_consumer *self,
 				unsigned long func,
-				struct pt_regs *regs);
+				struct pt_regs *regs,
+				unsigned long *data);
 	bool (*filter)(struct uprobe_consumer *self,
 				enum uprobe_filter_ctx ctx,
 				struct mm_struct *mm);
 
 	struct uprobe_consumer *next;
+	bool is_session;
+	unsigned int id;
 };
 
 #ifdef CONFIG_UPROBES
@@ -80,6 +84,12 @@ struct uprobe_task {
 	unsigned int			depth;
 };
 
+struct session_consumer {
+	long cookie;
+	unsigned int id;
+	int rc;
+};
+
 struct return_instance {
 	struct uprobe		*uprobe;
 	unsigned long		func;
@@ -88,6 +98,8 @@ struct return_instance {
 	bool			chained;	/* true, if instance is nested */
 
 	struct return_instance	*next;		/* keep as stack */
+	int			session_cnt;
+	struct session_consumer	sc[1];		/* 1 for zero item marking the end */
 };
 
 enum rp_check {
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index 2c83ba776fc7..cbd71dc06ef0 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		session_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->is_session) {
+		uprobe->session_cnt++;
+		uc->id = ++session_id ?: ++session_id;
+	}
+}
+
+static void
+uprobe_consumer_unaccount(struct uprobe *uprobe, struct uprobe_consumer *uc)
+{
+	if (uc->is_session)
+		uprobe->session_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 session_cnt)
+{
+	struct return_instance *ri __maybe_unused;
+
+	return sizeof(*ri) + session_cnt * sizeof(ri->sc[0]);
+}
+
+static struct return_instance *alloc_return_instance(int session_cnt)
+{
+	struct return_instance *ri;
+
+	ri = kzalloc(ri_size(session_cnt), GFP_KERNEL);
+	if (ri)
+		ri->session_cnt = session_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->session_cnt);
 		if (!n)
 			return -ENOMEM;
 
-		*n = *o;
+		memcpy(n, o, ri_size(o->session_cnt));
 		get_uprobe(n->uprobe);
 		n->next = NULL;
 
@@ -1853,35 +1892,38 @@ 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 struct return_instance *
+prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
+		  struct return_instance *ri, int session_cnt)
 {
-	struct return_instance *ri;
 	struct uprobe_task *utask;
 	unsigned long orig_ret_vaddr, trampoline_vaddr;
 	bool chained;
 
 	if (!get_xol_area())
-		return;
+		return ri;
 
 	utask = get_utask();
 	if (!utask)
-		return;
+		return ri;
 
 	if (utask->depth >= MAX_URETPROBE_DEPTH) {
 		printk_ratelimited(KERN_INFO "uprobe: omit uretprobe due to"
 				" nestedness limit pid/tgid=%d/%d\n",
 				current->pid, current->tgid);
-		return;
+		return ri;
 	}
 
-	ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
-	if (!ri)
-		return;
+	if (!ri) {
+		ri = alloc_return_instance(session_cnt);
+		if (!ri)
+			return NULL;
+	}
 
 	trampoline_vaddr = get_trampoline_vaddr();
 	orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
 	if (orig_ret_vaddr == -1)
-		goto fail;
+		return ri;
 
 	/* drop the entries invalidated by longjmp() */
 	chained = (orig_ret_vaddr == trampoline_vaddr);
@@ -1899,7 +1941,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 			 * attack from user-space.
 			 */
 			uprobe_warn(current, "handle tail call");
-			goto fail;
+			return ri;
 		}
 		orig_ret_vaddr = utask->return_instances->orig_ret_vaddr;
 	}
@@ -1914,9 +1956,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 	ri->next = utask->return_instances;
 	utask->return_instances = ri;
 
-	return;
- fail:
-	kfree(ri);
+	return NULL;
 }
 
 /* Prepare to single-step probed instruction out of line. */
@@ -2069,44 +2109,90 @@ 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);
-	for (uc = uprobe->consumers; uc; uc = uc->next) {
+	if (uprobe->session_cnt) {
+		ri = alloc_return_instance(uprobe->session_cnt);
+		if (!ri)
+			goto out;
+	}
+	for (uc = uprobe->consumers, sc = &ri->sc[0]; uc; uc = uc->next) {
 		int rc = 0;
 
 		if (uc->handler) {
-			rc = uc->handler(uc, regs);
+			rc = uc->handler(uc, regs, uc->is_session ? &sc->cookie : NULL);
 			WARN(rc & ~UPROBE_HANDLER_MASK,
 				"bad rc=0x%x from %ps()\n", rc, uc->handler);
 		}
 
-		if (uc->ret_handler)
+		if (uc->is_session) {
+			need_prep |= !rc;
+			remove = 0;
+			sc->id = uc->id;
+			sc->rc = rc;
+			sc++;
+		} else if (uc->ret_handler) {
 			need_prep = true;
+		}
 
 		remove &= rc;
 	}
 
 	if (need_prep && !remove)
-		prepare_uretprobe(uprobe, regs); /* put bp at return */
+		ri = prepare_uretprobe(uprobe, regs, ri, uprobe->session_cnt); /* put bp at return */
+	kfree(ri);
 
 	if (remove && uprobe->consumers) {
 		WARN_ON(!uprobe_is_active(uprobe));
 		unapply_uprobe(uprobe, current->mm);
 	}
+ out:
 	up_read(&uprobe->register_rwsem);
 }
 
+static struct session_consumer *
+consumer_find(struct session_consumer *sc, struct uprobe_consumer *uc)
+{
+	for (; sc && sc->id; sc++) {
+		if (sc->id == uc->id)
+			return sc;
+	}
+	return NULL;
+}
+
 static void
 handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
 {
 	struct uprobe *uprobe = ri->uprobe;
+	struct session_consumer *sc, *tmp;
 	struct uprobe_consumer *uc;
 
 	down_read(&uprobe->register_rwsem);
-	for (uc = uprobe->consumers; uc; uc = uc->next) {
-		if (uc->ret_handler)
-			uc->ret_handler(uc, ri->func, regs);
+	for (uc = uprobe->consumers, sc = &ri->sc[0]; uc; uc = uc->next) {
+		long *cookie = NULL;
+		int rc = 0;
+
+		if (uc->is_session) {
+			/*
+			 * session_consumers are in order with uprobe_consumers,
+			 * we just need to reflect that any uprobe_consumer could
+			 * be removed or added
+			 */
+			tmp = consumer_find(sc, uc);
+			if (tmp) {
+				rc = tmp->rc;
+				cookie = &tmp->cookie;
+				sc = tmp + 1;
+			} else {
+				rc = 1;
+			}
+		}
+
+		if (!rc && uc->ret_handler)
+			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 f5154c051d2c..ae7c35379e4a 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3329,7 +3329,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,
+			  unsigned long *data)
 {
 	struct bpf_uprobe *uprobe;
 
@@ -3338,7 +3339,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,
+			      unsigned long *data)
 {
 	struct bpf_uprobe *uprobe;
 
diff --git a/kernel/trace/trace_uprobe.c b/kernel/trace/trace_uprobe.c
index 8541fa1494ae..f7b17f08344c 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,
+			     unsigned long *data);
 static int uretprobe_dispatcher(struct uprobe_consumer *con,
-				unsigned long func, struct pt_regs *regs);
+				unsigned long func, struct pt_regs *regs,
+				unsigned long *data);
 
 #ifdef CONFIG_STACK_GROWSUP
 static unsigned long adjust_stack_addr(unsigned long addr, unsigned int n)
@@ -1500,7 +1502,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,
+			     unsigned long *data)
 {
 	struct trace_uprobe *tu;
 	struct uprobe_dispatch_data udd;
@@ -1530,7 +1533,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,
+				unsigned long *data)
 {
 	struct trace_uprobe *tu;
 	struct uprobe_dispatch_data udd;

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Andrii Nakryiko]

On Mon, Jun 10, 2024 at 4:06 AM Jiri Olsa <olsajiri@gmail.com> wrote:
>
> On Thu, Jun 06, 2024 at 09:52:39AM -0700, Andrii Nakryiko wrote:
> > On Thu, Jun 6, 2024 at 9:46 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> > >
> > > On Wed, Jun 05, 2024 at 10:50:11PM +0200, Jiri Olsa wrote:
> > > > On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> > > > > On 06/05, Andrii Nakryiko wrote:
> > > > > >
> > > > > > so any such
> > > > > > limitations will cause problems, issue reports, investigation, etc.
> > > > >
> > > > > Agreed...
> > > > >
> > > > > > As one possible solution, what if we do
> > > > > >
> > > > > > struct return_instance {
> > > > > >     ...
> > > > > >     u64 session_cookies[];
> > > > > > };
> > > > > >
> > > > > > and allocate sizeof(struct return_instance) + 8 *
> > > > > > <num-of-session-consumers> and then at runtime pass
> > > > > > &session_cookies[i] as data pointer to session-aware callbacks?
> > > > >
> > > > > I too thought about this, but I guess it is not that simple.
> > > > >
> > > > > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > > > > What if uprobe_unregister(C1) comes before the probed function
> > > > > returns?
> > > > >
> > > > > We need something like map_cookie_to_consumer().
> > > >
> > > > I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?
> > >
> > > ok, hash table is probably too big for this.. I guess some solution that
> > > would iterate consumers and cookies made sure it matches would be fine
> > >
> >
> > Yes, I was hoping to avoid hash tables for this, and in the common
> > case have no added overhead.
>
> hi,
> here's first stab on that.. the change below:
>   - extends current handlers with extra argument rather than adding new
>     set of handlers
>   - store session consumers objects within return_instance object and
>   - iterate these objects ^^^ in handle_uretprobe_chain
>
> I guess it could be still polished, but I wonder if this could
> be the right direction to do this.. thoughts? ;-)

Yeah, I think this is the right direction. It's a bit sad that this
makes getting rid of rw_sem on hot path even harder, but that's a
separate problem.

>
> thanks,
> jirka
>
>
> ---
> diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> index f46e0ca0169c..4e40e8352eac 100644
> --- a/include/linux/uprobes.h
> +++ b/include/linux/uprobes.h
> @@ -34,15 +34,19 @@ 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,
> +                       unsigned long *data);

can we use __u64 here? This long vs __u64 might cause problems for BPF
when the host is 32-bit architecture (BPF is always 64-bit).

>         int (*ret_handler)(struct uprobe_consumer *self,
>                                 unsigned long func,
> -                               struct pt_regs *regs);
> +                               struct pt_regs *regs,
> +                               unsigned long *data);
>         bool (*filter)(struct uprobe_consumer *self,
>                                 enum uprobe_filter_ctx ctx,
>                                 struct mm_struct *mm);
>

[...]

>  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->session_cnt);
>                 if (!n)
>                         return -ENOMEM;
>
> -               *n = *o;
> +               memcpy(n, o, ri_size(o->session_cnt));
>                 get_uprobe(n->uprobe);
>                 n->next = NULL;
>
> @@ -1853,35 +1892,38 @@ 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 struct return_instance *
> +prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
> +                 struct return_instance *ri, int session_cnt)

you have struct uprobe, why do you need to pass session_cnt? Also,
given return_instance is cached, it seems more natural to have

struct return_instance **ri as in/out parameter, and keep the function
itself as void

>  {
> -       struct return_instance *ri;
>         struct uprobe_task *utask;
>         unsigned long orig_ret_vaddr, trampoline_vaddr;
>         bool chained;
>

[...]

>         if (need_prep && !remove)
> -               prepare_uretprobe(uprobe, regs); /* put bp at return */
> +               ri = prepare_uretprobe(uprobe, regs, ri, uprobe->session_cnt); /* put bp at return */
> +       kfree(ri);
>
>         if (remove && uprobe->consumers) {
>                 WARN_ON(!uprobe_is_active(uprobe));
>                 unapply_uprobe(uprobe, current->mm);
>         }
> + out:
>         up_read(&uprobe->register_rwsem);
>  }
>
> +static struct session_consumer *
> +consumer_find(struct session_consumer *sc, struct uprobe_consumer *uc)

why can't we keep track of remaining number of session_consumer items
instead of using entire extra entry as a terminating element? Seems
wasteful and unnecessary.

> +{
> +       for (; sc && sc->id; sc++) {
> +               if (sc->id == uc->id)
> +                       return sc;
> +       }
> +       return NULL;
> +}
> +

[...]

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Mon, Jun 17, 2024 at 03:53:50PM -0700, Andrii Nakryiko wrote:
> On Mon, Jun 10, 2024 at 4:06 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> >
> > On Thu, Jun 06, 2024 at 09:52:39AM -0700, Andrii Nakryiko wrote:
> > > On Thu, Jun 6, 2024 at 9:46 AM Jiri Olsa <olsajiri@gmail.com> wrote:
> > > >
> > > > On Wed, Jun 05, 2024 at 10:50:11PM +0200, Jiri Olsa wrote:
> > > > > On Wed, Jun 05, 2024 at 07:56:19PM +0200, Oleg Nesterov wrote:
> > > > > > On 06/05, Andrii Nakryiko wrote:
> > > > > > >
> > > > > > > so any such
> > > > > > > limitations will cause problems, issue reports, investigation, etc.
> > > > > >
> > > > > > Agreed...
> > > > > >
> > > > > > > As one possible solution, what if we do
> > > > > > >
> > > > > > > struct return_instance {
> > > > > > >     ...
> > > > > > >     u64 session_cookies[];
> > > > > > > };
> > > > > > >
> > > > > > > and allocate sizeof(struct return_instance) + 8 *
> > > > > > > <num-of-session-consumers> and then at runtime pass
> > > > > > > &session_cookies[i] as data pointer to session-aware callbacks?
> > > > > >
> > > > > > I too thought about this, but I guess it is not that simple.
> > > > > >
> > > > > > Just for example. Suppose we have 2 session-consumers C1 and C2.
> > > > > > What if uprobe_unregister(C1) comes before the probed function
> > > > > > returns?
> > > > > >
> > > > > > We need something like map_cookie_to_consumer().
> > > > >
> > > > > I guess we could have hash table in return_instance that gets 'consumer -> cookie' ?
> > > >
> > > > ok, hash table is probably too big for this.. I guess some solution that
> > > > would iterate consumers and cookies made sure it matches would be fine
> > > >
> > >
> > > Yes, I was hoping to avoid hash tables for this, and in the common
> > > case have no added overhead.
> >
> > hi,
> > here's first stab on that.. the change below:
> >   - extends current handlers with extra argument rather than adding new
> >     set of handlers
> >   - store session consumers objects within return_instance object and
> >   - iterate these objects ^^^ in handle_uretprobe_chain
> >
> > I guess it could be still polished, but I wonder if this could
> > be the right direction to do this.. thoughts? ;-)
> 
> Yeah, I think this is the right direction. It's a bit sad that this
> makes getting rid of rw_sem on hot path even harder, but that's a
> separate problem.
> 
> >
> > thanks,
> > jirka
> >
> >
> > ---
> > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > index f46e0ca0169c..4e40e8352eac 100644
> > --- a/include/linux/uprobes.h
> > +++ b/include/linux/uprobes.h
> > @@ -34,15 +34,19 @@ 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,
> > +                       unsigned long *data);
> 
> can we use __u64 here? This long vs __u64 might cause problems for BPF
> when the host is 32-bit architecture (BPF is always 64-bit).

ok

> 
> >         int (*ret_handler)(struct uprobe_consumer *self,
> >                                 unsigned long func,
> > -                               struct pt_regs *regs);
> > +                               struct pt_regs *regs,
> > +                               unsigned long *data);
> >         bool (*filter)(struct uprobe_consumer *self,
> >                                 enum uprobe_filter_ctx ctx,
> >                                 struct mm_struct *mm);
> >
> 
> [...]
> 
> >  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->session_cnt);
> >                 if (!n)
> >                         return -ENOMEM;
> >
> > -               *n = *o;
> > +               memcpy(n, o, ri_size(o->session_cnt));
> >                 get_uprobe(n->uprobe);
> >                 n->next = NULL;
> >
> > @@ -1853,35 +1892,38 @@ 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 struct return_instance *
> > +prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
> > +                 struct return_instance *ri, int session_cnt)
> 
> you have struct uprobe, why do you need to pass session_cnt? Also,
> given return_instance is cached, it seems more natural to have
> 
> struct return_instance **ri as in/out parameter, and keep the function
> itself as void

I tried that, but now I think it'd be better if we just let prepare_uretprobe
to allocate (if needed) and free ri in case it fails and do something like:

       if (need_prep && !remove)
               prepare_uretprobe(uprobe, regs, ri); /* put bp at return */
       else
               kfree(ri);

> 
> >  {
> > -       struct return_instance *ri;
> >         struct uprobe_task *utask;
> >         unsigned long orig_ret_vaddr, trampoline_vaddr;
> >         bool chained;
> >
> 
> [...]
> 
> >         if (need_prep && !remove)
> > -               prepare_uretprobe(uprobe, regs); /* put bp at return */
> > +               ri = prepare_uretprobe(uprobe, regs, ri, uprobe->session_cnt); /* put bp at return */
> > +       kfree(ri);
> >
> >         if (remove && uprobe->consumers) {
> >                 WARN_ON(!uprobe_is_active(uprobe));
> >                 unapply_uprobe(uprobe, current->mm);
> >         }
> > + out:
> >         up_read(&uprobe->register_rwsem);
> >  }
> >
> > +static struct session_consumer *
> > +consumer_find(struct session_consumer *sc, struct uprobe_consumer *uc)
> 
> why can't we keep track of remaining number of session_consumer items
> instead of using entire extra entry as a terminating element? Seems
> wasteful and unnecessary.

ok I think it's possible, will try that

thanks,
jirka

Re: [RFC bpf-next 01/10] uprobe: Add session callbacks to uprobe_consumer

[Jiri Olsa]

On Wed, Jun 05, 2024 at 10:25:56AM -0700, Andrii Nakryiko wrote:

SNIP

> > ---
> >  include/linux/uprobes.h | 18 +++++++++++
> >  kernel/events/uprobes.c | 69 +++++++++++++++++++++++++++++++++++------
> >  2 files changed, 78 insertions(+), 9 deletions(-)
> >
> > diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
> > index f46e0ca0169c..a2f2d5ac3cee 100644
> > --- a/include/linux/uprobes.h
> > +++ b/include/linux/uprobes.h
> > @@ -34,6 +34,12 @@ enum uprobe_filter_ctx {
> >  };
> >
> >  struct uprobe_consumer {
> > +       /*
> > +        * The handler callback return value controls removal of the uprobe.
> > +        *  0 on success, uprobe stays
> > +        *  1 on failure, remove the uprobe
> > +        *    console warning for anything else
> > +        */
> >         int (*handler)(struct uprobe_consumer *self, struct pt_regs *regs);
> >         int (*ret_handler)(struct uprobe_consumer *self,
> >                                 unsigned long func,
> > @@ -42,6 +48,17 @@ struct uprobe_consumer {
> >                                 enum uprobe_filter_ctx ctx,
> >                                 struct mm_struct *mm);
> >
> > +       /* The handler_session callback return value controls execution of
> > +        * the return uprobe and ret_handler_session callback.
> > +        *  0 on success
> > +        *  1 on failure, DO NOT install/execute the return uprobe
> > +        *    console warning for anything else
> > +        */
> > +       int (*handler_session)(struct uprobe_consumer *self, struct pt_regs *regs,
> > +                              unsigned long *data);
> > +       int (*ret_handler_session)(struct uprobe_consumer *self, unsigned long func,
> > +                                  struct pt_regs *regs, unsigned long *data);
> > +
> 
> We should try to avoid an alternative set of callbacks, IMO. Let's
> extend existing ones with `unsigned long *data`, but specify that
> unless consumer sets some flag on registration that it needs a session
> cookie, we'll pass NULL here? Or just allocate cookie data for each
> registered consumer for simplicity, don't know; given we don't expect
> many consumers on exactly the same uprobe, it might be ok to keep it
> simple.
>

ah, I did not want to break existing users.. but it's not uapi,
so we're good, ok makes sense

jirka
 
> 
> >         struct uprobe_consumer *next;
> >  };
> >
> > @@ -85,6 +102,7 @@ struct return_instance {
> >         unsigned long           func;
> >         unsigned long           stack;          /* stack pointer */
> >         unsigned long           orig_ret_vaddr; /* original return address */
> > +       unsigned long           data;
> >         bool                    chained;        /* true, if instance is nested */
> >
> >         struct return_instance  *next;          /* keep as stack */

SNIP

[RFC bpf-next 02/10] 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       | 50 +++++++++++++++++++++++++++++-----
 tools/include/uapi/linux/bpf.h |  1 +
 4 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 25ea393cf084..b400f50e2c3c 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 5070fa20d05c..71d279907a0c 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -4048,10 +4048,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:
@@ -5314,7 +5318,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 f5154c051d2c..53b111c8e887 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -1642,6 +1642,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)
 {
@@ -1659,13 +1670,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:
@@ -3346,6 +3357,26 @@ uprobe_multi_link_ret_handler(struct uprobe_consumer *con, unsigned long func, s
 	return uprobe_prog_run(uprobe, func, regs);
 }
 
+static int
+uprobe_multi_link_handler_session(struct uprobe_consumer *con, struct pt_regs *regs,
+				  unsigned long *data)
+{
+	struct bpf_uprobe *uprobe;
+
+	uprobe = container_of(con, struct bpf_uprobe, consumer);
+	return uprobe_prog_run(uprobe, instruction_pointer(regs), regs);
+}
+
+static int
+uprobe_multi_link_ret_handler_session(struct uprobe_consumer *con, unsigned long func,
+				      struct pt_regs *regs, unsigned long *data)
+{
+	struct bpf_uprobe *uprobe;
+
+	uprobe = container_of(con, struct bpf_uprobe, consumer);
+	return uprobe_prog_run(uprobe, func, regs);
+}
+
 static u64 bpf_uprobe_multi_entry_ip(struct bpf_run_ctx *ctx)
 {
 	struct bpf_uprobe_multi_run_ctx *run_ctx;
@@ -3382,7 +3413,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;
@@ -3460,10 +3491,15 @@ 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
-			uprobes[i].consumer.handler = uprobe_multi_link_handler;
+		if (is_uprobe_session(prog)) {
+			uprobes[i].consumer.handler_session = uprobe_multi_link_handler_session;
+			uprobes[i].consumer.ret_handler_session = uprobe_multi_link_ret_handler_session;
+		} else {
+			if (flags & BPF_F_UPROBE_MULTI_RETURN)
+				uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
+			else
+				uprobes[i].consumer.handler = uprobe_multi_link_handler;
+		}
 
 		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 25ea393cf084..b400f50e2c3c 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.1

[RFC bpf-next 03/10] 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 | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 53b111c8e887..4392807ee8d9 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3181,7 +3181,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;
 };
@@ -3294,10 +3294,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,
 	};
@@ -3316,7 +3321,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);
 
@@ -3345,7 +3350,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, NULL);
 }
 
 static int
@@ -3354,7 +3359,7 @@ 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, NULL);
 }
 
 static int
@@ -3364,7 +3369,7 @@ uprobe_multi_link_handler_session(struct uprobe_consumer *con, struct pt_regs *r
 	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
@@ -3374,14 +3379,14 @@ uprobe_multi_link_ret_handler_session(struct uprobe_consumer *con, unsigned long
 	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;
 }
 
@@ -3389,7 +3394,7 @@ 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;
 }
 
@@ -3586,7 +3591,8 @@ 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.1

[RFC bpf-next 04/10] 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 d1627a2ca30b..a0044448a708 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -9328,6 +9328,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);
 
@@ -9346,6 +9347,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),
@@ -11682,6 +11684,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)
 {
@@ -11916,10 +11952,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;
@@ -11990,12 +12028,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();
@@ -12009,7 +12055,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 26e4e35528c5..04cdb38f527f 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.1

[RFC bpf-next 05/10] 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 a0044448a708..702c2fb7e4df 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.1

[RFC bpf-next 06/10] selftests/bpf: Move ARRAY_SIZE to bpf_misc.h

[Jiri Olsa]

ARRAY_SIZE is used on multiple places, move its definition in
bpf_misc.h header.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 tools/testing/selftests/bpf/progs/bpf_misc.h                 | 2 ++
 tools/testing/selftests/bpf/progs/iters.c                    | 2 --
 tools/testing/selftests/bpf/progs/kprobe_multi_session.c     | 3 +--
 tools/testing/selftests/bpf/progs/linked_list.c              | 5 +----
 tools/testing/selftests/bpf/progs/netif_receive_skb.c        | 5 +----
 tools/testing/selftests/bpf/progs/profiler.inc.h             | 5 +----
 tools/testing/selftests/bpf/progs/setget_sockopt.c           | 5 +----
 tools/testing/selftests/bpf/progs/test_bpf_ma.c              | 4 ----
 tools/testing/selftests/bpf/progs/test_sysctl_loop1.c        | 5 +----
 tools/testing/selftests/bpf/progs/test_sysctl_loop2.c        | 5 +----
 tools/testing/selftests/bpf/progs/test_sysctl_prog.c         | 5 +----
 .../testing/selftests/bpf/progs/test_tcp_custom_syncookie.c  | 1 +
 .../testing/selftests/bpf/progs/test_tcp_custom_syncookie.h  | 2 --
 .../testing/selftests/bpf/progs/verifier_subprog_precision.c | 2 --
 14 files changed, 11 insertions(+), 40 deletions(-)

diff --git a/tools/testing/selftests/bpf/progs/bpf_misc.h b/tools/testing/selftests/bpf/progs/bpf_misc.h
index fb2f5513e29e..70f56af94513 100644
--- a/tools/testing/selftests/bpf/progs/bpf_misc.h
+++ b/tools/testing/selftests/bpf/progs/bpf_misc.h
@@ -135,4 +135,6 @@
 /* make it look to compiler like value is read and written */
 #define __sink(expr) asm volatile("" : "+g"(expr))
 
+#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
+
 #endif
diff --git a/tools/testing/selftests/bpf/progs/iters.c b/tools/testing/selftests/bpf/progs/iters.c
index fe65e0952a1e..16bdc3e25591 100644
--- a/tools/testing/selftests/bpf/progs/iters.c
+++ b/tools/testing/selftests/bpf/progs/iters.c
@@ -7,8 +7,6 @@
 #include "bpf_misc.h"
 #include "bpf_compiler.h"
 
-#define ARRAY_SIZE(x) (int)(sizeof(x) / sizeof((x)[0]))
-
 static volatile int zero = 0;
 
 int my_pid;
diff --git a/tools/testing/selftests/bpf/progs/kprobe_multi_session.c b/tools/testing/selftests/bpf/progs/kprobe_multi_session.c
index bbba9eb46551..bd8b7fb7061e 100644
--- a/tools/testing/selftests/bpf/progs/kprobe_multi_session.c
+++ b/tools/testing/selftests/bpf/progs/kprobe_multi_session.c
@@ -4,8 +4,7 @@
 #include <bpf/bpf_tracing.h>
 #include <stdbool.h>
 #include "bpf_kfuncs.h"
-
-#define ARRAY_SIZE(x) (int)(sizeof(x) / sizeof((x)[0]))
+#include "bpf_misc.h"
 
 char _license[] SEC("license") = "GPL";
 
diff --git a/tools/testing/selftests/bpf/progs/linked_list.c b/tools/testing/selftests/bpf/progs/linked_list.c
index f69bf3e30321..421f40835acd 100644
--- a/tools/testing/selftests/bpf/progs/linked_list.c
+++ b/tools/testing/selftests/bpf/progs/linked_list.c
@@ -4,10 +4,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_core_read.h>
 #include "bpf_experimental.h"
-
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (int)(sizeof(x) / sizeof((x)[0]))
-#endif
+#include "bpf_misc.h"
 
 #include "linked_list.h"
 
diff --git a/tools/testing/selftests/bpf/progs/netif_receive_skb.c b/tools/testing/selftests/bpf/progs/netif_receive_skb.c
index c0062645fc68..9e067dcbf607 100644
--- a/tools/testing/selftests/bpf/progs/netif_receive_skb.c
+++ b/tools/testing/selftests/bpf/progs/netif_receive_skb.c
@@ -5,6 +5,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
 #include <bpf/bpf_core_read.h>
+#include "bpf_misc.h"
 
 #include <errno.h>
 
@@ -23,10 +24,6 @@ bool skip = false;
 #define BADPTR			0
 #endif
 
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x)	(sizeof(x) / sizeof((x)[0]))
-#endif
-
 struct {
 	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
 	__uint(max_entries, 1);
diff --git a/tools/testing/selftests/bpf/progs/profiler.inc.h b/tools/testing/selftests/bpf/progs/profiler.inc.h
index 6957d9f2805e..8bd1ebd7d6af 100644
--- a/tools/testing/selftests/bpf/progs/profiler.inc.h
+++ b/tools/testing/selftests/bpf/progs/profiler.inc.h
@@ -9,6 +9,7 @@
 #include "err.h"
 #include "bpf_experimental.h"
 #include "bpf_compiler.h"
+#include "bpf_misc.h"
 
 #ifndef NULL
 #define NULL 0
@@ -133,10 +134,6 @@ struct {
 	__uint(max_entries, 16);
 } disallowed_exec_inodes SEC(".maps");
 
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(arr) (int)(sizeof(arr) / sizeof(arr[0]))
-#endif
-
 static INLINE bool IS_ERR(const void* ptr)
 {
 	return IS_ERR_VALUE((unsigned long)ptr);
diff --git a/tools/testing/selftests/bpf/progs/setget_sockopt.c b/tools/testing/selftests/bpf/progs/setget_sockopt.c
index 7a438600ae98..60518aed1ffc 100644
--- a/tools/testing/selftests/bpf/progs/setget_sockopt.c
+++ b/tools/testing/selftests/bpf/progs/setget_sockopt.c
@@ -6,10 +6,7 @@
 #include <bpf/bpf_core_read.h>
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
-
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
+#include "bpf_misc.h"
 
 extern unsigned long CONFIG_HZ __kconfig;
 
diff --git a/tools/testing/selftests/bpf/progs/test_bpf_ma.c b/tools/testing/selftests/bpf/progs/test_bpf_ma.c
index 3494ca30fa7f..4a4e0b8d9b72 100644
--- a/tools/testing/selftests/bpf/progs/test_bpf_ma.c
+++ b/tools/testing/selftests/bpf/progs/test_bpf_ma.c
@@ -7,10 +7,6 @@
 #include "bpf_experimental.h"
 #include "bpf_misc.h"
 
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
-
 struct generic_map_value {
 	void *data;
 };
diff --git a/tools/testing/selftests/bpf/progs/test_sysctl_loop1.c b/tools/testing/selftests/bpf/progs/test_sysctl_loop1.c
index 7f74077d6622..548660e299a5 100644
--- a/tools/testing/selftests/bpf/progs/test_sysctl_loop1.c
+++ b/tools/testing/selftests/bpf/progs/test_sysctl_loop1.c
@@ -10,10 +10,7 @@
 #include <bpf/bpf_helpers.h>
 
 #include "bpf_compiler.h"
-
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
+#include "bpf_misc.h"
 
 /* tcp_mem sysctl has only 3 ints, but this test is doing TCP_MEM_LOOPS */
 #define TCP_MEM_LOOPS 28  /* because 30 doesn't fit into 512 bytes of stack */
diff --git a/tools/testing/selftests/bpf/progs/test_sysctl_loop2.c b/tools/testing/selftests/bpf/progs/test_sysctl_loop2.c
index 68a75436e8af..81249d119a8b 100644
--- a/tools/testing/selftests/bpf/progs/test_sysctl_loop2.c
+++ b/tools/testing/selftests/bpf/progs/test_sysctl_loop2.c
@@ -10,10 +10,7 @@
 #include <bpf/bpf_helpers.h>
 
 #include "bpf_compiler.h"
-
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
+#include "bpf_misc.h"
 
 /* tcp_mem sysctl has only 3 ints, but this test is doing TCP_MEM_LOOPS */
 #define TCP_MEM_LOOPS 20  /* because 30 doesn't fit into 512 bytes of stack */
diff --git a/tools/testing/selftests/bpf/progs/test_sysctl_prog.c b/tools/testing/selftests/bpf/progs/test_sysctl_prog.c
index efc3c61f7852..bbdd08764789 100644
--- a/tools/testing/selftests/bpf/progs/test_sysctl_prog.c
+++ b/tools/testing/selftests/bpf/progs/test_sysctl_prog.c
@@ -10,6 +10,7 @@
 #include <bpf/bpf_helpers.h>
 
 #include "bpf_compiler.h"
+#include "bpf_misc.h"
 
 /* Max supported length of a string with unsigned long in base 10 (pow2 - 1). */
 #define MAX_ULONG_STR_LEN 0xF
@@ -17,10 +18,6 @@
 /* Max supported length of sysctl value string (pow2). */
 #define MAX_VALUE_STR_LEN 0x40
 
-#ifndef ARRAY_SIZE
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
-
 const char tcp_mem_name[] = "net/ipv4/tcp_mem";
 static __always_inline int is_tcp_mem(struct bpf_sysctl *ctx)
 {
diff --git a/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.c b/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.c
index c8e4553648bf..44ee0d037f95 100644
--- a/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.c
+++ b/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.c
@@ -9,6 +9,7 @@
 #include "bpf_kfuncs.h"
 #include "test_siphash.h"
 #include "test_tcp_custom_syncookie.h"
+#include "bpf_misc.h"
 
 #define MAX_PACKET_OFF 0xffff
 
diff --git a/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.h b/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.h
index 29a6a53cf229..f8b1b7e68d2e 100644
--- a/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.h
+++ b/tools/testing/selftests/bpf/progs/test_tcp_custom_syncookie.h
@@ -7,8 +7,6 @@
 #define __packed __attribute__((__packed__))
 #define __force
 
-#define ARRAY_SIZE(arr)	(sizeof(arr) / sizeof((arr)[0]))
-
 #define swap(a, b)				\
 	do {					\
 		typeof(a) __tmp = (a);		\
diff --git a/tools/testing/selftests/bpf/progs/verifier_subprog_precision.c b/tools/testing/selftests/bpf/progs/verifier_subprog_precision.c
index 4a58e0398e72..6a6fad625f7e 100644
--- a/tools/testing/selftests/bpf/progs/verifier_subprog_precision.c
+++ b/tools/testing/selftests/bpf/progs/verifier_subprog_precision.c
@@ -8,8 +8,6 @@
 #include "bpf_misc.h"
 #include <../../../tools/include/linux/filter.h>
 
-#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))
-
 int vals[] SEC(".data.vals") = {1, 2, 3, 4};
 
 __naked __noinline __used
-- 
2.45.1

[RFC bpf-next 07/10] 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        | 38 ++++++++++++++
 .../bpf/progs/uprobe_multi_session.c          | 52 +++++++++++++++++++
 2 files changed, 90 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 8269cdee33ae..fddca2597818 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -5,6 +5,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"
 
@@ -497,6 +498,41 @@ 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();
+
+	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_1_result");
+	ASSERT_EQ(skel->bss->uprobe_session_result[2], 1, "uprobe_multi_func_1_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;
@@ -585,4 +621,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..b382d7d29475
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session.c
@@ -0,0 +1,52 @@
+// 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;
+                }
+        }
+
+	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.1

[RFC bpf-next 08/10] selftests/bpf: Add uprobe session errors test

[Jiri Olsa]

Adding uprobe session test to check that just single
session instance is allowed or single uprobe.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

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 fddca2597818..4bff681f0d7d 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -533,6 +533,31 @@ static void test_session_skel_api(void)
 	uprobe_multi_session__destroy(skel);
 }
 
+static void test_session_error_multiple_instances(void)
+{
+	struct uprobe_multi_session *skel_1 = NULL, *skel_2 = NULL;
+	int err;
+
+	skel_1 = uprobe_multi_session__open_and_load();
+	if (!ASSERT_OK_PTR(skel_1, "fentry_raw_skel_load"))
+		goto cleanup;
+
+	err = uprobe_multi_session__attach(skel_1);
+	if (!ASSERT_OK(err, " kprobe_multi_session__attach"))
+		goto cleanup;
+
+	skel_2 = uprobe_multi_session__open_and_load();
+	if (!ASSERT_OK_PTR(skel_2, "fentry_raw_skel_load"))
+		goto cleanup;
+
+	err = uprobe_multi_session__attach(skel_2);
+	ASSERT_EQ(err, -EBUSY, " kprobe_multi_session__attach");
+
+cleanup:
+	uprobe_multi_session__destroy(skel_1);
+	uprobe_multi_session__destroy(skel_2);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -623,4 +648,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_errors"))
+		test_session_error_multiple_instances();
 }
-- 
2.45.1

[RFC bpf-next 09/10] 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   | 50 +++++++++++++++++++
 2 files changed, 81 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 4bff681f0d7d..34671253e130 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_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"
 
@@ -558,6 +559,34 @@ static void test_session_error_multiple_instances(void)
 	uprobe_multi_session__destroy(skel_2);
 }
 
+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;
@@ -650,4 +679,6 @@ void test_uprobe_multi_test(void)
 		test_session_skel_api();
 	if (test__start_subtest("session_errors"))
 		test_session_error_multiple_instances();
+	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..ea41503fad18
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_cookie.c
@@ -0,0 +1,50 @@
+// 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)
+{
+	long *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.1

[RFC bpf-next 10/10] 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 34671253e130..cdd7c327e16d 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_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"
 
@@ -27,6 +28,12 @@ noinline void uprobe_multi_func_3(void)
 	asm volatile ("");
 }
 
+noinline void uprobe_session_recursive(int i)
+{
+	if (i)
+		uprobe_session_recursive(i - 1);
+}
+
 struct child {
 	int go[2];
 	int pid;
@@ -587,6 +594,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;
@@ -681,4 +736,6 @@ void test_uprobe_multi_test(void)
 		test_session_error_multiple_instances();
 	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..7babc180c28f
--- /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)
+{
+	long *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.1