[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@xxxxxxxxxx>
---
 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