[PATCH bpf v4 0/3] bpf: Fix use-after-free of progs detached from busy trampolines

[PATCH bpf v4 0/3] bpf: Fix use-after-free of progs detached from busy trampolines

[Florent Revest (Anthropic)]

A task running in a trampoline image can call a prog that was detached
and freed in the meantime, when it slept in a sleepable prog before
reaching the detached one or was preempted right before calling it.
Patch 1 handles the preempted case with an RCU tasks grace period,
patch 2 handles the sleeping case by patching detached progs out of the
images that still call them and patch 3 adds a selftest for the sleeping
case.

v3 patched all the nops of an image when it was put. bpf-ci pointed out
that this lets a task skip fmod_ret and LSM progs that are still
attached, for as long as it sleeps in an earlier prog, and Alexei asked
to go back to what v2 did: only the nop of the prog that is detached is
patched, in every image that isn't freed yet. Patch 2 explains why that
takes a list of images and not just the current one. ip_after_call is
gone in both versions, and the call to the original function is never
skipped.

sashiko noted that on riscv and loongarch a task can be preempted in
the middle of a multi-instruction patch site. ip_after_call has this too
and it isn't addressed here, pending input from the JIT maintainers on
whether a single instruction site is fine for these in-image jumps.

Tested on x86_64 under KVM and on arm64, s390x, powerpc64le, riscv64
and loongarch64 under qemu TCG, all with KASAN: the new selftest crashes
the unpatched kernel and passes with the series on every one of them,
along with trampoline_count, fentry/fexit, fentry_fexit, modify_return
and lsm_cgroup (and the full test_progs on x86_64). On x86_64, 18
fsession progs with cookies on an 11 argument function still fit in the
image. Same on bpf-next, where patch 2 has a trivial conflict in x86's
arch_bpf_trampoline_size().

Changes since v3
(https://lore.kernel.org/bpf/20260924170543.1017048-1-florent.revest@linux.dev/):
- Only patch the nop of the prog that is detached, in all the images
  that aren't freed yet, like v2 did, and explain why a list of images
  is needed (Alexei, bpf-ci)
- Lower BPF_MAX_TRAMP_LINKS to 36 on x86, the worst case no longer fit
  in a page with the nops (bpf-ci, Alexei)
- selftest: wait for the detached progs to really be freed before
  releasing the task, the grace period of patch 1 made the previous wait
  too short (bpf-ci). Detach two progs one after the other, so that the
  second detach has to reach an image that isn't the current one anymore
- Keep a single comment block in bpf_tramp_image_put() (bpf-ci)

Changes since v2
(https://lore.kernel.org/bpf/20260912095924.866254-1-florent.revest@linux.dev/):
- Patch all the nops in bpf_tramp_image_put() instead of the detached
  prog's nop at detach time, drop the image list, the trampoline
  backpointer and ip_after_call (Alexei)
- Wait for an RCU tasks grace period before freeing progs that were
  linked to a trampoline, for tasks preempted right before the enter
  helper (Junseo, sashiko)
- Initialize the jit ctx in loongarch's arch_bpf_trampoline_size() and
  the dummy image in every arch_bpf_trampoline_size() (bpf-ci)
- selftest: move the userfaultfd helper to testing_helpers.c and share
  it with bpf_mod_race, comment fixes (bpf-ci), add a subtest where the
  original function runs between the sleeping prog and the detached one

Changes since v1
(https://lore.kernel.org/bpf/20260819122252.1782790-1-florent.revest@linux.dev/):
- Patch nops in front of detached progs instead of taking prog
  references from the image (Alexei)
- Lower BPF_MAX_TRAMP_LINKS on arm64, loongarch and powerpc so the
  image still fits in a page
- Explain that the sleepable case doesn't depend on CONFIG_PREEMPTION
  (Kumar, Alexei)
- Add a selftest (Jiri, Alexei)
- Add Sechang's Reported-by (Junseo, Kumar)
- Drop Leon's and Kumar's acks since the code changed entirely

Florent Revest (Anthropic) (3):
  bpf: Wait for an RCU tasks grace period before freeing trampoline
    progs
  bpf: Skip detached progs in trampoline images that are still in use
  selftests/bpf: Detach a trampoline prog while a task sleeps before it

 arch/arm64/net/bpf_jit_comp.c                 |  33 +--
 arch/loongarch/net/bpf_jit.c                  |  45 +++--
 arch/powerpc/net/bpf_jit_comp.c               |  45 +++--
 arch/riscv/net/bpf_jit_comp64.c               |  39 ++--
 arch/s390/net/bpf_jit_comp.c                  |  46 +++--
 arch/x86/net/bpf_jit_comp.c                   |  26 ++-
 include/linux/bpf.h                           |  46 ++++-
 kernel/bpf/syscall.c                          |  19 +-
 kernel/bpf/trampoline.c                       |  75 +++++--
 .../selftests/bpf/prog_tests/bpf_mod_race.c   |  34 +---
 .../bpf/prog_tests/tramp_prog_detach.c        | 188 ++++++++++++++++++
 .../selftests/bpf/progs/tramp_prog_detach.c   |  56 ++++++
 tools/testing/selftests/bpf/testing_helpers.c |  28 +++
 tools/testing/selftests/bpf/testing_helpers.h |   2 +
 14 files changed, 526 insertions(+), 156 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/tramp_prog_detach.c
 create mode 100644 tools/testing/selftests/bpf/progs/tramp_prog_detach.c


base-commit: 5fc5768c7ca92895ccd1de94dc521e5a55ae7896
-- 
2.55.0

[PATCH bpf v4 1/3] bpf: Wait for an RCU tasks grace period before freeing trampoline progs

[Florent Revest (Anthropic)]

When a prog is detached from a trampoline, it is freed after an RCU
grace period, or an RCU tasks trace one if it is sleepable. This covers
the tasks that are running the prog, since the prog's enter helper takes
the matching read lock before the prog is called. On a preemptible
kernel, it doesn't cover a task that was preempted in the trampoline
just before the enter helper. That task holds no lock yet, and it calls
the prog after it was freed:

  BUG: KASAN: vmalloc-out-of-bounds in __bpf_prog_enter_recur+0x3a5/0x3f0
  Read of size 8 at addr ffffc90000055040 by task candidate/110
  CPU: 1 UID: 0 PID: 110 Comm: candidate Not tainted 7.3.0-rc2-00014-g15071f2a1263-dirty #2 PREEMPT(full)
  Call Trace:
   <TASK>
   __bpf_prog_enter_recur+0x3a5/0x3f0
   bpf_trampoline_6442509193+0x37/0xf1
   __x64_sys_futex+0x9/0x410
   do_syscall_64+0xb0/0x530
   ...

Wait for an RCU tasks grace period before the existing one when freeing
a prog that was linked to a trampoline. An RCU tasks grace period only
ends once the tasks that were preempted have run again, and
bpf_tramp_image_put() already relies on it to free the image. It doesn't
wait for tasks that sleep in the trampoline, the next commit takes care
of those.

Only progs that were linked to a trampoline can be called this way, so
bpf_trampoline_add_prog() marks them and other progs are still freed as
before.

Fixes: e21aa341785c ("bpf: Fix fexit trampoline.")
Reported-by: Junseo Lim <zirajs7@gmail.com>
Closes: https://lore.kernel.org/bpf/aqdrwVpanH3WGurX@omen-arch/
Assisted-by: Claude:unspecified
Signed-off-by: Florent Revest (Anthropic) <florent.revest@linux.dev>
---
 include/linux/bpf.h     |  1 +
 kernel/bpf/syscall.c    | 19 ++++++++++++++++++-
 kernel/bpf/trampoline.c |  1 +
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 1d2676782d70..d4f732996b47 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1770,6 +1770,7 @@ struct bpf_prog_aux {
 	bool offload_requested; /* Program is bound and offloaded to the netdev. */
 	bool attach_btf_trace; /* true if attaching to BTF-enabled raw tp */
 	bool attach_tracing_prog; /* true if tracing another tracing program */
+	bool tramp_linked; /* true if it was ever called from a trampoline */
 	bool func_proto_unreliable;
 	bool tail_call_reachable;
 	bool xdp_has_frags;
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 244a939b9d2d..96217b99399d 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2448,6 +2448,21 @@ static void __bpf_prog_put_rcu(struct rcu_head *rcu)
 	bpf_prog_free(aux->prog);
 }
 
+/*
+ * Progs called from a trampoline can also be reached by a task that was
+ * preempted in the trampoline before the prog's enter helper took its RCU
+ * read lock, wait for those first.
+ */
+static void __bpf_prog_put_rcu_tasks(struct rcu_head *rcu)
+{
+	struct bpf_prog *prog = container_of(rcu, struct bpf_prog_aux, rcu)->prog;
+
+	if (prog->sleepable)
+		call_rcu_tasks_trace(rcu, __bpf_prog_put_rcu);
+	else
+		call_rcu(rcu, __bpf_prog_put_rcu);
+}
+
 static void __bpf_prog_put_noref(struct bpf_prog *prog, bool deferred)
 {
 	bpf_prog_kallsyms_del_all(prog);
@@ -2461,7 +2476,9 @@ static void __bpf_prog_put_noref(struct bpf_prog *prog, bool deferred)
 		btf_put(prog->aux->attach_btf);
 
 	if (deferred) {
-		if (prog->sleepable)
+		if (IS_ENABLED(CONFIG_TASKS_RCU) && prog->aux->tramp_linked)
+			call_rcu_tasks(&prog->aux->rcu, __bpf_prog_put_rcu_tasks);
+		else if (prog->sleepable)
 			call_rcu_tasks_trace(&prog->aux->rcu, __bpf_prog_put_rcu);
 		else
 			call_rcu(&prog->aux->rcu, __bpf_prog_put_rcu);
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 90b70ea0d370..9d69c066a817 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -907,6 +907,7 @@ static int bpf_trampoline_add_prog(struct bpf_trampoline *tr,
 	}
 
 	hlist_add_head(&node->tramp_hlist, prog_list);
+	node->link->prog->aux->tramp_linked = true;
 	if (kind == BPF_TRAMP_FSESSION) {
 		tr->progs_cnt[BPF_TRAMP_FENTRY]++;
 		fexit = fsession_exit(node);
-- 
2.55.0

Re: [PATCH bpf v4 1/3] bpf: Wait for an RCU tasks grace period before freeing trampoline progs

[bot+bpf-ci]

[-- Attachment #1: Type: text/plain, Size: 1800 bytes --]

> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> index 1d2676782d70..d4f732996b47 100644
> --- a/include/linux/bpf.h
> +++ b/include/linux/bpf.h
> @@ -1770,6 +1770,7 @@ struct bpf_prog_aux {
>  	bool offload_requested; /* Program is bound and offloaded to the netdev. */
>  	bool attach_btf_trace; /* true if attaching to BTF-enabled raw tp */
>  	bool attach_tracing_prog; /* true if tracing another tracing program */
> +	bool tramp_linked; /* true if it was ever called from a trampoline */
>  	bool func_proto_unreliable;

Does this comment accurately describe when the flag is set?

Looking at bpf_trampoline_add_prog() in kernel/bpf/trampoline.c, the
flag is set at link time, right after hlist_add_head():

	hlist_add_head(&node->tramp_hlist, prog_list);
	node->link->prog->aux->tramp_linked = true;

This happens before bpf_trampoline_update() has generated an image that
could call the prog. If bpf_trampoline_update() then fails,
__bpf_trampoline_link_prog() calls bpf_trampoline_remove_prog() and the
flag stays true even though no trampoline ever called the prog.

Setting the flag conservatively is correct for the free path, since it
only adds an RCU Tasks grace period. But the comment describes a
different condition than the one the code tracks. The field name and the
commit message both say "linked to a trampoline".

Would something like /* true if it was ever linked to a trampoline */
match the code better, and stop a later reader from assuming the flag is
only set once a trampoline image has actually run the prog?

[ ... ]


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/36122863248

[PATCH bpf v4 2/3] bpf: Skip detached progs in trampoline images that are still in use

[Florent Revest (Anthropic)]

bpf_tramp_image_put() makes sure a trampoline image is not freed while
a task may still be running in it, but nothing similar is done for the
progs called by that image. Detach drops the last prog reference right
away and the prog is freed after grace periods, on the basis that a
task still in the traced function skips the fexit progs once the nop at
ip_after_call is patched to a jump.

That leaves out a task sleeping in a sleepable prog that runs before
the detached one, which no grace period waits for:

  CPU 0                               CPU 1
  in image I, sleeping in prog S
                                      detach P from I's trampoline
                                       -> new image, bpf_tramp_image_put(I)
                                      bpf_prog_put(P), last ref
                                      grace periods, P freed
  back from S
  __bpf_prog_enter(P)
  call P->bpf_func

If S and P are fexit progs the task is already past the patched jump,
and fentry only images don't have one. On x86 this is an int3 in
poisoned bpf_prog_pack memory:

  Oops: int3: 0000 [#1] SMP NOPTI
  CPU: 18 UID: 0 PID: 94573 Comm: x169 Not tainted 6.18.44 #1 PREEMPT(lazy)
  RIP: 0010:0xffffffffc0601d8d
  Call Trace:
   <TASK>
   ? bpf_trampoline_6442515411+0x1a4/0x21b
   bpf_lsm_bprm_committed_creds+0x5/0x10
   security_bprm_committed_creds+0x5f/0x70
   begin_new_exec+0x2d6/0x410
   ...

We hit this in production when progs attached through trampolines got
detached while their hooks were busy, and it was independently found
with a fuzzer and KASAN.

Have the JITs emit a patchable nop in front of each prog call sequence
and record it in the image. When a prog is detached, patch its nop to a
jump over the call sequence. Progs that stay attached keep running for
the tasks that are in the image, and ip_after_call isn't needed anymore.

The task can be in any image that isn't freed yet, not only in the
current one. It sleeps in image I1 that calls S, P and Q, then P is
detached and the trampoline moves to image I2, then Q is detached and
its call is still in I1. So the trampoline keeps a list of its images
until they are freed, and detaching a prog patches its nop in all of
them. Images hold a reference on the trampoline for that long.

With the extra nops, BPF_MAX_TRAMP_LINKS progs no longer fit in a page
on x86, arm64 and loongarch (and already didn't on powerpc), so lower
the limit there like s390 does.

Fixes: e21aa341785c ("bpf: Fix fexit trampoline.")
Reported-by: Sechang Lim <rhkrqnwk98@gmail.com>
Closes: https://lore.kernel.org/bpf/20260815071927.147049-1-zirajs7@gmail.com/
Suggested-by: Alexei Starovoitov <ast@kernel.org>
Assisted-by: Claude:unspecified
Signed-off-by: Florent Revest (Anthropic) <florent.revest@linux.dev>
---
 arch/arm64/net/bpf_jit_comp.c   | 33 ++++++++-------
 arch/loongarch/net/bpf_jit.c    | 45 ++++++++++----------
 arch/powerpc/net/bpf_jit_comp.c | 45 ++++++++++----------
 arch/riscv/net/bpf_jit_comp64.c | 39 +++++++++--------
 arch/s390/net/bpf_jit_comp.c    | 46 +++++++++++++-------
 arch/x86/net/bpf_jit_comp.c     | 26 +++++++-----
 include/linux/bpf.h             | 45 ++++++++++++++++++--
 kernel/bpf/trampoline.c         | 74 +++++++++++++++++++++++++--------
 8 files changed, 231 insertions(+), 122 deletions(-)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index c5f55d6161fe..bd73179c00ec 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -2418,10 +2418,11 @@ bool bpf_jit_supports_subprog_tailcalls(void)
 	return true;
 }
 
-static void invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *node,
-			    int bargs_off, int retval_off, int run_ctx_off,
-			    bool save_ret)
+static void invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_image *im,
+			    struct bpf_tramp_node *node, int bargs_off,
+			    int retval_off, int run_ctx_off, bool save_ret)
 {
+	void *skip;
 	__le32 *branch;
 	u64 enter_prog;
 	u64 exit_prog;
@@ -2431,6 +2432,10 @@ static void invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *node,
 	enter_prog = (u64)bpf_trampoline_enter(p);
 	exit_prog = (u64)bpf_trampoline_exit(p);
 
+	/* nop, patched to skip this prog when it is detached */
+	skip = ctx->ro_image + ctx->idx;
+	emit(A64_NOP, ctx);
+
 	if (node->cookie == 0) {
 		/* if cookie is zero, one instruction is enough to store it */
 		emit(A64_STR64I(A64_ZR, A64_SP, run_ctx_off + cookie_off), ctx);
@@ -2483,11 +2488,13 @@ static void invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *node,
 	emit(A64_ADD_I(1, A64_R(2), A64_SP, run_ctx_off), ctx);
 
 	emit_call(exit_prog, ctx);
+
+	bpf_tramp_image_add_skip(im, p, skip, ctx->ro_image + ctx->idx);
 }
 
-static void invoke_bpf_mod_ret(struct jit_ctx *ctx, struct bpf_tramp_nodes *tn,
-			       int bargs_off, int retval_off, int run_ctx_off,
-			       __le32 **branches)
+static void invoke_bpf_mod_ret(struct jit_ctx *ctx, struct bpf_tramp_image *im,
+			       struct bpf_tramp_nodes *tn, int bargs_off,
+			       int retval_off, int run_ctx_off, __le32 **branches)
 {
 	int i;
 
@@ -2496,7 +2503,7 @@ static void invoke_bpf_mod_ret(struct jit_ctx *ctx, struct bpf_tramp_nodes *tn,
 	 */
 	emit(A64_STR64I(A64_ZR, A64_SP, retval_off), ctx);
 	for (i = 0; i < tn->nr_nodes; i++) {
-		invoke_bpf_prog(ctx, tn->nodes[i], bargs_off, retval_off,
+		invoke_bpf_prog(ctx, im, tn->nodes[i], bargs_off, retval_off,
 				run_ctx_off, true);
 		/* if (*(u64 *)(sp + retval_off) !=  0)
 		 *	goto do_fexit;
@@ -2882,7 +2889,7 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
 			store_func_meta(ctx, meta, func_meta_off);
 			cookie_bargs_off--;
 		}
-		invoke_bpf_prog(ctx, fentry->nodes[i], bargs_off,
+		invoke_bpf_prog(ctx, im, fentry->nodes[i], bargs_off,
 				retval_off, run_ctx_off,
 				flags & BPF_TRAMP_F_RET_FENTRY_RET);
 	}
@@ -2893,7 +2900,7 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
 		if (!branches)
 			return -ENOMEM;
 
-		invoke_bpf_mod_ret(ctx, fmod_ret, bargs_off, retval_off,
+		invoke_bpf_mod_ret(ctx, im, fmod_ret, bargs_off, retval_off,
 				   run_ctx_off, branches);
 	}
 
@@ -2906,9 +2913,6 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
 		emit(A64_RET(A64_R(10)), ctx);
 		/* store return value */
 		emit(A64_STR64I(A64_R(0), A64_SP, retval_off), ctx);
-		/* reserve a nop for bpf_tramp_image_put */
-		im->ip_after_call = ctx->ro_image + ctx->idx;
-		emit(A64_NOP, ctx);
 	}
 
 	/* update the branches saved in invoke_bpf_mod_ret with cbnz */
@@ -2930,12 +2934,11 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
 			store_func_meta(ctx, meta, func_meta_off);
 			cookie_bargs_off--;
 		}
-		invoke_bpf_prog(ctx, fexit->nodes[i], bargs_off, retval_off,
+		invoke_bpf_prog(ctx, im, fexit->nodes[i], bargs_off, retval_off,
 				run_ctx_off, false);
 	}
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		im->ip_epilogue = ctx->ro_image + ctx->idx;
 		/* for the first pass, assume the worst case */
 		if (!ctx->image)
 			ctx->idx += 4;
@@ -2994,7 +2997,7 @@ int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 		.image = NULL,
 		.idx = 0,
 	};
-	struct bpf_tramp_image im;
+	struct bpf_tramp_image im = {};
 	struct arg_aux aaux;
 	int ret;
 
diff --git a/arch/loongarch/net/bpf_jit.c b/arch/loongarch/net/bpf_jit.c
index 4da278900938..df8abfa055d1 100644
--- a/arch/loongarch/net/bpf_jit.c
+++ b/arch/loongarch/net/bpf_jit.c
@@ -1696,13 +1696,19 @@ static void restore_stk_args(struct jit_ctx *ctx, int nr_stk_args, int args_off,
 	}
 }
 
-static int invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *n,
-			   int args_off, int retval_off, int run_ctx_off, bool save_ret)
+static int invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_image *im,
+			   struct bpf_tramp_node *n, int args_off, int retval_off,
+			   int run_ctx_off, bool save_ret)
 {
-	int ret;
+	int i, ret;
 	u32 *branch;
 	struct bpf_prog *p = n->link->prog;
 	int cookie_off = offsetof(struct bpf_tramp_run_ctx, bpf_cookie);
+	void *skip = ctx->ro_image + ctx->idx;
+
+	/* nops for move_imm+jirl, patched to skip this prog when it is detached */
+	for (i = 0; i < LOONGARCH_LONG_JUMP_NINSNS; i++)
+		emit_insn(ctx, nop);
 
 	if (n->cookie)
 		emit_store_stack_imm64(ctx, LOONGARCH_GPR_T1,
@@ -1755,13 +1761,17 @@ static int invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *n,
 	/* arg3: &run_ctx */
 	emit_insn(ctx, addid, LOONGARCH_GPR_A2, LOONGARCH_GPR_FP, -run_ctx_off);
 	ret = emit_call(ctx, (const u64)bpf_trampoline_exit(p));
+	if (ret)
+		return ret;
 
-	return ret;
+	bpf_tramp_image_add_skip(im, p, skip, ctx->ro_image + ctx->idx);
+	return 0;
 }
 
-static int invoke_bpf(struct jit_ctx *ctx, struct bpf_tramp_nodes *tn,
-		      int args_off, int retval_off, int run_ctx_off,
-		      int func_meta_off, bool save_ret, u64 func_meta, int cookie_off)
+static int invoke_bpf(struct jit_ctx *ctx, struct bpf_tramp_image *im,
+		      struct bpf_tramp_nodes *tn, int args_off, int retval_off,
+		      int run_ctx_off, int func_meta_off, bool save_ret,
+		      u64 func_meta, int cookie_off)
 {
 	int i, cur_cookie = (cookie_off - args_off) / 8;
 
@@ -1774,7 +1784,8 @@ static int invoke_bpf(struct jit_ctx *ctx, struct bpf_tramp_nodes *tn,
 			emit_store_stack_imm64(ctx, LOONGARCH_GPR_T1, -func_meta_off, meta);
 			cur_cookie--;
 		}
-		err = invoke_bpf_prog(ctx, tn->nodes[i], args_off, retval_off, run_ctx_off, save_ret);
+		err = invoke_bpf_prog(ctx, im, tn->nodes[i], args_off, retval_off,
+				      run_ctx_off, save_ret);
 		if (err)
 			return err;
 	}
@@ -2017,7 +2028,7 @@ static int __arch_prepare_bpf_trampoline(struct jit_ctx *ctx, struct bpf_tramp_i
 	}
 
 	if (fentry->nr_nodes) {
-		ret = invoke_bpf(ctx, fentry, args_off, retval_off, run_ctx_off, func_meta_off,
+		ret = invoke_bpf(ctx, im, fentry, args_off, retval_off, run_ctx_off, func_meta_off,
 				 flags & BPF_TRAMP_F_RET_FENTRY_RET, func_meta, cookie_off);
 		if (ret)
 			return ret;
@@ -2029,7 +2040,7 @@ static int __arch_prepare_bpf_trampoline(struct jit_ctx *ctx, struct bpf_tramp_i
 
 		emit_insn(ctx, std, LOONGARCH_GPR_ZERO, LOONGARCH_GPR_FP, -retval_off);
 		for (i = 0; i < fmod_ret->nr_nodes; i++) {
-			ret = invoke_bpf_prog(ctx, fmod_ret->nodes[i],
+			ret = invoke_bpf_prog(ctx, im, fmod_ret->nodes[i],
 					      args_off, retval_off, run_ctx_off, true);
 			if (ret)
 				goto out;
@@ -2051,10 +2062,6 @@ static int __arch_prepare_bpf_trampoline(struct jit_ctx *ctx, struct bpf_tramp_i
 			goto out;
 		emit_insn(ctx, std, LOONGARCH_GPR_A0, LOONGARCH_GPR_FP, -retval_off);
 		emit_insn(ctx, std, regmap[BPF_REG_0], LOONGARCH_GPR_FP, -(retval_off - 8));
-		im->ip_after_call = ctx->ro_image + ctx->idx;
-		/* Reserve space for the move_imm + jirl instruction */
-		for (i = 0; i < LOONGARCH_LONG_JUMP_NINSNS; i++)
-			emit_insn(ctx, nop);
 	}
 
 	for (i = 0; ctx->image && i < fmod_ret->nr_nodes; i++) {
@@ -2068,14 +2075,13 @@ static int __arch_prepare_bpf_trampoline(struct jit_ctx *ctx, struct bpf_tramp_i
 		emit_store_stack_imm64(ctx, LOONGARCH_GPR_T1, -func_meta_off, func_meta);
 
 	if (fexit->nr_nodes) {
-		ret = invoke_bpf(ctx, fexit, args_off, retval_off, run_ctx_off,
+		ret = invoke_bpf(ctx, im, fexit, args_off, retval_off, run_ctx_off,
 				 func_meta_off, false, func_meta, cookie_off);
 		if (ret)
 			goto out;
 	}
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		im->ip_epilogue = ctx->ro_image + ctx->idx;
 		move_addr(ctx, LOONGARCH_GPR_A0, (const u64)im);
 		ret = emit_call(ctx, (const u64)__bpf_tramp_exit);
 		if (ret)
@@ -2177,11 +2183,8 @@ int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *func_addr)
 {
 	int ret;
-	struct jit_ctx ctx;
-	struct bpf_tramp_image im;
-
-	ctx.image = NULL;
-	ctx.idx = 0;
+	struct jit_ctx ctx = {};
+	struct bpf_tramp_image im = {};
 
 	ret = __arch_prepare_bpf_trampoline(&ctx, &im, m, tnodes, func_addr, flags);
 
diff --git a/arch/powerpc/net/bpf_jit_comp.c b/arch/powerpc/net/bpf_jit_comp.c
index 7b07b43575f1..7a612688e7a2 100644
--- a/arch/powerpc/net/bpf_jit_comp.c
+++ b/arch/powerpc/net/bpf_jit_comp.c
@@ -602,14 +602,18 @@ int arch_protect_bpf_trampoline(void *image, unsigned int size)
 }
 
 static int invoke_bpf_prog(u32 *image, u32 *ro_image, struct codegen_context *ctx,
-			   struct bpf_tramp_node *n, int regs_off, int retval_off,
-			   int run_ctx_off, bool save_ret)
+			   struct bpf_tramp_image *im, struct bpf_tramp_node *n,
+			   int regs_off, int retval_off, int run_ctx_off, bool save_ret)
 {
 	struct bpf_prog *p = n->link->prog;
 	ppc_inst_t branch_insn;
-	u32 jmp_idx;
+	u32 jmp_idx, skip_idx;
 	int ret = 0;
 
+	/* nop, patched to skip this prog when it is detached */
+	skip_idx = ctx->idx;
+	EMIT(PPC_RAW_NOP());
+
 	/* Save cookie */
 	if (IS_ENABLED(CONFIG_PPC64)) {
 		PPC_LI64(_R3, n->cookie);
@@ -679,13 +683,17 @@ static int invoke_bpf_prog(u32 *image, u32 *ro_image, struct codegen_context *ct
 	EMIT(PPC_RAW_ADDI(_R5, _R1, run_ctx_off));
 	ret = bpf_jit_emit_func_call_rel(image, ro_image, ctx,
 					 (unsigned long)bpf_trampoline_exit(p));
+	if (ret)
+		return ret;
 
-	return ret;
+	if (ro_image) /* image is NULL for dummy pass */
+		bpf_tramp_image_add_skip(im, p, &ro_image[skip_idx], &ro_image[ctx->idx]);
+	return 0;
 }
 
 static int invoke_bpf_mod_ret(u32 *image, u32 *ro_image, struct codegen_context *ctx,
-			      struct bpf_tramp_nodes *tn, int regs_off, int retval_off,
-			      int run_ctx_off, u32 *branches)
+			      struct bpf_tramp_image *im, struct bpf_tramp_nodes *tn,
+			      int regs_off, int retval_off, int run_ctx_off, u32 *branches)
 {
 	int i;
 
@@ -696,8 +704,8 @@ static int invoke_bpf_mod_ret(u32 *image, u32 *ro_image, struct codegen_context
 	EMIT(PPC_RAW_LI(_R3, 0));
 	EMIT(PPC_RAW_STL(_R3, _R1, retval_off));
 	for (i = 0; i < tn->nr_nodes; i++) {
-		if (invoke_bpf_prog(image, ro_image, ctx, tn->nodes[i], regs_off, retval_off,
-				    run_ctx_off, true))
+		if (invoke_bpf_prog(image, ro_image, ctx, im, tn->nodes[i], regs_off,
+				    retval_off, run_ctx_off, true))
 			return -EINVAL;
 
 		/*
@@ -1043,8 +1051,8 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 			cookie_ctx_off--;
 		}
 
-		if (invoke_bpf_prog(image, ro_image, ctx, fentry->nodes[i], regs_off, retval_off,
-				    run_ctx_off, flags & BPF_TRAMP_F_RET_FENTRY_RET))
+		if (invoke_bpf_prog(image, ro_image, ctx, im, fentry->nodes[i], regs_off,
+				    retval_off, run_ctx_off, flags & BPF_TRAMP_F_RET_FENTRY_RET))
 			return -EINVAL;
 	}
 
@@ -1053,7 +1061,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		if (!branches)
 			return -ENOMEM;
 
-		if (invoke_bpf_mod_ret(image, ro_image, ctx, fmod_ret, regs_off, retval_off,
+		if (invoke_bpf_mod_ret(image, ro_image, ctx, im, fmod_ret, regs_off, retval_off,
 				       run_ctx_off, branches)) {
 			ret = -EINVAL;
 			goto cleanup;
@@ -1090,11 +1098,6 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		/* Restore updated tail_call_cnt */
 		if (flags & BPF_TRAMP_F_TAIL_CALL_CTX)
 			bpf_trampoline_restore_tail_call_cnt(image, ctx, bpf_frame_size, r4_off);
-
-		/* Reserve space to patch branch instruction to skip fexit progs */
-		if (ro_image) /* image is NULL for dummy pass */
-			im->ip_after_call = &((u32 *)ro_image)[ctx->idx];
-		EMIT(PPC_RAW_NOP());
 	}
 
 	/* Update branches saved in invoke_bpf_mod_ret with address of do_fexit */
@@ -1123,16 +1126,14 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 			cookie_ctx_off--;
 		}
 
-		if (invoke_bpf_prog(image, ro_image, ctx, fexit->nodes[i], regs_off, retval_off,
-				    run_ctx_off, false)) {
+		if (invoke_bpf_prog(image, ro_image, ctx, im, fexit->nodes[i], regs_off,
+				    retval_off, run_ctx_off, false)) {
 			ret = -EINVAL;
 			goto cleanup;
 		}
 	}
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		if (ro_image) /* image is NULL for dummy pass */
-			im->ip_epilogue = &((u32 *)ro_image)[ctx->idx];
 		PPC_LI_ADDR(_R3, im);
 		ret = bpf_jit_emit_func_call_rel(image, ro_image, ctx,
 						 (unsigned long)__bpf_tramp_exit);
@@ -1192,7 +1193,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *func_addr)
 {
-	struct bpf_tramp_image im;
+	struct bpf_tramp_image im = {};
 	int ret;
 
 	ret = __arch_prepare_bpf_trampoline(&im, NULL, NULL, NULL, m, flags, tnodes, func_addr);
@@ -1320,7 +1321,7 @@ int bpf_arch_text_poke(void *ip, enum bpf_text_poke_type old_t,
 
 	/*
 	 * If we are not poking at bpf prog entry, then we are simply patching in/out
-	 * an unconditional branch instruction at im->ip_after_call
+	 * an unconditional branch instruction in a trampoline image
 	 */
 	if (offset) {
 		if (old_t == BPF_MOD_CALL || new_t == BPF_MOD_CALL) {
diff --git a/arch/riscv/net/bpf_jit_comp64.c b/arch/riscv/net/bpf_jit_comp64.c
index 151031e97a24..558e02dc2fe3 100644
--- a/arch/riscv/net/bpf_jit_comp64.c
+++ b/arch/riscv/net/bpf_jit_comp64.c
@@ -904,12 +904,18 @@ static void emit_store_stack_imm64(u8 reg, int stack_off, u64 imm64,
 	emit_sd(RV_REG_FP, stack_off, reg, ctx);
 }
 
-static int invoke_bpf_prog(struct bpf_tramp_node *node, int args_off, int retval_off,
-			   int run_ctx_off, bool save_ret, struct rv_jit_context *ctx)
+static int invoke_bpf_prog(struct bpf_tramp_image *im, struct bpf_tramp_node *node,
+			   int args_off, int retval_off, int run_ctx_off, bool save_ret,
+			   struct rv_jit_context *ctx)
 {
 	int ret, branch_off;
 	struct bpf_prog *p = node->link->prog;
 	int cookie_off = offsetof(struct bpf_tramp_run_ctx, bpf_cookie);
+	void *skip = ctx->ro_insns + ctx->ninsns;
+
+	/* 2 nops for auipc+jalr, patched to skip this prog when it is detached */
+	emit(rv_nop(), ctx);
+	emit(rv_nop(), ctx);
 
 	if (node->cookie)
 		emit_store_stack_imm64(RV_REG_T1, -run_ctx_off + cookie_off, node->cookie, ctx);
@@ -962,13 +968,17 @@ static int invoke_bpf_prog(struct bpf_tramp_node *node, int args_off, int retval
 	/* arg3: &run_ctx */
 	emit_addi(RV_REG_A2, RV_REG_FP, -run_ctx_off, ctx);
 	ret = emit_call((const u64)bpf_trampoline_exit(p), true, ctx);
+	if (ret)
+		return ret;
 
-	return ret;
+	bpf_tramp_image_add_skip(im, p, skip, ctx->ro_insns + ctx->ninsns);
+	return 0;
 }
 
-static int invoke_bpf(struct bpf_tramp_nodes *tn, int args_off, int retval_off,
-		      int run_ctx_off, int func_meta_off, bool save_ret, u64 func_meta,
-		      int cookie_off, struct rv_jit_context *ctx)
+static int invoke_bpf(struct bpf_tramp_image *im, struct bpf_tramp_nodes *tn,
+		      int args_off, int retval_off, int run_ctx_off, int func_meta_off,
+		      bool save_ret, u64 func_meta, int cookie_off,
+		      struct rv_jit_context *ctx)
 {
 	int i, cur_cookie = (cookie_off - args_off) / 8;
 
@@ -981,8 +991,8 @@ static int invoke_bpf(struct bpf_tramp_nodes *tn, int args_off, int retval_off,
 			emit_store_stack_imm64(RV_REG_T1, -func_meta_off, meta, ctx);
 			cur_cookie--;
 		}
-		err = invoke_bpf_prog(tn->nodes[i], args_off, retval_off, run_ctx_off,
-				      save_ret, ctx);
+		err = invoke_bpf_prog(im, tn->nodes[i], args_off, retval_off,
+				      run_ctx_off, save_ret, ctx);
 		if (err)
 			return err;
 	}
@@ -1170,7 +1180,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 	}
 
 	if (fentry->nr_nodes) {
-		ret = invoke_bpf(fentry, args_off, retval_off, run_ctx_off, func_meta_off,
+		ret = invoke_bpf(im, fentry, args_off, retval_off, run_ctx_off, func_meta_off,
 				 flags & BPF_TRAMP_F_RET_FENTRY_RET, func_meta, cookie_off, ctx);
 		if (ret)
 			return ret;
@@ -1184,7 +1194,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 		/* cleanup to avoid garbage return value confusion */
 		emit_sd(RV_REG_FP, -retval_off, RV_REG_ZERO, ctx);
 		for (i = 0; i < fmod_ret->nr_nodes; i++) {
-			ret = invoke_bpf_prog(fmod_ret->nodes[i], args_off, retval_off,
+			ret = invoke_bpf_prog(im, fmod_ret->nodes[i], args_off, retval_off,
 					      run_ctx_off, true, ctx);
 			if (ret)
 				goto out;
@@ -1211,10 +1221,6 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 			emit_sd(RV_REG_FP, -tcc_off, RV_REG_TCC, ctx);
 		emit_sd(RV_REG_FP, -retval_off, RV_REG_A0, ctx);
 		emit_sd(RV_REG_FP, -(retval_off - 8), regmap[BPF_REG_0], ctx);
-		im->ip_after_call = ctx->ro_insns + ctx->ninsns;
-		/* 2 nops reserved for auipc+jalr pair */
-		emit(rv_nop(), ctx);
-		emit(rv_nop(), ctx);
 	}
 
 	/* update branches saved in invoke_bpf_mod_ret with bnez */
@@ -1230,14 +1236,13 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 		emit_store_stack_imm64(RV_REG_T1, -func_meta_off, func_meta, ctx);
 
 	if (fexit->nr_nodes) {
-		ret = invoke_bpf(fexit, args_off, retval_off, run_ctx_off, func_meta_off,
+		ret = invoke_bpf(im, fexit, args_off, retval_off, run_ctx_off, func_meta_off,
 				 false, func_meta, cookie_off, ctx);
 		if (ret)
 			goto out;
 	}
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		im->ip_epilogue = ctx->ro_insns + ctx->ninsns;
 		emit_imm(RV_REG_A0, ctx->insns ? (const s64)im : RV_MAX_COUNT_IMM, ctx);
 		ret = emit_call((const u64)__bpf_tramp_exit, true, ctx);
 		if (ret)
@@ -1299,7 +1304,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *func_addr)
 {
-	struct bpf_tramp_image im;
+	struct bpf_tramp_image im = {};
 	struct rv_jit_context ctx;
 	int ret;
 
diff --git a/arch/s390/net/bpf_jit_comp.c b/arch/s390/net/bpf_jit_comp.c
index c4b47070bb59..1b2566012b63 100644
--- a/arch/s390/net/bpf_jit_comp.c
+++ b/arch/s390/net/bpf_jit_comp.c
@@ -2566,6 +2566,8 @@ struct bpf_tramp_jit {
 	int r14_off;		/* Offset of saved %r14, has to be at the
 				 * bottom */
 	int do_fexit;		/* do_fexit: label */
+	int skip[BPF_MAX_TRAMP_LINKS];	/* skip: labels after each prog */
+	int nr_progs;
 };
 
 static void load_imm64(struct bpf_jit *jit, int dst_reg, u64 val)
@@ -2584,6 +2586,7 @@ static void emit_store_stack_imm64(struct bpf_jit *jit, int tmp_reg, int stack_o
 }
 
 static int invoke_bpf_prog(struct bpf_tramp_jit *tjit,
+			   struct bpf_tramp_image *im,
 			   const struct btf_func_model *m,
 			   struct bpf_tramp_node *node, bool save_ret)
 {
@@ -2591,8 +2594,20 @@ static int invoke_bpf_prog(struct bpf_tramp_jit *tjit,
 	int cookie_off = tjit->run_ctx_off +
 			 offsetof(struct bpf_tramp_run_ctx, bpf_cookie);
 	struct bpf_prog *p = node->link->prog;
+	void *skip = jit->prg_buf + jit->prg;
+	int idx = tjit->nr_progs++;
 	int patch;
 
+	if (idx >= ARRAY_SIZE(tjit->skip))
+		return -E2BIG;
+
+	/*
+	 * nop, patched to skip this prog when it is detached
+	 */
+
+	/* brcl 0,skip */
+	EMIT6_PCREL_RILC(0xc0040000, 0, tjit->skip[idx]);
+
 	/*
 	 * run_ctx.cookie = node->cookie;
 	 */
@@ -2652,10 +2667,15 @@ static int invoke_bpf_prog(struct bpf_tramp_jit *tjit,
 	/* brasl %r14,__bpf_prog_exit */
 	EMIT6_PCREL_RILB_PTR(0xc0050000, REG_14, bpf_trampoline_exit(p));
 
+	/* skip: */
+	tjit->skip[idx] = jit->prg;
+	bpf_tramp_image_add_skip(im, p, skip, jit->prg_buf + jit->prg);
+
 	return 0;
 }
 
 static int invoke_bpf(struct bpf_tramp_jit *tjit,
+		      struct bpf_tramp_image *im,
 		      const struct btf_func_model *m,
 		      struct bpf_tramp_nodes *tn, bool save_ret,
 		      u64 func_meta, int cookie_off)
@@ -2670,7 +2690,7 @@ static int invoke_bpf(struct bpf_tramp_jit *tjit,
 			emit_store_stack_imm64(jit, REG_0, tjit->func_meta_off, meta);
 			cur_cookie--;
 		}
-		if (invoke_bpf_prog(tjit, m, tn->nodes[i], save_ret))
+		if (invoke_bpf_prog(tjit, im, m, tn->nodes[i], save_ret))
 			return -EINVAL;
 	}
 
@@ -2712,6 +2732,11 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 	u64 func_meta;
 	int i, j;
 
+	/* The skip labels are taken from the previous pass. */
+	tjit->nr_progs = 0;
+	if (im)
+		im->nr_skips = 0;
+
 	/* Support as many stack arguments as "mvc" instruction can handle. */
 	nr_reg_args = min_t(int, m->nr_args, MAX_NR_REG_ARGS);
 	nr_stack_args = m->nr_args - nr_reg_args;
@@ -2875,7 +2900,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 		emit_store_stack_imm64(jit, REG_0, tjit->retval_off, 0);
 	}
 
-	if (invoke_bpf(tjit, m, fentry, flags & BPF_TRAMP_F_RET_FENTRY_RET,
+	if (invoke_bpf(tjit, im, m, fentry, flags & BPF_TRAMP_F_RET_FENTRY_RET,
 		       func_meta, cookie_off))
 		return -EINVAL;
 
@@ -2889,7 +2914,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 		       0xf000 | tjit->retval_off);
 
 		for (i = 0; i < fmod_ret->nr_nodes; i++) {
-			if (invoke_bpf_prog(tjit, m, fmod_ret->nodes[i], true))
+			if (invoke_bpf_prog(tjit, im, m, fmod_ret->nodes[i], true))
 				return -EINVAL;
 
 			/*
@@ -2943,15 +2968,6 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 		/* mvc tccnt_off(%r15),tail_call_cnt(4,%r15) */
 		_EMIT6(0xd203f000 | tjit->tccnt_off,
 		       0xf000 | offsetof(struct prog_frame, tail_call_cnt));
-
-		im->ip_after_call = jit->prg_buf + jit->prg;
-
-		/*
-		 * The following nop will be patched by bpf_tramp_image_put().
-		 */
-
-		/* brcl 0,im->ip_epilogue */
-		EMIT6_PCREL_RILC(0xc0040000, 0, (u64)im->ip_epilogue);
 	}
 
 	/* Set the "is_return" flag for fsession. */
@@ -2962,12 +2978,10 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 
 	/* do_fexit: */
 	tjit->do_fexit = jit->prg;
-	if (invoke_bpf(tjit, m, fexit, false, func_meta, cookie_off))
+	if (invoke_bpf(tjit, im, m, fexit, false, func_meta, cookie_off))
 		return -EINVAL;
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		im->ip_epilogue = jit->prg_buf + jit->prg;
-
 		/*
 		 * __bpf_tramp_exit(im);
 		 */
@@ -3016,7 +3030,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im,
 int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *orig_call)
 {
-	struct bpf_tramp_image im;
+	struct bpf_tramp_image im = {};
 	struct bpf_tramp_jit tjit;
 	int ret;
 
diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 2853e87797a7..6bca87457e87 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -3217,16 +3217,21 @@ static void restore_regs(const struct btf_func_model *m, u8 **prog,
 }
 
 static int invoke_bpf_prog(const struct btf_func_model *m, u8 **pprog,
+			   struct bpf_tramp_image *im,
 			   struct bpf_tramp_node *node, int stack_size,
 			   int run_ctx_off, bool save_ret,
 			   void *image, void *rw_image)
 {
 	u8 *prog = *pprog;
-	u8 *jmp_insn;
+	u8 *jmp_insn, *skip;
 	int ctx_cookie_off = offsetof(struct bpf_tramp_run_ctx, bpf_cookie);
 	struct bpf_prog *p = node->link->prog;
 	u64 cookie = node->cookie;
 
+	/* nop, patched to skip this prog when it is detached */
+	skip = image + (prog - (u8 *)rw_image);
+	emit_nops(&prog, X86_PATCH_SIZE);
+
 	/* mov rdi, cookie */
 	emit_mov_imm64(&prog, BPF_REG_1, (long) cookie >> 32, (u32) (long) cookie);
 
@@ -3301,6 +3306,8 @@ static int invoke_bpf_prog(const struct btf_func_model *m, u8 **pprog,
 	if (emit_rsb_call(&prog, bpf_trampoline_exit(p), image + (prog - (u8 *)rw_image)))
 		return -EINVAL;
 
+	bpf_tramp_image_add_skip(im, p, skip, image + (prog - (u8 *)rw_image));
+
 	*pprog = prog;
 	return 0;
 }
@@ -3332,6 +3339,7 @@ static int emit_cond_near_jump(u8 **pprog, void *func, void *ip, u8 jmp_cond)
 }
 
 static int invoke_bpf(const struct btf_func_model *m, u8 **pprog,
+		      struct bpf_tramp_image *im,
 		      struct bpf_tramp_nodes *tl, int stack_size,
 		      int run_ctx_off, int func_meta_off, bool save_ret,
 		      void *image, void *rw_image, u64 func_meta,
@@ -3346,7 +3354,7 @@ static int invoke_bpf(const struct btf_func_model *m, u8 **pprog,
 				func_meta | (cur_cookie << BPF_TRAMP_COOKIE_INDEX_SHIFT));
 			cur_cookie--;
 		}
-		if (invoke_bpf_prog(m, &prog, tl->nodes[i], stack_size,
+		if (invoke_bpf_prog(m, &prog, im, tl->nodes[i], stack_size,
 				    run_ctx_off, save_ret, image, rw_image))
 			return -EINVAL;
 	}
@@ -3355,6 +3363,7 @@ static int invoke_bpf(const struct btf_func_model *m, u8 **pprog,
 }
 
 static int invoke_bpf_mod_ret(const struct btf_func_model *m, u8 **pprog,
+			      struct bpf_tramp_image *im,
 			      struct bpf_tramp_nodes *tl, int stack_size,
 			      int run_ctx_off, u8 **branches,
 			      void *image, void *rw_image)
@@ -3368,7 +3377,7 @@ static int invoke_bpf_mod_ret(const struct btf_func_model *m, u8 **pprog,
 	emit_mov_imm32(&prog, false, BPF_REG_0, 0);
 	emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_0, -8);
 	for (i = 0; i < tl->nr_nodes; i++) {
-		if (invoke_bpf_prog(m, &prog, tl->nodes[i], stack_size, run_ctx_off, true,
+		if (invoke_bpf_prog(m, &prog, im, tl->nodes[i], stack_size, run_ctx_off, true,
 				    image, rw_image))
 			return -EINVAL;
 
@@ -3640,7 +3649,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	}
 
 	if (fentry->nr_nodes) {
-		if (invoke_bpf(m, &prog, fentry, regs_off, run_ctx_off, func_meta_off,
+		if (invoke_bpf(m, &prog, im, fentry, regs_off, run_ctx_off, func_meta_off,
 			       flags & BPF_TRAMP_F_RET_FENTRY_RET, image, rw_image,
 			       func_meta, cookie_off))
 			return -EINVAL;
@@ -3652,7 +3661,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		if (!branches)
 			return -ENOMEM;
 
-		if (invoke_bpf_mod_ret(m, &prog, fmod_ret, regs_off,
+		if (invoke_bpf_mod_ret(m, &prog, im, fmod_ret, regs_off,
 				       run_ctx_off, branches, image, rw_image)) {
 			ret = -EINVAL;
 			goto cleanup;
@@ -3682,8 +3691,6 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		}
 		/* remember return value in a stack for bpf prog to access */
 		emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_0, -8);
-		im->ip_after_call = image + (prog - (u8 *)rw_image);
-		emit_nops(&prog, X86_PATCH_SIZE);
 	}
 
 	if (fmod_ret->nr_nodes) {
@@ -3708,7 +3715,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		emit_store_stack_imm64(&prog, BPF_REG_0, -func_meta_off, func_meta);
 
 	if (fexit->nr_nodes) {
-		if (invoke_bpf(m, &prog, fexit, regs_off, run_ctx_off, func_meta_off,
+		if (invoke_bpf(m, &prog, im, fexit, regs_off, run_ctx_off, func_meta_off,
 			       false, image, rw_image, func_meta, cookie_off)) {
 			ret = -EINVAL;
 			goto cleanup;
@@ -3723,7 +3730,6 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	 * restored to R0.
 	 */
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		im->ip_epilogue = image + (prog - (u8 *)rw_image);
 		/* arg1: mov rdi, im */
 		emit_mov_imm64(&prog, BPF_REG_1, (long) im >> 32, (u32) (long) im);
 		if (emit_rsb_call(&prog, __bpf_tramp_exit, image + (prog - (u8 *)rw_image))) {
@@ -3811,7 +3817,7 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *func_addr)
 {
-	struct bpf_tramp_image im;
+	struct bpf_tramp_image im = {};
 	void *image;
 	int ret;
 
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index d4f732996b47..d5b507f090eb 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1258,11 +1258,17 @@ struct btf_func_model {
 #define BPF_TRAMP_F_INDIRECT		BIT(8)
 
 /* Each call __bpf_prog_enter + call bpf_func + call __bpf_prog_exit is ~50
- * bytes on x86.
+ * bytes on x86. The trampoline image has to fit in PAGE_SIZE.
  */
 enum {
-#if defined(__s390x__)
+#if defined(__s390x__) || defined(__powerpc64__)
 	BPF_MAX_TRAMP_LINKS = 27,
+#elif defined(__loongarch__)
+	BPF_MAX_TRAMP_LINKS = 33,
+#elif defined(__x86_64__)
+	BPF_MAX_TRAMP_LINKS = 36,
+#elif defined(__aarch64__)
+	BPF_MAX_TRAMP_LINKS = 37,
 #else
 	BPF_MAX_TRAMP_LINKS = 38,
 #endif
@@ -1363,19 +1369,48 @@ enum bpf_tramp_prog_type {
 	BPF_TRAMP_FSESSION,
 };
 
+/*
+ * Each prog call in a trampoline image is preceded by a nop. When the prog is
+ * detached, the nop is patched to a jump to target, right after the call, so
+ * that tasks still running in the image skip the prog.
+ */
+struct bpf_tramp_skip {
+	struct bpf_prog *prog;
+	void *nop;
+	void *target;
+};
+
 struct bpf_tramp_image {
 	void *image;
 	int size;
 	struct bpf_ksym ksym;
 	struct percpu_ref pcref;
-	void *ip_after_call;
-	void *ip_epilogue;
+	bool call_orig;
+	/* entry in tr->images, the image holds a reference on tr */
+	struct bpf_trampoline *tr;
+	struct list_head list;
+	int nr_skips;
+	struct bpf_tramp_skip *skips;
 	union {
 		struct rcu_head rcu;
 		struct work_struct work;
 	};
 };
 
+static inline void bpf_tramp_image_add_skip(struct bpf_tramp_image *im, struct bpf_prog *prog,
+					    void *nop, void *target)
+{
+	struct bpf_tramp_skip *skip;
+
+	/* struct_ops trampolines and arch_bpf_trampoline_size() have no image */
+	if (!im || !im->skips)
+		return;
+	skip = &im->skips[im->nr_skips++];
+	skip->prog = prog;
+	skip->nop = nop;
+	skip->target = target;
+}
+
 struct bpf_trampoline {
 	/* hlist for trampoline_key_table */
 	struct hlist_node hlist_key;
@@ -1402,6 +1437,8 @@ struct bpf_trampoline {
 	int progs_cnt[BPF_TRAMP_MAX];
 	/* Executable image of trampoline */
 	struct bpf_tramp_image *cur_image;
+	/* Images not freed yet, cur_image and older ones still in use */
+	struct list_head images;
 	/* Used as temporary old image storage for multi_attach */
 	struct {
 		struct bpf_tramp_image *old_image;
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 9d69c066a817..ea8e67a38919 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -401,6 +401,7 @@ static struct bpf_trampoline *bpf_trampoline_lookup(u64 key, unsigned long ip)
 	head = &trampoline_ip_table[hash_64(tr->ip, TRAMPOLINE_HASH_BITS)];
 	hlist_add_head(&tr->hlist_ip, head);
 	refcount_set(&tr->refcnt, 1);
+	INIT_LIST_HEAD(&tr->images);
 	for (i = 0; i < BPF_TRAMP_MAX; i++)
 		INIT_HLIST_HEAD(&tr->progs_hlist[i]);
 out:
@@ -565,15 +566,22 @@ static void bpf_tramp_image_free(struct bpf_tramp_image *im)
 	arch_free_bpf_trampoline(im->image, im->size);
 	bpf_jit_uncharge_modmem(im->size);
 	percpu_ref_exit(&im->pcref);
+	kfree(im->skips);
 	kfree_rcu(im, rcu);
 }
 
 static void __bpf_tramp_image_put_deferred(struct work_struct *work)
 {
 	struct bpf_tramp_image *im;
+	struct bpf_trampoline *tr;
 
 	im = container_of(work, struct bpf_tramp_image, work);
+	tr = im->tr;
+	trampoline_lock(tr);
+	list_del(&im->list);
+	trampoline_unlock(tr);
 	bpf_tramp_image_free(im);
+	bpf_trampoline_put(tr);
 }
 
 /* callback, fexit step 3 or fentry step 2 */
@@ -601,7 +609,7 @@ static void __bpf_tramp_image_put_rcu_tasks(struct rcu_head *rcu)
 	struct bpf_tramp_image *im;
 
 	im = container_of(rcu, struct bpf_tramp_image, rcu);
-	if (im->ip_after_call)
+	if (im->call_orig)
 		/* the case of fmod_ret/fexit trampoline and CONFIG_PREEMPTION=y */
 		percpu_ref_kill(&im->pcref);
 	else
@@ -621,9 +629,9 @@ static void bpf_tramp_image_put(struct bpf_tramp_image *im)
 	 *
 	 * The trampoline is unreachable before bpf_tramp_image_put().
 	 *
-	 * First, patch the trampoline to avoid calling into fexit progs.
-	 * The progs will be freed even if the original function is still
-	 * executing or sleeping.
+	 * Progs are patched out of the image when they are detached, see
+	 * bpf_trampoline_skip_prog(), so they can be freed even if a task is
+	 * still in the image.
 	 * In case of CONFIG_PREEMPT=y use call_rcu_tasks() to wait on
 	 * first few asm instructions to execute and call into
 	 * __bpf_tramp_enter->percpu_ref_get.
@@ -637,11 +645,7 @@ static void bpf_tramp_image_put(struct bpf_tramp_image *im)
 	 * percpu_ref_kill will be waiting for. Hence the first
 	 * call_rcu_tasks() is not necessary.
 	 */
-	if (im->ip_after_call) {
-		int err = bpf_arch_text_poke(im->ip_after_call, BPF_MOD_NOP,
-					     BPF_MOD_JUMP, NULL,
-					     im->ip_epilogue);
-		WARN_ON(err);
+	if (im->call_orig) {
 		if (IS_ENABLED(CONFIG_TASKS_RCU))
 			call_rcu_tasks(&im->rcu, __bpf_tramp_image_put_rcu_tasks);
 		else
@@ -658,7 +662,7 @@ static void bpf_tramp_image_put(struct bpf_tramp_image *im)
 	call_rcu_tasks_trace(&im->rcu, __bpf_tramp_image_put_rcu_tasks);
 }
 
-static struct bpf_tramp_image *bpf_tramp_image_alloc(u64 key, int size)
+static struct bpf_tramp_image *bpf_tramp_image_alloc(u64 key, int size, int nr_progs)
 {
 	struct bpf_tramp_image *im;
 	struct bpf_ksym *ksym;
@@ -669,6 +673,10 @@ static struct bpf_tramp_image *bpf_tramp_image_alloc(u64 key, int size)
 	if (!im)
 		goto out;
 
+	im->skips = kzalloc_objs(*im->skips, nr_progs);
+	if (!im->skips)
+		goto out_free_im;
+
 	err = bpf_jit_charge_modmem(size);
 	if (err)
 		goto out_free_im;
@@ -695,6 +703,7 @@ static struct bpf_tramp_image *bpf_tramp_image_alloc(u64 key, int size)
 out_uncharge:
 	bpf_jit_uncharge_modmem(size);
 out_free_im:
+	kfree(im->skips);
 	kfree(im);
 out:
 	return ERR_PTR(err);
@@ -771,11 +780,12 @@ static int bpf_trampoline_update(struct bpf_trampoline *tr, bool lock_direct_mut
 		goto out;
 	}
 
-	im = bpf_tramp_image_alloc(tr->key, size);
+	im = bpf_tramp_image_alloc(tr->key, size, total);
 	if (IS_ERR(im)) {
 		err = PTR_ERR(im);
 		goto out;
 	}
+	im->call_orig = tr->flags & BPF_TRAMP_F_CALL_ORIG;
 
 	err = arch_prepare_bpf_trampoline(im, im->image, im->image + size,
 					  &tr->func.model, tr->flags, tnodes,
@@ -806,8 +816,14 @@ static int bpf_trampoline_update(struct bpf_trampoline *tr, bool lock_direct_mut
 #endif
 
 out_free:
-	if (err)
+	if (err) {
 		bpf_tramp_image_free(im);
+	} else {
+		/* track the image until it is freed, for bpf_trampoline_skip_prog() */
+		refcount_inc(&tr->refcnt);
+		im->tr = tr;
+		list_add(&im->list, &tr->images);
+	}
 out:
 	/* If any error happens, restore previous flags */
 	if (err)
@@ -921,6 +937,32 @@ static int bpf_trampoline_add_prog(struct bpf_trampoline *tr,
 	return 0;
 }
 
+/*
+ * prog was detached and can be freed, but tasks may still be running in images
+ * that call it, sleeping in an earlier prog for example. They can be in any
+ * image that is not freed yet, not only in cur_image, so patch all of them to
+ * jump over prog.
+ */
+static void bpf_trampoline_skip_prog(struct bpf_trampoline *tr, struct bpf_prog *prog)
+{
+	struct bpf_tramp_image *im;
+	int i, err;
+
+	list_for_each_entry(im, &tr->images, list) {
+		for (i = 0; i < im->nr_skips; i++) {
+			struct bpf_tramp_skip *skip = &im->skips[i];
+
+			if (skip->prog != prog)
+				continue;
+			err = bpf_arch_text_poke(skip->nop, BPF_MOD_NOP, BPF_MOD_JUMP,
+						 NULL, skip->target);
+			WARN_ON_ONCE(err);
+			/* not a nop anymore, and prog's address can be reused */
+			skip->prog = NULL;
+		}
+	}
+}
+
 static void bpf_trampoline_remove_prog(struct bpf_trampoline *tr,
 				       struct bpf_tramp_node *node)
 {
@@ -938,6 +980,7 @@ static void bpf_trampoline_remove_prog(struct bpf_trampoline *tr,
 	}
 	hlist_del_init(&node->tramp_hlist);
 	tr->progs_cnt[kind]--;
+	bpf_trampoline_skip_prog(tr, node->link->prog);
 }
 
 static int __bpf_trampoline_link_prog(struct bpf_tramp_node *node,
@@ -1246,11 +1289,8 @@ void bpf_trampoline_put(struct bpf_trampoline *tr)
 		if (WARN_ON_ONCE(!hlist_empty(&tr->progs_hlist[i])))
 			goto out;
 
-	/* This code will be executed even when the last bpf_tramp_image
-	 * is alive. All progs are detached from the trampoline and the
-	 * trampoline image is patched with jmp into epilogue to skip
-	 * fexit progs. The fentry-only trampoline will be freed via
-	 * multiple rcu callbacks.
+	/* All progs are detached and the last image has been freed, images
+	 * hold a reference on the trampoline until then.
 	 */
 	hlist_del(&tr->hlist_key);
 	hlist_del(&tr->hlist_ip);
-- 
2.55.0

Re: [PATCH bpf v4 2/3] bpf: Skip detached progs in trampoline images that are still in use

[bot+bpf-ci]

[-- Attachment #1: Type: text/plain, Size: 6826 bytes --]

> diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
> index c5f55d6161fe..bd73179c00ec 100644
> --- a/arch/arm64/net/bpf_jit_comp.c
> +++ b/arch/arm64/net/bpf_jit_comp.c

[ ... ]

> @@ -2431,6 +2432,10 @@ static void invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *node,
>  	enter_prog = (u64)bpf_trampoline_enter(p);
>  	exit_prog = (u64)bpf_trampoline_exit(p);
>
> +	/* nop, patched to skip this prog when it is detached */
> +	skip = ctx->ro_image + ctx->idx;
> +	emit(A64_NOP, ctx);
> +

Is a context synchronization event needed on the other CPUs after the skip
NOP is patched?

A subsystem pattern flags this as potentially concerning: the new per-prog
NOP becomes a B when the prog is detached. bpf_trampoline_skip_prog() ->
bpf_arch_text_poke() -> aarch64_insn_patch_text_nosync() does:

    ret = aarch64_insn_write(tp, insn);
    if (ret == 0)
        caches_clean_inval_pou((uintptr_t)tp,
                               (uintptr_t)tp + AARCH64_INSN_SIZE);

Nothing forces an isb or exception on the other CPUs. Jump label patching
of the same NOP<->B kind does this through
arch_jump_label_transform_apply() -> kick_all_cpus_sync().

The new site is a correctness point: after the detach, the trampoline
relies on this jump so that a task still in the image does not call into
the prog, which is freed after the grace periods. A CPU that has not taken
a CSE may still execute the old NOP and call the prog.

The existing comment in bpf_arch_text_poke() justifies the missing sync
only for trampoline images, by pointing at bpf_tramp_image_put(). It does
not cover the progs called from a live image.

> diff --git a/arch/loongarch/net/bpf_jit.c b/arch/loongarch/net/bpf_jit.c
> index 4da278900938..df8abfa055d1 100644
> --- a/arch/loongarch/net/bpf_jit.c
> +++ b/arch/loongarch/net/bpf_jit.c

[ ... ]

> @@ -1696,13 +1696,19 @@ static void restore_stk_args(struct jit_ctx *ctx, int nr_stk_args, int args_off,
>  	}
>  }
>
> -static int invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_node *n,
> -			   int args_off, int retval_off, int run_ctx_off, bool save_ret)
> +static int invoke_bpf_prog(struct jit_ctx *ctx, struct bpf_tramp_image *im,
> +			   struct bpf_tramp_node *n, int args_off, int retval_off,
> +			   int run_ctx_off, bool save_ret)
>  {
> -	int ret;
> +	int i, ret;
>  	u32 *branch;
>  	struct bpf_prog *p = n->link->prog;
>  	int cookie_off = offsetof(struct bpf_tramp_run_ctx, bpf_cookie);
> +	void *skip = ctx->ro_image + ctx->idx;
> +
> +	/* nops for move_imm+jirl, patched to skip this prog when it is detached */
> +	for (i = 0; i < LOONGARCH_LONG_JUMP_NINSNS; i++)
> +		emit_insn(ctx, nop);

Can a task preempted inside the skip sled resume into half of the new jump
sequence and jump to a garbage address?

The skip site is a 5-instruction nop sled. bpf_trampoline_skip_prog() turns
it into a multi-instruction sequence with bpf_arch_text_poke(skip->nop,
BPF_MOD_NOP, BPF_MOD_JUMP, NULL, skip->target). That call goes through
emit_jump_or_nops() -> move_imm(T1, target) + 'jirl zero, t1, 0' and
writes the result with:

    ret = larch_insn_text_copy(ip, new_insns, LOONGARCH_LONG_JUMP_NBYTES);

This runs on detach, while the current image and any older images still on
tr->images are live. Every call of the traced function runs through the
sled. larch_insn_text_copy() uses stop_machine_cpuslocked(), which only
parks CPUs that are running. It cannot help a task that was preempted
between two sled instructions.

The trampoline does not disable preemption before __bpf_prog_enter*(). With
CONFIG_PREEMPT, PREEMPT_LAZY (ARCH_HAS_PREEMPT_LAZY is selected in
arch/loongarch/Kconfig) or PREEMPT_RT, waking the stopper thread preempts
whatever is running on each CPU at IPI return. That includes a task that
has executed only some of the nops.

Concrete path:

  CPU0 (task A): traced func -> trampoline -> sled nop0, nop1
                 | IPI, preempted
  CPU1: bpf_link release -> bpf_trampoline_remove_prog()
        -> bpf_trampoline_skip_prog() -> bpf_arch_text_poke()
        -> larch_insn_text_copy() writes
           lu12i.w t1 / ori t1 / lu32i.d t1 / jirl zero,t1,0 / nop
  CPU0: A resumes at sled+8 and executes 'lu32i.d t1; jirl zero, t1, 0'

Bits 31:0 of t1 still hold whatever the register contained before, for
example a value left behind by the previous prog's __bpf_prog_exit call.
The indirect jump goes to an arbitrary kernel address, causing an oops or
running unintended code.

Before this change, only the prog-entry sled and ip_after_call had this
5-instruction live poke. fentry-only trampolines had no patched site inside
the image at all. Now every detach patches one sled per image on the
hottest path. The x86, arm64 and s390 skip sites are patched with a single
atomic instruction change, so they don't have this problem.

Could the LoongArch skip site use a single instruction instead? For
example, a 'b' to the skip target, which is always within the +/-128MB
range of the same image, patched from nop to 'b'. The rest of the sled
would then be unnecessary.

> diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> index 9d69c066a817..ea8e67a38919 100644
> --- a/kernel/bpf/trampoline.c
> +++ b/kernel/bpf/trampoline.c

[ ... ]

> @@ -1246,11 +1289,8 @@ void bpf_trampoline_put(struct bpf_trampoline *tr)
>  		if (WARN_ON_ONCE(!hlist_empty(&tr->progs_hlist[i])))
>  			goto out;
>
> -	/* This code will be executed even when the last bpf_tramp_image
> -	 * is alive. All progs are detached from the trampoline and the
> -	 * trampoline image is patched with jmp into epilogue to skip
> -	 * fexit progs. The fentry-only trampoline will be freed via
> -	 * multiple rcu callbacks.
> +	/* All progs are detached and the last image has been freed, images
> +	 * hold a reference on the trampoline until then.
>  	 */

This isn't a bug, but the rewritten comment puts its text on the same line
as the opening /*:

	/* All progs are detached and the last image has been freed, images
	 * hold a reference on the trampoline until then.
	 */

The BPF subsystem's comment style asks for the opening /* of a multi-line
comment to be on its own line in kernel/bpf/, even where nearby code still
uses the old style. The patch replaces this comment's text, so the line
counts as modified. The other new multi-line comment in this file, above
bpf_trampoline_skip_prog(), already uses the preferred style. The same
comment in that style would be:

	/*
	 * All progs are detached and the last image has been freed, images
	 * hold a reference on the trampoline until then.
	 */


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/36122863248

[PATCH bpf v4 3/3] selftests/bpf: Detach a trampoline prog while a task sleeps before it

[Florent Revest (Anthropic)]

Add a test for the use-after-free fixed by the previous commit. A
sleepable prog on bpf_fentry_test1() blocks a task on a userfaultfd
page, like bpf_mod_race does, while the two progs that run after it in
the same trampoline are detached and freed one after the other, which
the test knows from their .bss maps going away. After the first detach
the task is in an image that isn't the trampoline's current one
anymore. The task is then released and must not call into the freed
progs. This is done with fentry progs, with fexit progs, where the task
is already past the original function, and with a task sleeping in a
fentry prog while fexit progs are detached, where the original function
must still be called. bpf_mod_race's userfaultfd helper moves to
testing_helpers.c so that both tests can use it.

Without the fix, on a kernel with KASAN:

  BUG: KASAN: vmalloc-out-of-bounds in __bpf_prog_enter_recur+0xed/0x1e0
  Read of size 8 at addr ffa0000000144040 by task test_progs/171
  CPU: 6 UID: 0 PID: 171 Comm: test_progs Tainted: G           OE       7.2.0+ #1 PREEMPT(full)
  Call Trace:
   <TASK>
   __bpf_prog_enter_recur+0xed/0x1e0
   bpf_trampoline_6442545468+0x72/0xe3
   bpf_fentry_test1+0x9/0x20
   bpf_prog_test_run_tracing+0x183/0x3e0
   __sys_bpf+0xd3f/0x38f0
   ...

Assisted-by: Claude:unspecified
Signed-off-by: Florent Revest (Anthropic) <florent.revest@linux.dev>
---
 .../selftests/bpf/prog_tests/bpf_mod_race.c   |  34 +---
 .../bpf/prog_tests/tramp_prog_detach.c        | 188 ++++++++++++++++++
 .../selftests/bpf/progs/tramp_prog_detach.c   |  56 ++++++
 tools/testing/selftests/bpf/testing_helpers.c |  28 +++
 tools/testing/selftests/bpf/testing_helpers.h |   2 +
 5 files changed, 275 insertions(+), 33 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/tramp_prog_detach.c
 create mode 100644 tools/testing/selftests/bpf/progs/tramp_prog_detach.c

diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_mod_race.c b/tools/testing/selftests/bpf/prog_tests/bpf_mod_race.c
index ecc3d47919ad..f8497e764beb 100644
--- a/tools/testing/selftests/bpf/prog_tests/bpf_mod_race.c
+++ b/tools/testing/selftests/bpf/prog_tests/bpf_mod_race.c
@@ -55,38 +55,6 @@ static void *load_module_thread(void *p)
 	return p;
 }
 
-static int sys_userfaultfd(int flags)
-{
-	return syscall(__NR_userfaultfd, flags);
-}
-
-static int test_setup_uffd(void *fault_addr)
-{
-	struct uffdio_register uffd_register = {};
-	struct uffdio_api uffd_api = {};
-	int uffd;
-
-	uffd = sys_userfaultfd(O_CLOEXEC);
-	if (uffd < 0)
-		return -errno;
-
-	uffd_api.api = UFFD_API;
-	uffd_api.features = 0;
-	if (ioctl(uffd, UFFDIO_API, &uffd_api)) {
-		close(uffd);
-		return -1;
-	}
-
-	uffd_register.range.start = (unsigned long)fault_addr;
-	uffd_register.range.len = getpagesize();
-	uffd_register.mode = UFFDIO_REGISTER_MODE_MISSING;
-	if (ioctl(uffd, UFFDIO_REGISTER, &uffd_register)) {
-		close(uffd);
-		return -1;
-	}
-	return uffd;
-}
-
 static void test_bpf_mod_race_config(const struct test_config *config)
 {
 	void *fault_addr, *skel_fail;
@@ -117,7 +85,7 @@ static void test_bpf_mod_race_config(const struct test_config *config)
 	if (!ASSERT_OK(bpf_mod_race__attach(skel), "bpf_mod_kfunc_race__attach"))
 		goto end_destroy;
 
-	uffd = test_setup_uffd(fault_addr);
+	uffd = uffd_block_page(fault_addr);
 	if (!ASSERT_GE(uffd, 0, "userfaultfd open + register address"))
 		goto end_destroy;
 
diff --git a/tools/testing/selftests/bpf/prog_tests/tramp_prog_detach.c b/tools/testing/selftests/bpf/prog_tests/tramp_prog_detach.c
new file mode 100644
index 000000000000..1eb0d7237605
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/tramp_prog_detach.c
@@ -0,0 +1,188 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <test_progs.h>
+#include <pthread.h>
+#include <poll.h>
+#include <sys/mman.h>
+#include <linux/userfaultfd.h>
+#include "tramp_prog_detach.skel.h"
+#include "testing_helpers.h"
+
+/*
+ * Detach and free progs while a task sleeps in the prog that runs before them
+ * in the same trampoline image, then let that task continue through the
+ * image. It must not call into the freed progs.
+ *
+ * The task is held in a sleepable prog with userfaultfd, like bpf_mod_race
+ * does.
+ */
+
+static struct bpf_program *pick_prog(struct tramp_prog_detach *skel,
+				     bool fexit, bool sleepable)
+{
+	if (fexit)
+		return sleepable ? skel->progs.fexit_sleepable :
+				   skel->progs.fexit_victim;
+	return sleepable ? skel->progs.fentry_sleepable :
+			   skel->progs.fentry_victim;
+}
+
+static struct bpf_program *sleepable_prog;
+
+static void *run_sleepable(void *arg)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, topts);
+
+	/* calls bpf_fentry_test1() */
+	return (void *)(long)bpf_prog_test_run_opts(bpf_program__fd(sleepable_prog),
+						    &topts);
+}
+
+static struct tramp_prog_detach *load_one(bool fexit, bool sleepable)
+{
+	struct tramp_prog_detach *skel;
+	int err;
+
+	skel = tramp_prog_detach__open();
+	if (!ASSERT_OK_PTR(skel, "open"))
+		return NULL;
+
+	bpf_program__set_autoload(pick_prog(skel, fexit, sleepable), true);
+	err = tramp_prog_detach__load(skel);
+	if (!ASSERT_OK(err, "load"))
+		goto err;
+	skel->bss->pid = getpid();
+	err = tramp_prog_detach__attach(skel);
+	if (!ASSERT_OK(err, "attach"))
+		goto err;
+	return skel;
+err:
+	tramp_prog_detach__destroy(skel);
+	return NULL;
+}
+
+/* The .bss map of a destroyed skeleton goes away when its prog is freed */
+static bool wait_for_map_free(__u32 map_id)
+{
+	int i, fd;
+
+	for (i = 0; i < 100; i++) {
+		fd = bpf_map_get_fd_by_id(map_id);
+		if (fd < 0)
+			return true;
+		close(fd);
+		usleep(100 * 1000);
+	}
+	return false;
+}
+
+static void test_detach(bool sleepable_fexit, bool victim_fexit)
+{
+	struct tramp_prog_detach *sleepable = NULL, *victims[2] = {};
+	struct pollfd pfd = { .events = POLLIN };
+	struct uffdio_copy uffd_copy = {};
+	struct bpf_map_info map_info = {};
+	__u32 map_info_len = sizeof(map_info);
+	struct uffd_msg uffd_msg;
+	void *fault_page, *src_page = MAP_FAILED;
+	long page_size = getpagesize();
+	bool started = false;
+	void *thread_ret;
+	pthread_t thread;
+	int i, uffd = -1;
+
+	fault_page = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
+			  MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (!ASSERT_NEQ(fault_page, MAP_FAILED, "mmap fault_page"))
+		return;
+	src_page = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (!ASSERT_NEQ(src_page, MAP_FAILED, "mmap src_page"))
+		goto out;
+
+	/* The most recently attached prog runs first */
+	for (i = 0; i < ARRAY_SIZE(victims); i++) {
+		victims[i] = load_one(victim_fexit, false);
+		if (!victims[i])
+			goto out;
+	}
+	sleepable = load_one(sleepable_fexit, true);
+	if (!sleepable)
+		goto out;
+	sleepable_prog = pick_prog(sleepable, sleepable_fexit, true);
+
+	/* Not armed yet so this doesn't block, make sure sleepable runs first */
+	if (!ASSERT_OK((long)run_sleepable(NULL), "dry run"))
+		goto out;
+	for (i = 0; i < ARRAY_SIZE(victims); i++)
+		if (!ASSERT_LT(sleepable->bss->ts, victims[i]->bss->ts, "prog order"))
+			goto out;
+
+	uffd = uffd_block_page(fault_page);
+	if (!ASSERT_GE(uffd, 0, "userfaultfd open + register address"))
+		goto out;
+	sleepable->bss->fault_addr = fault_page;
+
+	if (!ASSERT_OK(pthread_create(&thread, NULL, run_sleepable, NULL),
+		       "pthread_create"))
+		goto out;
+	started = true;
+
+	/* Wait for the thread to sleep in bpf_copy_from_user() */
+	pfd.fd = uffd;
+	if (!ASSERT_EQ(poll(&pfd, 1, 10000), 1, "poll uffd"))
+		goto out;
+	if (!ASSERT_EQ(read(uffd, &uffd_msg, sizeof(uffd_msg)), sizeof(uffd_msg),
+		       "read uffd"))
+		goto out;
+	if (!ASSERT_EQ(uffd_msg.event, UFFD_EVENT_PAGEFAULT, "uffd pagefault"))
+		goto out;
+
+	/*
+	 * Detach and unload the victim progs and wait for them to be freed.
+	 * After the first one, the task is in an image that isn't the
+	 * trampoline's current one anymore.
+	 */
+	for (i = 0; i < ARRAY_SIZE(victims); i++) {
+		if (!ASSERT_OK(bpf_map_get_info_by_fd(bpf_map__fd(victims[i]->maps.bss),
+						      &map_info, &map_info_len),
+			       "victim bss info"))
+			goto out;
+		tramp_prog_detach__destroy(victims[i]);
+		victims[i] = NULL;
+		if (!ASSERT_TRUE(wait_for_map_free(map_info.id), "victim freed"))
+			goto out;
+	}
+
+out:
+	/* Let the thread proceed with the rest of the trampoline */
+	if (uffd >= 0) {
+		uffd_copy.dst = (unsigned long)fault_page;
+		uffd_copy.src = (unsigned long)src_page;
+		uffd_copy.len = page_size;
+		ASSERT_OK(ioctl(uffd, UFFDIO_COPY, &uffd_copy), "uffd copy");
+		close(uffd);
+	}
+	if (started &&
+	    ASSERT_OK(pthread_join(thread, &thread_ret), "pthread_join"))
+		ASSERT_NULL(thread_ret, "blocking run");
+
+	for (i = 0; i < ARRAY_SIZE(victims); i++)
+		tramp_prog_detach__destroy(victims[i]);
+	tramp_prog_detach__destroy(sleepable);
+	if (src_page != MAP_FAILED)
+		munmap(src_page, page_size);
+	munmap(fault_page, page_size);
+}
+
+void serial_test_tramp_prog_detach(void)
+{
+	/* a task sleeping before the original function is called */
+	if (test__start_subtest("fentry"))
+		test_detach(false, false);
+	/* a task sleeping after the original function returned */
+	if (test__start_subtest("fexit"))
+		test_detach(true, true);
+	/* the original function runs in between and must still be called */
+	if (test__start_subtest("fentry_fexit"))
+		test_detach(false, true);
+}
diff --git a/tools/testing/selftests/bpf/progs/tramp_prog_detach.c b/tools/testing/selftests/bpf/progs/tramp_prog_detach.c
new file mode 100644
index 000000000000..507d372167ec
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/tramp_prog_detach.c
@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+char _license[] SEC("license") = "GPL";
+
+int pid;
+void *fault_addr;
+__u64 ts;
+
+static int do_sleepable(void)
+{
+	char dst;
+
+	if (bpf_get_current_pid_tgid() >> 32 != pid)
+		return 0;
+
+	ts = bpf_ktime_get_ns();
+	/* blocks for as long as user space wants when fault_addr is armed */
+	bpf_copy_from_user(&dst, sizeof(dst), fault_addr);
+	return 0;
+}
+
+static int do_victim(void)
+{
+	if (bpf_get_current_pid_tgid() >> 32 != pid)
+		return 0;
+
+	ts = bpf_ktime_get_ns();
+	return 0;
+}
+
+SEC("?fentry.s/bpf_fentry_test1")
+int BPF_PROG(fentry_sleepable, int a)
+{
+	return do_sleepable();
+}
+
+SEC("?fentry/bpf_fentry_test1")
+int BPF_PROG(fentry_victim, int a)
+{
+	return do_victim();
+}
+
+SEC("?fexit.s/bpf_fentry_test1")
+int BPF_PROG(fexit_sleepable, int a, int ret)
+{
+	return do_sleepable();
+}
+
+SEC("?fexit/bpf_fentry_test1")
+int BPF_PROG(fexit_victim, int a, int ret)
+{
+	return do_victim();
+}
diff --git a/tools/testing/selftests/bpf/testing_helpers.c b/tools/testing/selftests/bpf/testing_helpers.c
index c970e7793dfc..9f672d0c9dda 100644
--- a/tools/testing/selftests/bpf/testing_helpers.c
+++ b/tools/testing/selftests/bpf/testing_helpers.c
@@ -13,6 +13,7 @@
 #include "test_progs.h"
 #include "testing_helpers.h"
 #include <linux/membarrier.h>
+#include <linux/userfaultfd.h>
 
 int parse_num_list(const char *s, bool **num_set, int *num_set_len)
 {
@@ -459,6 +460,33 @@ int kern_sync_rcu(void)
 	return syscall(__NR_membarrier, MEMBARRIER_CMD_SHARED, 0, 0);
 }
 
+int uffd_block_page(void *fault_addr)
+{
+	struct uffdio_register uffd_register = {};
+	struct uffdio_api uffd_api = {};
+	int uffd;
+
+	uffd = syscall(__NR_userfaultfd, O_CLOEXEC);
+	if (uffd < 0)
+		return -errno;
+
+	uffd_api.api = UFFD_API;
+	uffd_api.features = 0;
+	if (ioctl(uffd, UFFDIO_API, &uffd_api)) {
+		close(uffd);
+		return -1;
+	}
+
+	uffd_register.range.start = (unsigned long)fault_addr;
+	uffd_register.range.len = getpagesize();
+	uffd_register.mode = UFFDIO_REGISTER_MODE_MISSING;
+	if (ioctl(uffd, UFFDIO_REGISTER, &uffd_register)) {
+		close(uffd);
+		return -1;
+	}
+	return uffd;
+}
+
 int get_xlated_program(int fd_prog, struct bpf_insn **buf, __u32 *cnt)
 {
 	__u32 buf_element_size = sizeof(struct bpf_insn);
diff --git a/tools/testing/selftests/bpf/testing_helpers.h b/tools/testing/selftests/bpf/testing_helpers.h
index 2edc6fb7fc52..1f83b6080716 100644
--- a/tools/testing/selftests/bpf/testing_helpers.h
+++ b/tools/testing/selftests/bpf/testing_helpers.h
@@ -36,6 +36,8 @@ __u64 read_perf_max_sample_freq(void);
 int load_bpf_testmod(bool verbose);
 int unload_bpf_testmod(bool verbose);
 int kern_sync_rcu(void);
+/* returns a userfaultfd that makes accesses to fault_addr's page block */
+int uffd_block_page(void *fault_addr);
 int finit_module(int fd, const char *param_values, int flags);
 int delete_module(const char *name, int flags);
 int load_module(const char *path, bool verbose);
-- 
2.55.0