On Wed, Feb 28, 2024 at 1:03 AM Jiri Olsa <jolsa@xxxxxxxxxx> wrote:
>
> Adding bpf_kprobe_multi_is_return kfunc that returns true if the
> bpf program is executed from the exit probe of the kprobe multi
> link attached in wrapper mode. It returns false otherwise.
>
> Adding new kprobe hook for kprobe program type.
>
> Signed-off-by: Jiri Olsa <jolsa@xxxxxxxxxx>
> ---
> kernel/bpf/btf.c | 3 +++
> kernel/trace/bpf_trace.c | 49 +++++++++++++++++++++++++++++++++++++---
> 2 files changed, 49 insertions(+), 3 deletions(-)
>
> diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
> index 6ff0bd1a91d5..5ab55720e881 100644
> --- a/kernel/bpf/btf.c
> +++ b/kernel/bpf/btf.c
> @@ -218,6 +218,7 @@ enum btf_kfunc_hook {
> BTF_KFUNC_HOOK_SOCKET_FILTER,
> BTF_KFUNC_HOOK_LWT,
> BTF_KFUNC_HOOK_NETFILTER,
> + BTF_KFUNC_HOOK_KPROBE,
> BTF_KFUNC_HOOK_MAX,
> };
>
> @@ -8112,6 +8113,8 @@ static int bpf_prog_type_to_kfunc_hook(enum bpf_prog_type prog_type)
> return BTF_KFUNC_HOOK_LWT;
> case BPF_PROG_TYPE_NETFILTER:
> return BTF_KFUNC_HOOK_NETFILTER;
> + case BPF_PROG_TYPE_KPROBE:
> + return BTF_KFUNC_HOOK_KPROBE;
> default:
> return BTF_KFUNC_HOOK_MAX;
> }
> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index 726a8c71f0da..cb801c94b8fa 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -2594,6 +2594,7 @@ struct bpf_kprobe_multi_run_ctx {
> struct bpf_run_ctx run_ctx;
> struct bpf_kprobe_multi_link *link;
> unsigned long entry_ip;
> + bool is_return;
> };
>
> struct user_syms {
> @@ -2793,11 +2794,13 @@ static u64 bpf_kprobe_multi_entry_ip(struct bpf_run_ctx *ctx)
>
> static int
> kprobe_multi_link_prog_run(struct bpf_kprobe_multi_link *link,
> - unsigned long entry_ip, struct pt_regs *regs)
> + unsigned long entry_ip, struct pt_regs *regs,
> + bool is_return)
> {
> struct bpf_kprobe_multi_run_ctx run_ctx = {
> .link = link,
> .entry_ip = entry_ip,
> + .is_return = is_return,
> };
> struct bpf_run_ctx *old_run_ctx;
> int err;
> @@ -2830,7 +2833,7 @@ kprobe_multi_link_handler(struct fprobe *fp, unsigned long fentry_ip,
> int err;
>
> link = container_of(fp, struct bpf_kprobe_multi_link, fp);
> - err = kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
> + err = kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs, false);
> return link->is_wrapper ? err : 0;
> }
>
> @@ -2842,7 +2845,7 @@ kprobe_multi_link_exit_handler(struct fprobe *fp, unsigned long fentry_ip,
> struct bpf_kprobe_multi_link *link;
>
> link = container_of(fp, struct bpf_kprobe_multi_link, fp);
> - kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
> + kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs, true);
> }
>
> static int symbols_cmp_r(const void *a, const void *b, const void *priv)
> @@ -3111,6 +3114,46 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
> kvfree(cookies);
> return err;
> }
> +
> +__bpf_kfunc_start_defs();
> +
> +__bpf_kfunc bool bpf_kprobe_multi_is_return(void)
and for uprobes we'll have bpf_uprobe_multi_is_return?...
BTW, have you tried implementing a "session cookie" idea?
> +{
> + struct bpf_kprobe_multi_run_ctx *run_ctx;
> +
> + run_ctx = container_of(current->bpf_ctx, struct bpf_kprobe_multi_run_ctx, run_ctx);
> + return run_ctx->is_return;
> +}
> +
> +__bpf_kfunc_end_defs();
> +
> +BTF_KFUNCS_START(kprobe_multi_kfunc_set_ids)
> +BTF_ID_FLAGS(func, bpf_kprobe_multi_is_return)
> +BTF_KFUNCS_END(kprobe_multi_kfunc_set_ids)
> +
> +static int bpf_kprobe_multi_filter(const struct bpf_prog *prog, u32 kfunc_id)
> +{
> + if (!btf_id_set8_contains(&kprobe_multi_kfunc_set_ids, kfunc_id))
> + return 0;
> +
> + if (prog->expected_attach_type != BPF_TRACE_KPROBE_MULTI)
> + return -EACCES;
> +
> + return 0;
> +}
> +
> +static const struct btf_kfunc_id_set bpf_kprobe_multi_kfunc_set = {
> + .owner = THIS_MODULE,
> + .set = &kprobe_multi_kfunc_set_ids,
> + .filter = bpf_kprobe_multi_filter,
> +};
> +
> +static int __init bpf_kprobe_multi_kfuncs_init(void)
> +{
> + return register_btf_kfunc_id_set(BPF_PROG_TYPE_KPROBE, &bpf_kprobe_multi_kfunc_set);
> +}
> +
> +late_initcall(bpf_kprobe_multi_kfuncs_init);
> #else /* !CONFIG_FPROBE */
> int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *prog)
> {
> --
> 2.43.2
>
