[bpf-next,2/5] bpf: tracing: support to attach program to multi hooks

Commit Message

梦龙董 Feb. 20, 2024, 3:51 a.m. UTC

In this commit, we add the support to allow attaching a tracing BPF
program to multi hooks.

In the commit 4a1e7c0c63e0 ("bpf: Support attaching freplace programs to
multiple attach points"), the freplace BPF program is made to support
attach to multiple attach points. And in this commit, we extend it to
fentry/fexit/raw_tp/...

The use case is obvious. For now, we have to create a BPF program for each
kernel function, for which we want to trace, even through all the program
have the same (or similar logic). This can consume extra memory, and make
the program loading slow if we have plenty of kernel function to trace.
The KPROBE_MULTI maybe a alternative, but it can't what TRACING do. For
example, the kretprobe can't obtain the function args, but the FEXIT can.

Now, we need to hold the reference for the target btf and kernel module
in the bpf link, as a program can have multiple target. Therefore, we
introduce the attach_btf and mod field to the struct bpf_tracing_link.
During the attach, we will check the target is compatible with the
program, which means that the function args that the program accessed
in the target function prototype should be the same as the origin target.

Signed-off-by: Menglong Dong <[e-mail]>
---
 include/linux/bpf.h      |   2 +
 include/uapi/linux/bpf.h |   1 +
 kernel/bpf/syscall.c     | 117 +++++++++++++++++++++++++++++++--------
 3 files changed, 98 insertions(+), 22 deletions(-)

Comments

Jiri Olsa Feb. 20, 2024, 5:18 p.m. UTC | #1

On Tue, Feb 20, 2024 at 11:51:02AM +0800, Menglong Dong wrote:

SNIP

> @@ -3228,7 +3260,9 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  	struct bpf_link_primer link_primer;
>  	struct bpf_prog *tgt_prog = NULL;
>  	struct bpf_trampoline *tr = NULL;
> +	struct btf *attach_btf = NULL;
>  	struct bpf_tracing_link *link;
> +	struct module *mod = NULL;
>  	u64 key = 0;
>  	int err;
>  
> @@ -3258,31 +3292,50 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  		goto out_put_prog;
>  	}
>  
> -	if (!!tgt_prog_fd != !!btf_id) {
> -		err = -EINVAL;
> -		goto out_put_prog;
> -	}
> -
>  	if (tgt_prog_fd) {
> -		/*
> -		 * For now we only allow new targets for BPF_PROG_TYPE_EXT. If this
> -		 * part would be changed to implement the same for
> -		 * BPF_PROG_TYPE_TRACING, do not forget to update the way how
> -		 * attach_tracing_prog flag is set.
> -		 */
> -		if (prog->type != BPF_PROG_TYPE_EXT) {
> +		if (!btf_id) {
>  			err = -EINVAL;
>  			goto out_put_prog;
>  		}
> -
>  		tgt_prog = bpf_prog_get(tgt_prog_fd);
>  		if (IS_ERR(tgt_prog)) {
> -			err = PTR_ERR(tgt_prog);
>  			tgt_prog = NULL;
> -			goto out_put_prog;
> +			/* tgt_prog_fd is the fd of the kernel module BTF */
> +			attach_btf = btf_get_by_fd(tgt_prog_fd);

I think we should pass the btf_fd through attr, like add
link_create.tracing_btf_fd instead, this seems confusing

> +			if (IS_ERR(attach_btf)) {
> +				attach_btf = NULL;
> +				err = -EINVAL;
> +				goto out_put_prog;
> +			}
> +			if (!btf_is_kernel(attach_btf)) {
> +				btf_put(attach_btf);
> +				err = -EOPNOTSUPP;
> +				goto out_put_prog;
> +			}
> +		} else if (prog->type == BPF_PROG_TYPE_TRACING &&
> +			   tgt_prog->type == BPF_PROG_TYPE_TRACING) {
> +			prog->aux->attach_tracing_prog = true;
>  		}

could you please add comment on why this check is in here?

> -
> -		key = bpf_trampoline_compute_key(tgt_prog, NULL, btf_id);
> +		key = bpf_trampoline_compute_key(tgt_prog, attach_btf,
> +						 btf_id);
> +	} else if (btf_id) {
> +		attach_btf = bpf_get_btf_vmlinux();
> +		if (IS_ERR(attach_btf)) {
> +			attach_btf = NULL;
> +			err = PTR_ERR(attach_btf);
> +			goto out_unlock;
> +		}
> +		if (!attach_btf) {
> +			err = -EINVAL;
> +			goto out_unlock;
> +		}
> +		btf_get(attach_btf);
> +		key = bpf_trampoline_compute_key(NULL, attach_btf, btf_id);
> +	} else {
> +		attach_btf = prog->aux->attach_btf;
> +		/* get the reference of the btf for bpf link */
> +		if (attach_btf)
> +			btf_get(attach_btf);
>  	}
>  
>  	link = kzalloc(sizeof(*link), GFP_USER);
> @@ -3319,7 +3372,7 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  	 *   are NULL, then program was already attached and user did not provide
>  	 *   tgt_prog_fd so we have no way to find out or create trampoline
>  	 */
> -	if (!prog->aux->dst_trampoline && !tgt_prog) {
> +	if (!prog->aux->dst_trampoline && !tgt_prog && !btf_id) {
>  		/*
>  		 * Allow re-attach for TRACING and LSM programs. If it's
>  		 * currently linked, bpf_trampoline_link_prog will fail.
> @@ -3346,17 +3399,27 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  		 * different from the destination specified at load time, we
>  		 * need a new trampoline and a check for compatibility
>  		 */
> +		struct btf *origin_btf = prog->aux->attach_btf;
>  		struct bpf_attach_target_info tgt_info = {};
>  
> +		/* use the new attach_btf to check the target */
> +		prog->aux->attach_btf = attach_btf;
>  		err = bpf_check_attach_target(NULL, prog, tgt_prog, btf_id,
>  					      &tgt_info);
> +		prog->aux->attach_btf = origin_btf;

could we pass the attach_btf as argument then?

jirka

>  		if (err)
>  			goto out_unlock;
>  
> -		if (tgt_info.tgt_mod) {
> -			module_put(prog->aux->mod);
> -			prog->aux->mod = tgt_info.tgt_mod;
> -		}
> +		mod = tgt_info.tgt_mod;
> +		/* the new target and the previous target are in the same
> +		 * module, release the reference once.
> +		 */
> +		if (mod && mod == prog->aux->mod)
> +			module_put(mod);
> +		err = bpf_tracing_check_multi(prog, tgt_prog, attach_btf,
> +					      tgt_info.tgt_type);
> +		if (err)
> +			goto out_unlock;
>  
>  		tr = bpf_trampoline_get(key, &tgt_info);
>  		if (!tr) {
> @@ -3373,6 +3436,7 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  		 */
>  		tr = prog->aux->dst_trampoline;
>  		tgt_prog = prog->aux->dst_prog;
> +		mod = prog->aux->mod;
>  	}
>  
>  	err = bpf_link_prime(&link->link.link, &link_primer);
> @@ -3388,6 +3452,8 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  
>  	link->tgt_prog = tgt_prog;
>  	link->trampoline = tr;
> +	link->attach_btf = attach_btf;
> +	link->mod = mod;
>  
>  	/* Always clear the trampoline and target prog from prog->aux to make
>  	 * sure the original attach destination is not kept alive after a
> @@ -3400,20 +3466,27 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
>  	if (prog->aux->dst_trampoline && tr != prog->aux->dst_trampoline)
>  		/* we allocated a new trampoline, so free the old one */
>  		bpf_trampoline_put(prog->aux->dst_trampoline);
> +	if (prog->aux->mod && mod != prog->aux->mod)
> +		/* the mod in prog is not used anywhere, move it to link */
> +		module_put(prog->aux->mod);
>  
>  	prog->aux->dst_prog = NULL;
>  	prog->aux->dst_trampoline = NULL;
> +	prog->aux->mod = NULL;
>  	mutex_unlock(&prog->aux->dst_mutex);
>  
>  	return bpf_link_settle(&link_primer);
>  out_unlock:
>  	if (tr && tr != prog->aux->dst_trampoline)
>  		bpf_trampoline_put(tr);
> +	if (mod && mod != prog->aux->mod)
> +		module_put(mod);
>  	mutex_unlock(&prog->aux->dst_mutex);
>  	kfree(link);
>  out_put_prog:
>  	if (tgt_prog_fd && tgt_prog)
>  		bpf_prog_put(tgt_prog);
> +	btf_put(attach_btf);
>  	return err;
>  }
>  
> -- 
> 2.39.2
>

Dan Carpenter Feb. 21, 2024, 5:31 a.m. UTC | #2

Hi Menglong,

kernel test robot noticed the following build warnings:

url:    https://github.com/intel-lab-lkp/linux/commits/Menglong-Dong/bpf-tracing-add-support-to-record-and-check-the-accessed-args/20240220-115317
base:   https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git master
patch link:    https://lore.kernel.org/r/20240220035105.34626-3-dongmenglong.8%40bytedance.com
patch subject: [PATCH bpf-next 2/5] bpf: tracing: support to attach program to multi hooks
config: m68k-randconfig-r071-20240220 (https://download.01.org/0day-ci/archive/20240221/[e-mail]/config)
compiler: m68k-linux-gcc (GCC) 13.2.0

If you fix the issue in a separate patch/commit (i.e. not just a new version of
the same patch/commit), kindly add following tags
| Reported-by: kernel test robot <[e-mail]>
| Reported-by: Dan Carpenter <[e-mail]>
| Closes: https://lore.kernel.org/r/[e-mail]/

smatch warnings:
kernel/bpf/syscall.c:3325 bpf_tracing_prog_attach() warn: passing zero to 'PTR_ERR'
kernel/bpf/syscall.c:3485 bpf_tracing_prog_attach() error: uninitialized symbol 'link'.

vim +/PTR_ERR +3325 kernel/bpf/syscall.c

4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3255  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3256  				   int tgt_prog_fd,
2fcc82411e74e5 Kui-Feng Lee           2022-05-10  3257  				   u32 btf_id,
2fcc82411e74e5 Kui-Feng Lee           2022-05-10  3258  				   u64 bpf_cookie)
fec56f5890d93f Alexei Starovoitov     2019-11-14  3259  {
a3b80e1078943d Andrii Nakryiko        2020-04-28  3260  	struct bpf_link_primer link_primer;
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3261  	struct bpf_prog *tgt_prog = NULL;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3262  	struct bpf_trampoline *tr = NULL;
5f80eb32851d7a Menglong Dong          2024-02-20  3263  	struct btf *attach_btf = NULL;
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3264  	struct bpf_tracing_link *link;
5f80eb32851d7a Menglong Dong          2024-02-20  3265  	struct module *mod = NULL;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3266  	u64 key = 0;
a3b80e1078943d Andrii Nakryiko        2020-04-28  3267  	int err;
fec56f5890d93f Alexei Starovoitov     2019-11-14  3268  
9e4e01dfd3254c KP Singh               2020-03-29  3269  	switch (prog->type) {
9e4e01dfd3254c KP Singh               2020-03-29  3270  	case BPF_PROG_TYPE_TRACING:
fec56f5890d93f Alexei Starovoitov     2019-11-14  3271  		if (prog->expected_attach_type != BPF_TRACE_FENTRY &&
be8704ff07d237 Alexei Starovoitov     2020-01-20  3272  		    prog->expected_attach_type != BPF_TRACE_FEXIT &&
9e4e01dfd3254c KP Singh               2020-03-29  3273  		    prog->expected_attach_type != BPF_MODIFY_RETURN) {
9e4e01dfd3254c KP Singh               2020-03-29  3274  			err = -EINVAL;
9e4e01dfd3254c KP Singh               2020-03-29  3275  			goto out_put_prog;
9e4e01dfd3254c KP Singh               2020-03-29  3276  		}
9e4e01dfd3254c KP Singh               2020-03-29  3277  		break;
9e4e01dfd3254c KP Singh               2020-03-29  3278  	case BPF_PROG_TYPE_EXT:
9e4e01dfd3254c KP Singh               2020-03-29  3279  		if (prog->expected_attach_type != 0) {
9e4e01dfd3254c KP Singh               2020-03-29  3280  			err = -EINVAL;
9e4e01dfd3254c KP Singh               2020-03-29  3281  			goto out_put_prog;
9e4e01dfd3254c KP Singh               2020-03-29  3282  		}
9e4e01dfd3254c KP Singh               2020-03-29  3283  		break;
9e4e01dfd3254c KP Singh               2020-03-29  3284  	case BPF_PROG_TYPE_LSM:
9e4e01dfd3254c KP Singh               2020-03-29  3285  		if (prog->expected_attach_type != BPF_LSM_MAC) {
9e4e01dfd3254c KP Singh               2020-03-29  3286  			err = -EINVAL;
9e4e01dfd3254c KP Singh               2020-03-29  3287  			goto out_put_prog;
9e4e01dfd3254c KP Singh               2020-03-29  3288  		}
9e4e01dfd3254c KP Singh               2020-03-29  3289  		break;
9e4e01dfd3254c KP Singh               2020-03-29  3290  	default:
fec56f5890d93f Alexei Starovoitov     2019-11-14  3291  		err = -EINVAL;
fec56f5890d93f Alexei Starovoitov     2019-11-14  3292  		goto out_put_prog;
fec56f5890d93f Alexei Starovoitov     2019-11-14  3293  	}
fec56f5890d93f Alexei Starovoitov     2019-11-14  3294  
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3295  	if (tgt_prog_fd) {
5f80eb32851d7a Menglong Dong          2024-02-20  3296  		if (!btf_id) {
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3297  			err = -EINVAL;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3298  			goto out_put_prog;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3299  		}
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3300  		tgt_prog = bpf_prog_get(tgt_prog_fd);
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3301  		if (IS_ERR(tgt_prog)) {
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3302  			tgt_prog = NULL;
5f80eb32851d7a Menglong Dong          2024-02-20  3303  			/* tgt_prog_fd is the fd of the kernel module BTF */
5f80eb32851d7a Menglong Dong          2024-02-20  3304  			attach_btf = btf_get_by_fd(tgt_prog_fd);
5f80eb32851d7a Menglong Dong          2024-02-20  3305  			if (IS_ERR(attach_btf)) {
5f80eb32851d7a Menglong Dong          2024-02-20  3306  				attach_btf = NULL;
5f80eb32851d7a Menglong Dong          2024-02-20  3307  				err = -EINVAL;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3308  				goto out_put_prog;
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3309  			}
5f80eb32851d7a Menglong Dong          2024-02-20  3310  			if (!btf_is_kernel(attach_btf)) {
5f80eb32851d7a Menglong Dong          2024-02-20  3311  				btf_put(attach_btf);
5f80eb32851d7a Menglong Dong          2024-02-20  3312  				err = -EOPNOTSUPP;
5f80eb32851d7a Menglong Dong          2024-02-20  3313  				goto out_put_prog;
5f80eb32851d7a Menglong Dong          2024-02-20  3314  			}
5f80eb32851d7a Menglong Dong          2024-02-20  3315  		} else if (prog->type == BPF_PROG_TYPE_TRACING &&
5f80eb32851d7a Menglong Dong          2024-02-20  3316  			   tgt_prog->type == BPF_PROG_TYPE_TRACING) {
5f80eb32851d7a Menglong Dong          2024-02-20  3317  			prog->aux->attach_tracing_prog = true;
5f80eb32851d7a Menglong Dong          2024-02-20  3318  		}
5f80eb32851d7a Menglong Dong          2024-02-20  3319  		key = bpf_trampoline_compute_key(tgt_prog, attach_btf,
5f80eb32851d7a Menglong Dong          2024-02-20  3320  						 btf_id);
5f80eb32851d7a Menglong Dong          2024-02-20  3321  	} else if (btf_id) {
5f80eb32851d7a Menglong Dong          2024-02-20  3322  		attach_btf = bpf_get_btf_vmlinux();
5f80eb32851d7a Menglong Dong          2024-02-20  3323  		if (IS_ERR(attach_btf)) {
5f80eb32851d7a Menglong Dong          2024-02-20  3324  			attach_btf = NULL;
                                                                                ^^^^^^^^^^^^^^^^^^
This needs to be done after the "err = " assignment on the next line.

5f80eb32851d7a Menglong Dong          2024-02-20 @3325  			err = PTR_ERR(attach_btf);
                                                                                ^^^^^^^^^^^^^^^^^^^^^^^^^^
Here.

5f80eb32851d7a Menglong Dong          2024-02-20  3326  			goto out_unlock;
5f80eb32851d7a Menglong Dong          2024-02-20  3327  		}
5f80eb32851d7a Menglong Dong          2024-02-20  3328  		if (!attach_btf) {
5f80eb32851d7a Menglong Dong          2024-02-20  3329  			err = -EINVAL;
5f80eb32851d7a Menglong Dong          2024-02-20  3330  			goto out_unlock;

"link" is not initialized on this goto path so it leads to an
uninitialized variable warning.

5f80eb32851d7a Menglong Dong          2024-02-20  3331  		}
5f80eb32851d7a Menglong Dong          2024-02-20  3332  		btf_get(attach_btf);
5f80eb32851d7a Menglong Dong          2024-02-20  3333  		key = bpf_trampoline_compute_key(NULL, attach_btf, btf_id);
5f80eb32851d7a Menglong Dong          2024-02-20  3334  	} else {
5f80eb32851d7a Menglong Dong          2024-02-20  3335  		attach_btf = prog->aux->attach_btf;
5f80eb32851d7a Menglong Dong          2024-02-20  3336  		/* get the reference of the btf for bpf link */
5f80eb32851d7a Menglong Dong          2024-02-20  3337  		if (attach_btf)
5f80eb32851d7a Menglong Dong          2024-02-20  3338  			btf_get(attach_btf);
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3339  	}
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3340  
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3341  	link = kzalloc(sizeof(*link), GFP_USER);
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3342  	if (!link) {
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3343  		err = -ENOMEM;
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3344  		goto out_put_prog;
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3345  	}
f7e0beaf39d386 Kui-Feng Lee           2022-05-10  3346  	bpf_link_init(&link->link.link, BPF_LINK_TYPE_TRACING,
f2e10bff16a0fd Andrii Nakryiko        2020-04-28  3347  		      &bpf_tracing_link_lops, prog);
f2e10bff16a0fd Andrii Nakryiko        2020-04-28  3348  	link->attach_type = prog->expected_attach_type;
2fcc82411e74e5 Kui-Feng Lee           2022-05-10  3349  	link->link.cookie = bpf_cookie;
70ed506c3bbcfa Andrii Nakryiko        2020-03-02  3350  

[ snip ]

3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3474  	prog->aux->dst_trampoline = NULL;
5f80eb32851d7a Menglong Dong          2024-02-20  3475  	prog->aux->mod = NULL;
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3476  	mutex_unlock(&prog->aux->dst_mutex);
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3477  
a3b80e1078943d Andrii Nakryiko        2020-04-28  3478  	return bpf_link_settle(&link_primer);
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3479  out_unlock:
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3480  	if (tr && tr != prog->aux->dst_trampoline)
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3481  		bpf_trampoline_put(tr);
5f80eb32851d7a Menglong Dong          2024-02-20  3482  	if (mod && mod != prog->aux->mod)
5f80eb32851d7a Menglong Dong          2024-02-20  3483  		module_put(mod);
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29  3484  	mutex_unlock(&prog->aux->dst_mutex);
3aac1ead5eb6b7 Toke Høiland-Jørgensen 2020-09-29 @3485  	kfree(link);
                                                                      ^^^^

fec56f5890d93f Alexei Starovoitov     2019-11-14  3486  out_put_prog:
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3487  	if (tgt_prog_fd && tgt_prog)
4a1e7c0c63e02d Toke Høiland-Jørgensen 2020-09-29  3488  		bpf_prog_put(tgt_prog);
5f80eb32851d7a Menglong Dong          2024-02-20  3489  	btf_put(attach_btf);
fec56f5890d93f Alexei Starovoitov     2019-11-14  3490  	return err;
fec56f5890d93f Alexei Starovoitov     2019-11-14  3491  }

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 0225b8dbdd9d..cf8f2df9afb9 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1606,6 +1606,8 @@  struct bpf_tracing_link {
 	enum bpf_attach_type attach_type;
 	struct bpf_trampoline *trampoline;
 	struct bpf_prog *tgt_prog;
+	struct btf *attach_btf;
+	struct module *mod;
 };
 
 struct bpf_link_primer {
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index d96708380e52..0ded10a85bfe 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1668,6 +1668,7 @@  union bpf_attr {
 		union {
 			__u32	target_fd;	/* target object to attach to or ... */
 			__u32	target_ifindex; /* target ifindex */
+			__u32	target_btf_obj_fd;
 		};
 		__u32		attach_type;	/* attach type */
 		__u32		flags;		/* extra flags */
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index b2750b79ac80..3b432fcd5bdb 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -3178,6 +3178,9 @@  static void bpf_tracing_link_dealloc(struct bpf_link *link)
 	struct bpf_tracing_link *tr_link =
 		container_of(link, struct bpf_tracing_link, link.link);
 
+
+	btf_put(tr_link->attach_btf);
+	module_put(tr_link->mod);
 	kfree(tr_link);
 }
 
@@ -3220,6 +3223,35 @@  static const struct bpf_link_ops bpf_tracing_link_lops = {
 	.fill_link_info = bpf_tracing_link_fill_link_info,
 };
 
+static int bpf_tracing_check_multi(struct bpf_prog *prog,
+				   struct bpf_prog *tgt_prog,
+				   struct btf *btf2,
+				   const struct btf_type *t2)
+{
+	struct btf *btf1 = prog->aux->attach_btf;
+	const struct btf_type *t1;
+
+	/* this case is already valided in bpf_check_attach_target() */
+	if (prog->type == BPF_PROG_TYPE_EXT)
+		return 0;
+
+	/* For now, noly support multi attach for kernel function attach
+	 * point.
+	 */
+	if (!btf1)
+		return -EOPNOTSUPP;
+
+	btf2 = btf2 ?: tgt_prog->aux->btf;
+	t1 = prog->aux->attach_func_proto;
+
+	/* the target is the same as the origin one, this is a re-attach */
+	if (t1 == t2)
+		return 0;
+
+	return btf_check_func_part_match(btf1, t1, btf2, t2,
+					 prog->aux->accessed_args);
+}
+
 static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 				   int tgt_prog_fd,
 				   u32 btf_id,
@@ -3228,7 +3260,9 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 	struct bpf_link_primer link_primer;
 	struct bpf_prog *tgt_prog = NULL;
 	struct bpf_trampoline *tr = NULL;
+	struct btf *attach_btf = NULL;
 	struct bpf_tracing_link *link;
+	struct module *mod = NULL;
 	u64 key = 0;
 	int err;
 
@@ -3258,31 +3292,50 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 		goto out_put_prog;
 	}
 
-	if (!!tgt_prog_fd != !!btf_id) {
-		err = -EINVAL;
-		goto out_put_prog;
-	}
-
 	if (tgt_prog_fd) {
-		/*
-		 * For now we only allow new targets for BPF_PROG_TYPE_EXT. If this
-		 * part would be changed to implement the same for
-		 * BPF_PROG_TYPE_TRACING, do not forget to update the way how
-		 * attach_tracing_prog flag is set.
-		 */
-		if (prog->type != BPF_PROG_TYPE_EXT) {
+		if (!btf_id) {
 			err = -EINVAL;
 			goto out_put_prog;
 		}
-
 		tgt_prog = bpf_prog_get(tgt_prog_fd);
 		if (IS_ERR(tgt_prog)) {
-			err = PTR_ERR(tgt_prog);
 			tgt_prog = NULL;
-			goto out_put_prog;
+			/* tgt_prog_fd is the fd of the kernel module BTF */
+			attach_btf = btf_get_by_fd(tgt_prog_fd);
+			if (IS_ERR(attach_btf)) {
+				attach_btf = NULL;
+				err = -EINVAL;
+				goto out_put_prog;
+			}
+			if (!btf_is_kernel(attach_btf)) {
+				btf_put(attach_btf);
+				err = -EOPNOTSUPP;
+				goto out_put_prog;
+			}
+		} else if (prog->type == BPF_PROG_TYPE_TRACING &&
+			   tgt_prog->type == BPF_PROG_TYPE_TRACING) {
+			prog->aux->attach_tracing_prog = true;
 		}
-
-		key = bpf_trampoline_compute_key(tgt_prog, NULL, btf_id);
+		key = bpf_trampoline_compute_key(tgt_prog, attach_btf,
+						 btf_id);
+	} else if (btf_id) {
+		attach_btf = bpf_get_btf_vmlinux();
+		if (IS_ERR(attach_btf)) {
+			attach_btf = NULL;
+			err = PTR_ERR(attach_btf);
+			goto out_unlock;
+		}
+		if (!attach_btf) {
+			err = -EINVAL;
+			goto out_unlock;
+		}
+		btf_get(attach_btf);
+		key = bpf_trampoline_compute_key(NULL, attach_btf, btf_id);
+	} else {
+		attach_btf = prog->aux->attach_btf;
+		/* get the reference of the btf for bpf link */
+		if (attach_btf)
+			btf_get(attach_btf);
 	}
 
 	link = kzalloc(sizeof(*link), GFP_USER);
@@ -3319,7 +3372,7 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 	 *   are NULL, then program was already attached and user did not provide
 	 *   tgt_prog_fd so we have no way to find out or create trampoline
 	 */
-	if (!prog->aux->dst_trampoline && !tgt_prog) {
+	if (!prog->aux->dst_trampoline && !tgt_prog && !btf_id) {
 		/*
 		 * Allow re-attach for TRACING and LSM programs. If it's
 		 * currently linked, bpf_trampoline_link_prog will fail.
@@ -3346,17 +3399,27 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 		 * different from the destination specified at load time, we
 		 * need a new trampoline and a check for compatibility
 		 */
+		struct btf *origin_btf = prog->aux->attach_btf;
 		struct bpf_attach_target_info tgt_info = {};
 
+		/* use the new attach_btf to check the target */
+		prog->aux->attach_btf = attach_btf;
 		err = bpf_check_attach_target(NULL, prog, tgt_prog, btf_id,
 					      &tgt_info);
+		prog->aux->attach_btf = origin_btf;
 		if (err)
 			goto out_unlock;
 
-		if (tgt_info.tgt_mod) {
-			module_put(prog->aux->mod);
-			prog->aux->mod = tgt_info.tgt_mod;
-		}
+		mod = tgt_info.tgt_mod;
+		/* the new target and the previous target are in the same
+		 * module, release the reference once.
+		 */
+		if (mod && mod == prog->aux->mod)
+			module_put(mod);
+		err = bpf_tracing_check_multi(prog, tgt_prog, attach_btf,
+					      tgt_info.tgt_type);
+		if (err)
+			goto out_unlock;
 
 		tr = bpf_trampoline_get(key, &tgt_info);
 		if (!tr) {
@@ -3373,6 +3436,7 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 		 */
 		tr = prog->aux->dst_trampoline;
 		tgt_prog = prog->aux->dst_prog;
+		mod = prog->aux->mod;
 	}
 
 	err = bpf_link_prime(&link->link.link, &link_primer);
@@ -3388,6 +3452,8 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 
 	link->tgt_prog = tgt_prog;
 	link->trampoline = tr;
+	link->attach_btf = attach_btf;
+	link->mod = mod;
 
 	/* Always clear the trampoline and target prog from prog->aux to make
 	 * sure the original attach destination is not kept alive after a
@@ -3400,20 +3466,27 @@  static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 	if (prog->aux->dst_trampoline && tr != prog->aux->dst_trampoline)
 		/* we allocated a new trampoline, so free the old one */
 		bpf_trampoline_put(prog->aux->dst_trampoline);
+	if (prog->aux->mod && mod != prog->aux->mod)
+		/* the mod in prog is not used anywhere, move it to link */
+		module_put(prog->aux->mod);
 
 	prog->aux->dst_prog = NULL;
 	prog->aux->dst_trampoline = NULL;
+	prog->aux->mod = NULL;
 	mutex_unlock(&prog->aux->dst_mutex);
 
 	return bpf_link_settle(&link_primer);
 out_unlock:
 	if (tr && tr != prog->aux->dst_trampoline)
 		bpf_trampoline_put(tr);
+	if (mod && mod != prog->aux->mod)
+		module_put(mod);
 	mutex_unlock(&prog->aux->dst_mutex);
 	kfree(link);
 out_put_prog:
 	if (tgt_prog_fd && tgt_prog)
 		bpf_prog_put(tgt_prog);
+	btf_put(attach_btf);
 	return err;
 }