Re: [PATCH 3/8] bpf: Add bpf_cookie support to fprobe

From: Andrii Nakryiko
Date: Tue Feb 08 2022 - 18:53:38 EST


On Tue, Feb 8, 2022 at 3:46 PM Jiri Olsa <olsajiri@xxxxxxxxx> wrote:
>
> On Tue, Feb 08, 2022 at 03:35:24PM -0800, Andrii Nakryiko wrote:
> > On Tue, Feb 8, 2022 at 1:07 AM Jiri Olsa <jolsa@xxxxxxxxxx> wrote:
> > >
> > > On Mon, Feb 07, 2022 at 10:59:21AM -0800, Andrii Nakryiko wrote:
> > > > On Wed, Feb 2, 2022 at 5:54 AM Jiri Olsa <jolsa@xxxxxxxxxx> wrote:
> > > > >
> > > > > Adding support to call bpf_get_attach_cookie helper from
> > > > > kprobe program attached by fprobe link.
> > > > >
> > > > > The bpf_cookie is provided by array of u64 values, where
> > > > > each value is paired with provided function address with
> > > > > the same array index.
> > > > >
> > > > > Suggested-by: Andrii Nakryiko <andrii@xxxxxxxxxx>
> > > > > Signed-off-by: Jiri Olsa <jolsa@xxxxxxxxxx>
> > > > > ---
> > > > > include/linux/bpf.h | 2 +
> > > > > include/uapi/linux/bpf.h | 1 +
> > > > > kernel/bpf/syscall.c | 83 +++++++++++++++++++++++++++++++++-
> > > > > kernel/trace/bpf_trace.c | 16 ++++++-
> > > > > tools/include/uapi/linux/bpf.h | 1 +
> > > > > 5 files changed, 100 insertions(+), 3 deletions(-)
> > > > >
> > > > > diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> > > > > index 6eb0b180d33b..7b65f05c0487 100644
> > > > > --- a/include/linux/bpf.h
> > > > > +++ b/include/linux/bpf.h
> > > > > @@ -1301,6 +1301,8 @@ static inline void bpf_reset_run_ctx(struct bpf_run_ctx *old_ctx)
> > > > > #endif
> > > > > }
> > > > >
> > > > > +u64 bpf_fprobe_cookie(struct bpf_run_ctx *ctx, u64 ip);
> > > > > +
> > > > > /* BPF program asks to bypass CAP_NET_BIND_SERVICE in bind. */
> > > > > #define BPF_RET_BIND_NO_CAP_NET_BIND_SERVICE (1 << 0)
> > > > > /* BPF program asks to set CN on the packet. */
> > > > > diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> > > > > index c0912f0a3dfe..0dc6aa4f9683 100644
> > > > > --- a/include/uapi/linux/bpf.h
> > > > > +++ b/include/uapi/linux/bpf.h
> > > > > @@ -1484,6 +1484,7 @@ union bpf_attr {
> > > > > __aligned_u64 addrs;
> > > > > __u32 cnt;
> > > > > __u32 flags;
> > > > > + __aligned_u64 bpf_cookies;
> > > >
> > > > maybe put it right after addrs, they are closely related and cnt
> > > > describes all of syms/addrs/cookies.
> > >
> > > ok
> > >
> > > >
> > > > > } fprobe;
> > > > > };
> > > > > } link_create;
> > > > > diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
> > > > > index 0cfbb112c8e1..6c5e74bc43b6 100644
> > > > > --- a/kernel/bpf/syscall.c
> > > > > +++ b/kernel/bpf/syscall.c
> > > > > @@ -33,6 +33,8 @@
> > > > > #include <linux/rcupdate_trace.h>
> > > > > #include <linux/memcontrol.h>
> > > > > #include <linux/fprobe.h>
> > > > > +#include <linux/bsearch.h>
> > > > > +#include <linux/sort.h>
> > > > >
> > > > > #define IS_FD_ARRAY(map) ((map)->map_type == BPF_MAP_TYPE_PERF_EVENT_ARRAY || \
> > > > > (map)->map_type == BPF_MAP_TYPE_CGROUP_ARRAY || \
> > > > > @@ -3025,10 +3027,18 @@ static int bpf_perf_link_attach(const union bpf_attr *attr, struct bpf_prog *pro
> > > > >
> > > > > #ifdef CONFIG_FPROBE
> > > > >
> > > > > +struct bpf_fprobe_cookie {
> > > > > + unsigned long addr;
> > > > > + u64 bpf_cookie;
> > > > > +};
> > > > > +
> > > > > struct bpf_fprobe_link {
> > > > > struct bpf_link link;
> > > > > struct fprobe fp;
> > > > > unsigned long *addrs;
> > > > > + struct bpf_run_ctx run_ctx;
> > > > > + struct bpf_fprobe_cookie *bpf_cookies;
> > > >
> > > > you already have all the addrs above, why keeping a second copy of
> > > > each addrs in bpf_fprobe_cookie. Let's have two arrays: addrs
> > > > (unsigned long) and cookies (u64) and make sure that they are sorted
> > > > together. Then lookup addrs, calculate index, use that index to fetch
> > > > cookie.
> > > >
> > > > Seems like sort_r() provides exactly the interface you'd need to do
> > > > this very easily. Having addrs separate from cookies also a bit
> > > > advantageous in terms of TLB misses (if you need any more persuasion
> > > > ;)
> > >
> > > no persuation needed, I actually tried that but it turned out sort_r
> > > is not ready yet ;-)
> > >
> > > because you can't pass priv pointer to the swap callback, so we can't
> > > swap the other array.. I did a change to allow that, but it's not trivial
> > > and will need some bigger testing/review because the original sort
> > > calls sort_r, and of course there are many 'sort' users ;-)
> >
> > Big sigh... :( Did you do something similar to _CMP_WRAPPER? You don't
> > need to change the interface of sort(), so it shouldn't require
> > extensive code refactoring. You'll just need to adjust priv to be not
> > just cmp_func, but cmp_func + swap_fun (need a small struct on the
> > stack in sort, probably). Or you did something else?
>
> I ended up with change below
>

exactly what I had in mind

> jirka
>
>
> ---
> include/linux/sort.h | 2 +-
> include/linux/types.h | 1 +
> lib/sort.c | 44 +++++++++++++++++++++++++++++++++----------
> 3 files changed, 36 insertions(+), 11 deletions(-)
>
> diff --git a/include/linux/sort.h b/include/linux/sort.h
> index b5898725fe9d..e163287ac6c1 100644
> --- a/include/linux/sort.h
> +++ b/include/linux/sort.h
> @@ -6,7 +6,7 @@
>
> void sort_r(void *base, size_t num, size_t size,
> cmp_r_func_t cmp_func,
> - swap_func_t swap_func,
> + swap_r_func_t swap_func,
> const void *priv);
>
> void sort(void *base, size_t num, size_t size,
> diff --git a/include/linux/types.h b/include/linux/types.h
> index ac825ad90e44..ea8cf60a8a79 100644
> --- a/include/linux/types.h
> +++ b/include/linux/types.h
> @@ -226,6 +226,7 @@ struct callback_head {
> typedef void (*rcu_callback_t)(struct rcu_head *head);
> typedef void (*call_rcu_func_t)(struct rcu_head *head, rcu_callback_t func);
>
> +typedef void (*swap_r_func_t)(void *a, void *b, int size, const void *priv);
> typedef void (*swap_func_t)(void *a, void *b, int size);
>
> typedef int (*cmp_r_func_t)(const void *a, const void *b, const void *priv);
> diff --git a/lib/sort.c b/lib/sort.c
> index aa18153864d2..f65078608c16 100644
> --- a/lib/sort.c
> +++ b/lib/sort.c
> @@ -122,16 +122,29 @@ static void swap_bytes(void *a, void *b, size_t n)
> * a pointer, but small integers make for the smallest compare
> * instructions.
> */
> -#define SWAP_WORDS_64 (swap_func_t)0
> -#define SWAP_WORDS_32 (swap_func_t)1
> -#define SWAP_BYTES (swap_func_t)2
> +#define SWAP_WORDS_64 (swap_r_func_t)0
> +#define SWAP_WORDS_32 (swap_r_func_t)1
> +#define SWAP_BYTES (swap_r_func_t)2
> +#define SWAP_WRAPPER (swap_r_func_t)3
> +
> +struct wrapper {
> + cmp_func_t cmp;
> + swap_func_t swap;
> +};
>
> /*
> * The function pointer is last to make tail calls most efficient if the
> * compiler decides not to inline this function.
> */
> -static void do_swap(void *a, void *b, size_t size, swap_func_t swap_func)
> +static void do_swap(void *a, void *b, size_t size, swap_r_func_t swap_func, const void *priv)
> {
> + const struct wrapper *w = priv;

I'd just move this under if

> +
> + if (swap_func == SWAP_WRAPPER) {

const struct wrapper *w = priv; here

> + w->swap(a, b, (int)size);
> + return;
> + }
> +
> if (swap_func == SWAP_WORDS_64)
> swap_words_64(a, b, size);
> else if (swap_func == SWAP_WORDS_32)
> @@ -139,15 +152,17 @@ static void do_swap(void *a, void *b, size_t size, swap_func_t swap_func)
> else if (swap_func == SWAP_BYTES)
> swap_bytes(a, b, size);
> else
> - swap_func(a, b, (int)size);
> + swap_func(a, b, (int)size, priv);
> }
>
> #define _CMP_WRAPPER ((cmp_r_func_t)0L)
>
> static int do_cmp(const void *a, const void *b, cmp_r_func_t cmp, const void *priv)
> {
> + const struct wrapper *w = priv;
> +
> if (cmp == _CMP_WRAPPER)
> - return ((cmp_func_t)(priv))(a, b);
> + return w->cmp(a, b);

same here, or just stick to the previous style with

return ((const struct wrapper *)priv)->cmd(a, b);

> return cmp(a, b, priv);
> }
>
> @@ -198,16 +213,20 @@ static size_t parent(size_t i, unsigned int lsbit, size_t size)
> */
> void sort_r(void *base, size_t num, size_t size,
> cmp_r_func_t cmp_func,
> - swap_func_t swap_func,
> + swap_r_func_t swap_func,
> const void *priv)
> {
> /* pre-scale counters for performance */
> size_t n = num * size, a = (num/2) * size;
> const unsigned int lsbit = size & -size; /* Used to find parent */
> + const struct wrapper *w = priv;
>
> if (!a) /* num < 2 || size == 0 */
> return;
>
> + if (swap_func == SWAP_WRAPPER && !w->swap)

same here, I'd probably do the cast right here to keep this wrapper
stuff as local as possible

> + swap_func = NULL;
> +
> if (!swap_func) {
> if (is_aligned(base, size, 8))
> swap_func = SWAP_WORDS_64;
> @@ -230,7 +249,7 @@ void sort_r(void *base, size_t num, size_t size,
> if (a) /* Building heap: sift down --a */
> a -= size;
> else if (n -= size) /* Sorting: Extract root to --n */
> - do_swap(base, base + n, size, swap_func);
> + do_swap(base, base + n, size, swap_func, priv);
> else /* Sort complete */
> break;
>
> @@ -257,7 +276,7 @@ void sort_r(void *base, size_t num, size_t size,
> c = b; /* Where "a" belongs */
> while (b != a) { /* Shift it into place */
> b = parent(b, lsbit, size);
> - do_swap(base + b, base + c, size, swap_func);
> + do_swap(base + b, base + c, size, swap_func, priv);
> }
> }
> }
> @@ -267,6 +286,11 @@ void sort(void *base, size_t num, size_t size,
> cmp_func_t cmp_func,
> swap_func_t swap_func)
> {
> - return sort_r(base, num, size, _CMP_WRAPPER, swap_func, cmp_func);
> + struct wrapper w = {
> + .cmp = cmp_func,
> + .swap = swap_func,
> + };
> +
> + return sort_r(base, num, size, _CMP_WRAPPER, SWAP_WRAPPER, &w);
> }
> EXPORT_SYMBOL(sort);
> --
> 2.34.1
>