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@xxxxxxxxxx> 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@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);
> +

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
>