This change allows the uprobe consumer to behave as session which
means that 'handler' and 'ret_handler' callbacks are connected in
a way that allows to:

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

The session 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.

To achive this we are adding new return values the uprobe consumer
can return from 'handler' callback:

  UPROBE_HANDLER_IGNORE
  - Ignore 'ret_handler' callback for this consumer.

  UPROBE_HANDLER_IWANTMYCOOKIE
  - Store cookie and pass it to 'ret_handler' (if defined).

We store shared data in the return_consumer object array as part of
the return_instance object. This way the handle_uretprobe_chain can
find related return_consumer and its shared data.

We also store entry handler return value, for cases when there are
multiple consumers on single uprobe and some of them are ignored and
some of them not, in which case the return probe gets installed and
we need to have a way to find out which consumer needs to be ignored.

The tricky part is when consumer is registered 'after' the uprobe
entry handler is hit. In such case this consumer's 'ret_handler' gets
executed as well, but it won't have the proper data pointer set,
so we can filter it out.

Suggested-by: Oleg Nesterov <oleg@redhat.com>
Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/uprobes.h |  25 ++++++-
 kernel/events/uprobes.c | 150 ++++++++++++++++++++++++++++++++--------
 2 files changed, 144 insertions(+), 31 deletions(-)

diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
index bb265a632b91..751b1dd4abe9 100644
--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h
@@ -23,8 +23,20 @@ struct inode;
 struct notifier_block;
 struct page;
 
+/*
+ * Allowed return values from uprobe consumer's handler callback
+ * with following meaning:
+ *
+ * UPROBE_HANDLER_REMOVE
+ * - Remove the uprobe breakpoint from current->mm.
+ * UPROBE_HANDLER_IGNORE
+ * - Ignore ret_handler callback for this consumer.
+ * UPROBE_HANDLER_IWANTMYCOOKIE
+ * - Store cookie and pass it to ret_handler (if defined).
+ */
 #define UPROBE_HANDLER_REMOVE		1
-#define UPROBE_HANDLER_MASK		1
+#define UPROBE_HANDLER_IGNORE		2
+#define UPROBE_HANDLER_IWANTMYCOOKIE	3
 
 #define MAX_URETPROBE_DEPTH		64
 
@@ -44,6 +56,8 @@ struct uprobe_consumer {
 	bool (*filter)(struct uprobe_consumer *self, struct mm_struct *mm);
 
 	struct list_head cons_node;
+
+	__u64 id;	/* set when uprobe_consumer is registered */
 };
 
 #ifdef CONFIG_UPROBES
@@ -83,14 +97,23 @@ struct uprobe_task {
 	unsigned int			depth;
 };
 
+struct return_consumer {
+	__u64	cookie;
+	__u64	id;
+	int	rc;
+};
+
 struct return_instance {
 	struct uprobe		*uprobe;
 	unsigned long		func;
 	unsigned long		stack;		/* stack pointer */
 	unsigned long		orig_ret_vaddr; /* original return address */
 	bool			chained;	/* true, if instance is nested */
+	int			consumers_cnt;
 
 	struct return_instance	*next;		/* keep as stack */
+
+	struct return_consumer	consumers[] __counted_by(consumers_cnt);
 };
 
 enum rp_check {
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index 4b7e590dc428..5d376cc9a983 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -65,7 +65,7 @@ struct uprobe {
 	struct rcu_head		rcu;
 	loff_t			offset;
 	loff_t			ref_ctr_offset;
-	unsigned long		flags;
+	unsigned long		flags;		/* "unsigned long" so bitops work */
 
 	/*
 	 * The generic code assumes that it has two members of unknown type
@@ -826,8 +826,11 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
 
 static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
 {
+	static atomic64_t id;
+
 	down_write(&uprobe->consumer_rwsem);
 	list_add_rcu(&uc->cons_node, &uprobe->consumers);
+	uc->id = (__u64) atomic64_inc_return(&id);
 	up_write(&uprobe->consumer_rwsem);
 }
 
@@ -1786,6 +1789,34 @@ static struct uprobe_task *get_utask(void)
 	return current->utask;
 }
 
+static size_t ri_size(int consumers_cnt)
+{
+	struct return_instance *ri;
+
+	return sizeof(*ri) + sizeof(ri->consumers[0]) * consumers_cnt;
+}
+
+#define DEF_CNT 4
+
+static struct return_instance *alloc_return_instance(void)
+{
+	struct return_instance *ri;
+
+	ri = kzalloc(ri_size(DEF_CNT), GFP_KERNEL);
+	if (!ri)
+		return ZERO_SIZE_PTR;
+
+	ri->consumers_cnt = DEF_CNT;
+	return ri;
+}
+
+static struct return_instance *dup_return_instance(struct return_instance *old)
+{
+	size_t size = ri_size(old->consumers_cnt);
+
+	return kmemdup(old, size, GFP_KERNEL);
+}
+
 static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
 {
 	struct uprobe_task *n_utask;
@@ -1798,11 +1829,10 @@ 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 = dup_return_instance(o);
 		if (!n)
 			return -ENOMEM;
 
-		*n = *o;
 		/*
 		 * uprobe's refcnt has to be positive at this point, kept by
 		 * utask->return_instances items; return_instances can't be
@@ -1895,39 +1925,35 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
 	utask->return_instances = ri;
 }
 
-static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
+static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
+			      struct return_instance *ri)
 {
-	struct return_instance *ri;
 	struct uprobe_task *utask;
 	unsigned long orig_ret_vaddr, trampoline_vaddr;
 	bool chained;
 
 	if (!get_xol_area())
-		return;
+		goto free;
 
 	utask = get_utask();
 	if (!utask)
-		return;
+		goto free;
 
 	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;
+		goto free;
 	}
 
 	/* we need to bump refcount to store uprobe in utask */
 	if (!try_get_uprobe(uprobe))
-		return;
-
-	ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
-	if (!ri)
-		goto fail;
+		goto free;
 
 	trampoline_vaddr = uprobe_get_trampoline_vaddr();
 	orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
 	if (orig_ret_vaddr == -1)
-		goto fail;
+		goto put;
 
 	/* drop the entries invalidated by longjmp() */
 	chained = (orig_ret_vaddr == trampoline_vaddr);
@@ -1945,7 +1971,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 			 * attack from user-space.
 			 */
 			uprobe_warn(current, "handle tail call");
-			goto fail;
+			goto put;
 		}
 		orig_ret_vaddr = utask->return_instances->orig_ret_vaddr;
 	}
@@ -1960,9 +1986,10 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 	utask->return_instances = ri;
 
 	return;
-fail:
-	kfree(ri);
+put:
 	put_uprobe(uprobe);
+free:
+	kfree(ri);
 }
 
 /* Prepare to single-step probed instruction out of line. */
@@ -2114,35 +2141,94 @@ static struct uprobe *find_active_uprobe_rcu(unsigned long bp_vaddr, int *is_swb
 	return uprobe;
 }
 
+static struct return_instance*
+push_consumer(struct return_instance *ri, int idx, __u64 id, __u64 cookie, int rc)
+{
+	if (unlikely(ri == ZERO_SIZE_PTR))
+		return ri;
+
+	if (unlikely(idx >= ri->consumers_cnt)) {
+		struct return_instance *old_ri = ri;
+
+		ri->consumers_cnt += DEF_CNT;
+		ri = krealloc(old_ri, ri_size(old_ri->consumers_cnt), GFP_KERNEL);
+		if (!ri) {
+			kfree(old_ri);
+			return ZERO_SIZE_PTR;
+		}
+	}
+
+	ri->consumers[idx].id = id;
+	ri->consumers[idx].cookie = cookie;
+	ri->consumers[idx].rc = rc;
+	return ri;
+}
+
+static struct return_consumer *
+return_consumer_find(struct return_instance *ri, int *iter, int id)
+{
+	struct return_consumer *ric;
+	int idx = *iter;
+
+	for (ric = &ri->consumers[idx]; idx < ri->consumers_cnt; idx++, ric++) {
+		if (ric->id == id) {
+			*iter = idx + 1;
+			return ric;
+		}
+	}
+	return NULL;
+}
+
 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 */
-	bool has_consumers = false;
+	bool has_consumers = false, remove = true, ignore = true;
+	struct return_instance *ri = NULL;
+	int push_idx = 0;
 
 	current->utask->auprobe = &uprobe->arch;
 
 	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
 				 srcu_read_lock_held(&uprobes_srcu)) {
+		__u64 cookie = 0;
 		int rc = 0;
 
 		if (uc->handler) {
-			rc = uc->handler(uc, regs);
-			WARN(rc & ~UPROBE_HANDLER_MASK,
+			rc = uc->handler(uc, regs, &cookie);
+			WARN(rc < 0 || rc > 3,
 				"bad rc=0x%x from %ps()\n", rc, uc->handler);
 		}
 
-		if (uc->ret_handler)
-			need_prep = true;
-
-		remove &= rc;
+		remove &= rc == UPROBE_HANDLER_REMOVE;
+		ignore &= rc == UPROBE_HANDLER_IGNORE;
 		has_consumers = true;
+
+		if (!uc->ret_handler || rc == UPROBE_HANDLER_REMOVE)
+			continue;
+
+		/*
+		 * If alloc_return_instance and push_consumer fail, the return probe
+		 * won't be prepared, but we'll finish to execute all entry handlers.
+		 *
+		 * We need to store handler's return value in case the return uprobe
+		 * gets installed and contains consumers that need to be ignored.
+		 */
+		if (!ri)
+			ri = alloc_return_instance();
+
+		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE || rc == UPROBE_HANDLER_IGNORE)
+			ri = push_consumer(ri, push_idx++, uc->id, cookie, rc);
 	}
 	current->utask->auprobe = NULL;
 
-	if (need_prep && !remove)
-		prepare_uretprobe(uprobe, regs); /* put bp at return */
+	if (!ignore && !ZERO_OR_NULL_PTR(ri)) {
+		/*
+		 * The push_idx value has the final number of return consumers,
+		 * and ri->consumers_cnt has number of allocated consumers.
+		 */
+		ri->consumers_cnt = push_idx;
+		prepare_uretprobe(uprobe, regs, ri);
+	}
 
 	if (remove && has_consumers) {
 		down_read(&uprobe->register_rwsem);
@@ -2161,14 +2247,18 @@ static void
 handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
 {
 	struct uprobe *uprobe = ri->uprobe;
+	struct return_consumer *ric;
 	struct uprobe_consumer *uc;
-	int srcu_idx;
+	int srcu_idx, ric_idx = 0;
 
 	srcu_idx = srcu_read_lock(&uprobes_srcu);
 	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
 				 srcu_read_lock_held(&uprobes_srcu)) {
+		ric = return_consumer_find(ri, &ric_idx, uc->id);
+		if (ric && ric->rc == UPROBE_HANDLER_IGNORE)
+			continue;
 		if (uc->ret_handler)
-			uc->ret_handler(uc, ri->func, regs);
+			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
 	}
 	srcu_read_unlock(&uprobes_srcu, srcu_idx);
 }
-- 
2.46.0

On 09/17, Jiri Olsa wrote:
>
>  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
...
> +	if (!ignore && !ZERO_OR_NULL_PTR(ri)) {
> +		/*
> +		 * The push_idx value has the final number of return consumers,
> +		 * and ri->consumers_cnt has number of allocated consumers.
> +		 */
> +		ri->consumers_cnt = push_idx;
> +		prepare_uretprobe(uprobe, regs, ri);
> +	}

This looks wrong. ri is not kfreed if ignore == true.

But see my previous email, if we change this code as I tried to suggest
the problem goes away and handler_chain() doesn't need "bool ignore".

Oleg.

I don't see anything wrong after a quick glance, but I don't
really understand the UPROBE_HANDLER_IGNORE logic, see below.

On 09/17, Jiri Olsa wrote:
>
> + * UPROBE_HANDLER_IWANTMYCOOKIE
> + * - Store cookie and pass it to ret_handler (if defined).

Cough ;) yes it was me who used this name in the previous discussion, but maybe

	UPROBE_HANDLER_COOKIE

will look a bit better? Feel free to ignore.

>  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
...
> +		if (!uc->ret_handler || rc == UPROBE_HANDLER_REMOVE)
> +			continue;
> +
> +		/*
> +		 * If alloc_return_instance and push_consumer fail, the return probe
> +		 * won't be prepared, but we'll finish to execute all entry handlers.
> +		 *
> +		 * We need to store handler's return value in case the return uprobe
> +		 * gets installed and contains consumers that need to be ignored.
> +		 */
> +		if (!ri)
> +			ri = alloc_return_instance();
> +
> +		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE || rc == UPROBE_HANDLER_IGNORE)
> +			ri = push_consumer(ri, push_idx++, uc->id, cookie, rc);

So this code allocates ri (which implies prepare_uretprobe!) and calls push_consumer()
even if rc == UPROBE_HANDLER_IGNORE.

Why? The comment in uprobes.h says:

	UPROBE_HANDLER_IGNORE
	- Ignore ret_handler callback for this consumer

but the ret_handler callback won't be ignored?

To me this code should do:

		if (!uc->ret_handler || UPROBE_HANDLER_REMOVE || UPROBE_HANDLER_IGNORE)
			continue;

		if (!ri)
			ri = alloc_return_instance();

		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE)
			ri = push_consumer(...);

And,

>  handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
...
>  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
>  				 srcu_read_lock_held(&uprobes_srcu)) {
> +		ric = return_consumer_find(ri, &ric_idx, uc->id);
> +		if (ric && ric->rc == UPROBE_HANDLER_IGNORE)
> +			continue;
>  		if (uc->ret_handler)
> -			uc->ret_handler(uc, ri->func, regs);
> +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
>  	}

the UPROBE_HANDLER_IGNORE check above and the new ric->rc member should die,

		if (!uc->ret_handler)
			continue;

		ric = return_consumer_find(...);
		uc->ret_handler(..., ric ? &ric->cookie : NULL);

as we have already discussed, the session ret_handler(data) can simply do

		// my ->handler() wasn't called or it didn't return
		// UPROBE_HANDLER_IWANTMYCOOKIE
		if (!data)
			return;

at the start.

Could you explain why this can't work?

Oleg.

On Tue, Sep 17, 2024 at 02:03:17PM +0200, Oleg Nesterov wrote:
> I don't see anything wrong after a quick glance, but I don't
> really understand the UPROBE_HANDLER_IGNORE logic, see below.
> 
> On 09/17, Jiri Olsa wrote:
> >
> > + * UPROBE_HANDLER_IWANTMYCOOKIE
> > + * - Store cookie and pass it to ret_handler (if defined).
> 
> Cough ;) yes it was me who used this name in the previous discussion, but maybe
> 
> 	UPROBE_HANDLER_COOKIE
> 
> will look a bit better? Feel free to ignore.

ok, no fun it is..

> 
> >  static void handler_chain(struct uprobe *uprobe, struct pt_regs *regs)
> ...
> > +		if (!uc->ret_handler || rc == UPROBE_HANDLER_REMOVE)
> > +			continue;
> > +
> > +		/*
> > +		 * If alloc_return_instance and push_consumer fail, the return probe
> > +		 * won't be prepared, but we'll finish to execute all entry handlers.
> > +		 *
> > +		 * We need to store handler's return value in case the return uprobe
> > +		 * gets installed and contains consumers that need to be ignored.
> > +		 */
> > +		if (!ri)
> > +			ri = alloc_return_instance();
> > +
> > +		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE || rc == UPROBE_HANDLER_IGNORE)
> > +			ri = push_consumer(ri, push_idx++, uc->id, cookie, rc);
> 
> So this code allocates ri (which implies prepare_uretprobe!) and calls push_consumer()
> even if rc == UPROBE_HANDLER_IGNORE.
> 
> Why? The comment in uprobes.h says:
> 
> 	UPROBE_HANDLER_IGNORE
> 	- Ignore ret_handler callback for this consumer
> 
> but the ret_handler callback won't be ignored?
> 
> To me this code should do:
> 
> 		if (!uc->ret_handler || UPROBE_HANDLER_REMOVE || UPROBE_HANDLER_IGNORE)
> 			continue;
> 
> 		if (!ri)
> 			ri = alloc_return_instance();
> 
> 		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE)
> 			ri = push_consumer(...);
> 
> And,
> 
> >  handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
> ...
> >  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
> >  				 srcu_read_lock_held(&uprobes_srcu)) {
> > +		ric = return_consumer_find(ri, &ric_idx, uc->id);
> > +		if (ric && ric->rc == UPROBE_HANDLER_IGNORE)
> > +			continue;
> >  		if (uc->ret_handler)
> > -			uc->ret_handler(uc, ri->func, regs);
> > +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
> >  	}
> 
> the UPROBE_HANDLER_IGNORE check above and the new ric->rc member should die,
> 
> 		if (!uc->ret_handler)
> 			continue;
> 
> 		ric = return_consumer_find(...);
> 		uc->ret_handler(..., ric ? &ric->cookie : NULL);
> 
> as we have already discussed, the session ret_handler(data) can simply do
> 
> 		// my ->handler() wasn't called or it didn't return
> 		// UPROBE_HANDLER_IWANTMYCOOKIE
> 		if (!data)
> 			return;
> 
> at the start.
> 
> Could you explain why this can't work?

I'll try ;-) it's for the case when consumer does not use UPROBE_HANDLER_IWANTMYCOOKIE

let's have 2 consumers on single uprobe, consumer-A returning UPROBE_HANDLER_IGNORE
and the consumer-B returning zero, so we want the return uprobe installed, but we
want just consumer-B to be executed

  - so uprobe gets installed and handle_uretprobe_chain goes over all consumers
    calling ret_handler callback

  - but we don't know consumer-A needs to be ignored, and it does not
    expect cookie so we have no way to find out it needs to be ignored

the change solves this by storing also return value for consumer

if all consumers ignore the ret_handler callback return uprobe is not installed

jirka

Damn, sorry for delay :/

And sorry, still can't understand, see below...

On 09/17, Jiri Olsa wrote:
>
> On Tue, Sep 17, 2024 at 02:03:17PM +0200, Oleg Nesterov wrote:
> >
> > To me this code should do:
> >
> > 		if (!uc->ret_handler || UPROBE_HANDLER_REMOVE || UPROBE_HANDLER_IGNORE)
> > 			continue;
> >
> > 		if (!ri)
> > 			ri = alloc_return_instance();
> >
> > 		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE)
> > 			ri = push_consumer(...);
> >
> > And,
> >
> > >  handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
> > ...
> > >  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
> > >  				 srcu_read_lock_held(&uprobes_srcu)) {
> > > +		ric = return_consumer_find(ri, &ric_idx, uc->id);
> > > +		if (ric && ric->rc == UPROBE_HANDLER_IGNORE)
> > > +			continue;
> > >  		if (uc->ret_handler)
> > > -			uc->ret_handler(uc, ri->func, regs);
> > > +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
> > >  	}
> >
> > the UPROBE_HANDLER_IGNORE check above and the new ric->rc member should die,
> >
> > 		if (!uc->ret_handler)
> > 			continue;
> >
> > 		ric = return_consumer_find(...);
> > 		uc->ret_handler(..., ric ? &ric->cookie : NULL);
> >
> > as we have already discussed, the session ret_handler(data) can simply do
> >
> > 		// my ->handler() wasn't called or it didn't return
> > 		// UPROBE_HANDLER_IWANTMYCOOKIE
> > 		if (!data)
> > 			return;
> >
> > at the start.
> >
> > Could you explain why this can't work?
>
> I'll try ;-) it's for the case when consumer does not use UPROBE_HANDLER_IWANTMYCOOKIE
>
> let's have 2 consumers on single uprobe, consumer-A returning UPROBE_HANDLER_IGNORE
> and the consumer-B returning zero, so we want the return uprobe installed, but we
> want just consumer-B to be executed
>
>   - so uprobe gets installed and handle_uretprobe_chain goes over all consumers
>     calling ret_handler callback
>
>   - but we don't know consumer-A needs to be ignored, and it does not
>     expect cookie so we have no way to find out it needs to be ignored

How does this differ from the case when consumer-A returns _REMOVE but another
consumer returns 0?

But what I really can't understand is

	and it does not
	expect cookie so we have no way to find out it needs to be ignored

If we change the code as I suggested above, push_consumer() won't be called
if consumer-A returns UPROBE_HANDLER_IGNORE.

This means that handle_uretprobe_chain() -> return_consumer_find() will
return NULL, so handle_uretprobe_chain() won't pass the valid cookie to
consumer-A's ret_handler callback, it will pass data => NULL.

So, again, why can't consumer-A's ret_handler callback do

	// my ->handler() wasn't called or it didn't return
	// UPROBE_HANDLER_IWANTMYCOOKIE
	if (!data)
		return;

at the start?

Why the UPROBE_HANDLER_IGNORE case is more problematic than the
UPROBE_HANDLER_REMOVE case?

Oleg.

On Sun, Sep 22, 2024 at 05:27:23PM +0200, Oleg Nesterov wrote:
> Damn, sorry for delay :/
> 
> And sorry, still can't understand, see below...
> 
> On 09/17, Jiri Olsa wrote:
> >
> > On Tue, Sep 17, 2024 at 02:03:17PM +0200, Oleg Nesterov wrote:
> > >
> > > To me this code should do:
> > >
> > > 		if (!uc->ret_handler || UPROBE_HANDLER_REMOVE || UPROBE_HANDLER_IGNORE)
> > > 			continue;
> > >
> > > 		if (!ri)
> > > 			ri = alloc_return_instance();
> > >
> > > 		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE)
> > > 			ri = push_consumer(...);
> > >
> > > And,
> > >
> > > >  handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
> > > ...
> > > >  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
> > > >  				 srcu_read_lock_held(&uprobes_srcu)) {
> > > > +		ric = return_consumer_find(ri, &ric_idx, uc->id);
> > > > +		if (ric && ric->rc == UPROBE_HANDLER_IGNORE)
> > > > +			continue;
> > > >  		if (uc->ret_handler)
> > > > -			uc->ret_handler(uc, ri->func, regs);
> > > > +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
> > > >  	}
> > >
> > > the UPROBE_HANDLER_IGNORE check above and the new ric->rc member should die,
> > >
> > > 		if (!uc->ret_handler)
> > > 			continue;
> > >
> > > 		ric = return_consumer_find(...);
> > > 		uc->ret_handler(..., ric ? &ric->cookie : NULL);
> > >
> > > as we have already discussed, the session ret_handler(data) can simply do
> > >
> > > 		// my ->handler() wasn't called or it didn't return
> > > 		// UPROBE_HANDLER_IWANTMYCOOKIE
> > > 		if (!data)
> > > 			return;
> > >
> > > at the start.
> > >
> > > Could you explain why this can't work?
> >
> > I'll try ;-) it's for the case when consumer does not use UPROBE_HANDLER_IWANTMYCOOKIE
> >
> > let's have 2 consumers on single uprobe, consumer-A returning UPROBE_HANDLER_IGNORE
> > and the consumer-B returning zero, so we want the return uprobe installed, but we
> > want just consumer-B to be executed
> >
> >   - so uprobe gets installed and handle_uretprobe_chain goes over all consumers
> >     calling ret_handler callback
> >
> >   - but we don't know consumer-A needs to be ignored, and it does not
> >     expect cookie so we have no way to find out it needs to be ignored
> 
> How does this differ from the case when consumer-A returns _REMOVE but another
> consumer returns 0?
> 
> But what I really can't understand is
> 
> 	and it does not
> 	expect cookie so we have no way to find out it needs to be ignored
> 
> If we change the code as I suggested above, push_consumer() won't be called
> if consumer-A returns UPROBE_HANDLER_IGNORE.
> 
> This means that handle_uretprobe_chain() -> return_consumer_find() will
> return NULL, so handle_uretprobe_chain() won't pass the valid cookie to
> consumer-A's ret_handler callback, it will pass data => NULL.
> 
> So, again, why can't consumer-A's ret_handler callback do
> 
> 	// my ->handler() wasn't called or it didn't return
> 	// UPROBE_HANDLER_IWANTMYCOOKIE
> 	if (!data)
> 		return;

ok, I think I understand the issue now.. it all depends on 'handler' to return
either UPROBE_HANDLER_IGNORE or UPROBE_HANDLER_IWANTMYCOOKIE

my idea was to make the interface more generic, so some future uprobe user won't
depend on handler callback to return UPROBE_HANDLER_IWANTMYCOOKIE and can just
return 0, but we can do that as follow up if it's ever needed

change below should do what you proposed originally

also on top of that.. I discussed with Andrii the possibility of dropping
the UPROBE_HANDLER_IWANTMYCOOKIE completely and setup cookie for any consumer
that has both 'handler' and 'ret_handler' defined, wdyt?

thanks,
jirka


---
diff --git a/include/linux/uprobes.h b/include/linux/uprobes.h
index bb265a632b91..5221080b1f5a 100644
--- a/include/linux/uprobes.h
+++ b/include/linux/uprobes.h
@@ -23,8 +23,20 @@ struct inode;
 struct notifier_block;
 struct page;
 
+/*
+ * Allowed return values from uprobe consumer's handler callback
+ * with following meaning:
+ *
+ * UPROBE_HANDLER_REMOVE
+ * - Remove the uprobe breakpoint from current->mm.
+ * UPROBE_HANDLER_IGNORE
+ * - Ignore ret_handler callback for this consumer.
+ * UPROBE_HANDLER_IWANTMYCOOKIE
+ * - Store cookie and pass it to ret_handler (if defined).
+ */
 #define UPROBE_HANDLER_REMOVE		1
-#define UPROBE_HANDLER_MASK		1
+#define UPROBE_HANDLER_IGNORE		2
+#define UPROBE_HANDLER_IWANTMYCOOKIE	3
 
 #define MAX_URETPROBE_DEPTH		64
 
@@ -44,6 +56,8 @@ struct uprobe_consumer {
 	bool (*filter)(struct uprobe_consumer *self, struct mm_struct *mm);
 
 	struct list_head cons_node;
+
+	__u64 id;	/* set when uprobe_consumer is registered */
 };
 
 #ifdef CONFIG_UPROBES
@@ -83,14 +97,22 @@ struct uprobe_task {
 	unsigned int			depth;
 };
 
+struct return_consumer {
+	__u64	cookie;
+	__u64	id;
+};
+
 struct return_instance {
 	struct uprobe		*uprobe;
 	unsigned long		func;
 	unsigned long		stack;		/* stack pointer */
 	unsigned long		orig_ret_vaddr; /* original return address */
 	bool			chained;	/* true, if instance is nested */
+	int			consumers_cnt;
 
 	struct return_instance	*next;		/* keep as stack */
+
+	struct return_consumer	consumers[] __counted_by(consumers_cnt);
 };
 
 enum rp_check {
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index 4b7e590dc428..0dca2f2ecf9c 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -65,7 +65,7 @@ struct uprobe {
 	struct rcu_head		rcu;
 	loff_t			offset;
 	loff_t			ref_ctr_offset;
-	unsigned long		flags;
+	unsigned long		flags;		/* "unsigned long" so bitops work */
 
 	/*
 	 * The generic code assumes that it has two members of unknown type
@@ -826,8 +826,11 @@ static struct uprobe *alloc_uprobe(struct inode *inode, loff_t offset,
 
 static void consumer_add(struct uprobe *uprobe, struct uprobe_consumer *uc)
 {
+	static atomic64_t id;
+
 	down_write(&uprobe->consumer_rwsem);
 	list_add_rcu(&uc->cons_node, &uprobe->consumers);
+	uc->id = (__u64) atomic64_inc_return(&id);
 	up_write(&uprobe->consumer_rwsem);
 }
 
@@ -1786,6 +1789,34 @@ static struct uprobe_task *get_utask(void)
 	return current->utask;
 }
 
+static size_t ri_size(int consumers_cnt)
+{
+	struct return_instance *ri;
+
+	return sizeof(*ri) + sizeof(ri->consumers[0]) * consumers_cnt;
+}
+
+#define DEF_CNT 4
+
+static struct return_instance *alloc_return_instance(void)
+{
+	struct return_instance *ri;
+
+	ri = kzalloc(ri_size(DEF_CNT), GFP_KERNEL);
+	if (!ri)
+		return ZERO_SIZE_PTR;
+
+	ri->consumers_cnt = DEF_CNT;
+	return ri;
+}
+
+static struct return_instance *dup_return_instance(struct return_instance *old)
+{
+	size_t size = ri_size(old->consumers_cnt);
+
+	return kmemdup(old, size, GFP_KERNEL);
+}
+
 static int dup_utask(struct task_struct *t, struct uprobe_task *o_utask)
 {
 	struct uprobe_task *n_utask;
@@ -1798,11 +1829,10 @@ 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 = dup_return_instance(o);
 		if (!n)
 			return -ENOMEM;
 
-		*n = *o;
 		/*
 		 * uprobe's refcnt has to be positive at this point, kept by
 		 * utask->return_instances items; return_instances can't be
@@ -1895,39 +1925,35 @@ static void cleanup_return_instances(struct uprobe_task *utask, bool chained,
 	utask->return_instances = ri;
 }
 
-static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
+static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs,
+			      struct return_instance *ri)
 {
-	struct return_instance *ri;
 	struct uprobe_task *utask;
 	unsigned long orig_ret_vaddr, trampoline_vaddr;
 	bool chained;
 
 	if (!get_xol_area())
-		return;
+		goto free;
 
 	utask = get_utask();
 	if (!utask)
-		return;
+		goto free;
 
 	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;
+		goto free;
 	}
 
 	/* we need to bump refcount to store uprobe in utask */
 	if (!try_get_uprobe(uprobe))
-		return;
-
-	ri = kmalloc(sizeof(struct return_instance), GFP_KERNEL);
-	if (!ri)
-		goto fail;
+		goto free;
 
 	trampoline_vaddr = uprobe_get_trampoline_vaddr();
 	orig_ret_vaddr = arch_uretprobe_hijack_return_addr(trampoline_vaddr, regs);
 	if (orig_ret_vaddr == -1)
-		goto fail;
+		goto put;
 
 	/* drop the entries invalidated by longjmp() */
 	chained = (orig_ret_vaddr == trampoline_vaddr);
@@ -1945,7 +1971,7 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 			 * attack from user-space.
 			 */
 			uprobe_warn(current, "handle tail call");
-			goto fail;
+			goto put;
 		}
 		orig_ret_vaddr = utask->return_instances->orig_ret_vaddr;
 	}
@@ -1960,9 +1986,10 @@ static void prepare_uretprobe(struct uprobe *uprobe, struct pt_regs *regs)
 	utask->return_instances = ri;
 
 	return;
-fail:
-	kfree(ri);
+put:
 	put_uprobe(uprobe);
+free:
+	kfree(ri);
 }
 
 /* Prepare to single-step probed instruction out of line. */
@@ -2114,35 +2141,90 @@ static struct uprobe *find_active_uprobe_rcu(unsigned long bp_vaddr, int *is_swb
 	return uprobe;
 }
 
+static struct return_instance*
+push_consumer(struct return_instance *ri, int idx, __u64 id, __u64 cookie)
+{
+	if (unlikely(ri == ZERO_SIZE_PTR))
+		return ri;
+
+	if (unlikely(idx >= ri->consumers_cnt)) {
+		struct return_instance *old_ri = ri;
+
+		ri->consumers_cnt += DEF_CNT;
+		ri = krealloc(old_ri, ri_size(old_ri->consumers_cnt), GFP_KERNEL);
+		if (!ri) {
+			kfree(old_ri);
+			return ZERO_SIZE_PTR;
+		}
+	}
+
+	ri->consumers[idx].id = id;
+	ri->consumers[idx].cookie = cookie;
+	return ri;
+}
+
+static struct return_consumer *
+return_consumer_find(struct return_instance *ri, int *iter, int id)
+{
+	struct return_consumer *ric;
+	int idx = *iter;
+
+	for (ric = &ri->consumers[idx]; idx < ri->consumers_cnt; idx++, ric++) {
+		if (ric->id == id) {
+			*iter = idx + 1;
+			return ric;
+		}
+	}
+	return NULL;
+}
+
+static bool ignore_ret_handler(int rc)
+{
+	return rc == UPROBE_HANDLER_REMOVE || rc == UPROBE_HANDLER_IGNORE;
+}
+
 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 */
-	bool has_consumers = false;
+	bool has_consumers = false, remove = true;
+	struct return_instance *ri = NULL;
+	int push_idx = 0;
 
 	current->utask->auprobe = &uprobe->arch;
 
 	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
 				 srcu_read_lock_held(&uprobes_srcu)) {
+		__u64 cookie = 0;
 		int rc = 0;
 
 		if (uc->handler) {
-			rc = uc->handler(uc, regs);
-			WARN(rc & ~UPROBE_HANDLER_MASK,
+			rc = uc->handler(uc, regs, &cookie);
+			WARN(rc < 0 || rc > 3,
 				"bad rc=0x%x from %ps()\n", rc, uc->handler);
 		}
 
-		if (uc->ret_handler)
-			need_prep = true;
-
-		remove &= rc;
+		remove &= rc == UPROBE_HANDLER_REMOVE;
 		has_consumers = true;
+
+		if (!uc->ret_handler || ignore_ret_handler(rc))
+			continue;
+
+		if (!ri)
+			ri = alloc_return_instance();
+
+		if (rc == UPROBE_HANDLER_IWANTMYCOOKIE)
+			ri = push_consumer(ri, push_idx++, uc->id, cookie);
 	}
 	current->utask->auprobe = NULL;
 
-	if (need_prep && !remove)
-		prepare_uretprobe(uprobe, regs); /* put bp at return */
+	if (!ZERO_OR_NULL_PTR(ri)) {
+		/*
+		 * The push_idx value has the final number of return consumers,
+		 * and ri->consumers_cnt has number of allocated consumers.
+		 */
+		ri->consumers_cnt = push_idx;
+		prepare_uretprobe(uprobe, regs, ri);
+	}
 
 	if (remove && has_consumers) {
 		down_read(&uprobe->register_rwsem);
@@ -2161,14 +2243,16 @@ static void
 handle_uretprobe_chain(struct return_instance *ri, struct pt_regs *regs)
 {
 	struct uprobe *uprobe = ri->uprobe;
+	struct return_consumer *ric;
 	struct uprobe_consumer *uc;
-	int srcu_idx;
+	int srcu_idx, ric_idx = 0;
 
 	srcu_idx = srcu_read_lock(&uprobes_srcu);
 	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
 				 srcu_read_lock_held(&uprobes_srcu)) {
+		ric = return_consumer_find(ri, &ric_idx, uc->id);
 		if (uc->ret_handler)
-			uc->ret_handler(uc, ri->func, regs);
+			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
 	}
 	srcu_read_unlock(&uprobes_srcu, srcu_idx);
 }

On 09/23, Jiri Olsa wrote:
>
> change below should do what you proposed originally

LGTM, just one nit below.

But I guess you need to do this on top of bpf/bpf.git, Andrii has already
applied your series.

And to remind, 02/14 must be fixed in any case unless I am totally confused,
handler_chain() can leak return_instance.

> also on top of that.. I discussed with Andrii the possibility of dropping
> the UPROBE_HANDLER_IWANTMYCOOKIE completely and setup cookie for any consumer
> that has both 'handler' and 'ret_handler' defined, wdyt?

Up to you. As I said from the very beginning I won't insist on _IWANTMYCOOKIE.

>  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
>  				 srcu_read_lock_held(&uprobes_srcu)) {
> +		ric = return_consumer_find(ri, &ric_idx, uc->id);
>  		if (uc->ret_handler)
> -			uc->ret_handler(uc, ri->func, regs);
> +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
>  	}
>  	srcu_read_unlock(&uprobes_srcu, srcu_idx);

return_consumer_find() makes no sense if !uc->ret_handler, can you move

		ric = return_consumer_find(ri, &ric_idx, uc->id);

into the "if (uc->ret_handler)" block?

Oleg.

On Mon, Sep 23, 2024 at 12:05:53PM +0200, Oleg Nesterov wrote:
> On 09/23, Jiri Olsa wrote:
> >
> > change below should do what you proposed originally
> 
> LGTM, just one nit below.
> 
> But I guess you need to do this on top of bpf/bpf.git, Andrii has already
> applied your series.

that seems confusing but looks like just that one fix with the
commit link in [1] was applied

[1] https://lore.kernel.org/bpf/172708047825.3261420.5126267811201364070.git-patchwork-notify@kernel.org/T/#mb065649b5ab8f7ea5b03c215bdc6555a0b76c0d7

> 
> And to remind, 02/14 must be fixed in any case unless I am totally confused,
> handler_chain() can leak return_instance.

yep it was missing kfree, but it's not needed in this new version

> 
> > also on top of that.. I discussed with Andrii the possibility of dropping
> > the UPROBE_HANDLER_IWANTMYCOOKIE completely and setup cookie for any consumer
> > that has both 'handler' and 'ret_handler' defined, wdyt?
> 
> Up to you. As I said from the very beginning I won't insist on _IWANTMYCOOKIE.

ok

> 
> >  	list_for_each_entry_srcu(uc, &uprobe->consumers, cons_node,
> >  				 srcu_read_lock_held(&uprobes_srcu)) {
> > +		ric = return_consumer_find(ri, &ric_idx, uc->id);
> >  		if (uc->ret_handler)
> > -			uc->ret_handler(uc, ri->func, regs);
> > +			uc->ret_handler(uc, ri->func, regs, ric ? &ric->cookie : NULL);
> >  	}
> >  	srcu_read_unlock(&uprobes_srcu, srcu_idx);
> 
> return_consumer_find() makes no sense if !uc->ret_handler, can you move
> 
> 		ric = return_consumer_find(ri, &ric_idx, uc->id);
> 
> into the "if (uc->ret_handler)" block?

ok, will move that

thanks,
jirka

On 09/23, Jiri Olsa wrote:
>
> On Mon, Sep 23, 2024 at 12:05:53PM +0200, Oleg Nesterov wrote:
> > On 09/23, Jiri Olsa wrote:
> > >
> > > change below should do what you proposed originally
> >
> > LGTM, just one nit below.
> >
> > But I guess you need to do this on top of bpf/bpf.git, Andrii has already
> > applied your series.
>
> that seems confusing

Yes ;)

> but looks like just that one fix with the
> commit link in [1] was applied

Ah, OK.

> > And to remind, 02/14 must be fixed in any case unless I am totally confused,
> > handler_chain() can leak return_instance.
>
> yep it was missing kfree, but it's not needed in this new version

Yes, yes, the new version looks fine to me.

Oleg.