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

[Jiri Olsa <olsajiri@gmail.com>]

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

SNIP

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

yes

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

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

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

ok

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

ok, will try

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

ok

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

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

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

ok

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

yes ;-) will change

jirka

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

SNIP