[PATCH v3 0/5] kprobes: Fixes for ftrace ops and optprobe state machine

[PATCH v3 0/5] kprobes: Fixes for ftrace ops and optprobe state machine

[Masami Hiramatsu (Google)]

Hi,

Here is v3 of the remaining bug fixes for kprobes and optprobes.
The previous version is here:

  https://lore.kernel.org/all/179427358690.291326.2078291446658598883.stgit@devnote2/

This version restores __arm_kprobe() fix as patch 3, and update the KUnit
test code to check it.
This includes following fixes.

1. Ftrace Ops Switching for Post-Handlers (Patch 1):
   - When adding a second probe with a post_handler to an existing ftrace
     kprobe, the aggregate probe's post_handler was set but remained armed
     on kprobe_ftrace_ops. Subsequent disarming attempted to disarm from
     kprobe_ipmodify_ops, failing with -ENOENT and leaking the armed probe.
     Switch ftrace ops safely in register_aggr_kprobe() outside of
     text_mutex and cpus_read_lock.

2. Optprobe State Machine & Tasks RCU Quiescence (Patches 2-4):
   - Forcibly unoptimized probes (unoptimize_kprobe(force = true)) were
     not queued to cooling_list, causing kprobe_disarmed() to return true
     prematurely and risking use-after-free in the optinsn detour buffer
     before Tasks RCU completed.
   - When re-arming a probe that is queued on unoptimizing_list before the
     optimizer thread runs, skip arch_arm_kprobe() in __arm_kprobe() so
     the active jump instruction is not overwritten with a breakpoint.
   - When a disabled cooling probe finishes cooling and transitions to
     safely disarmed in optprobe_finalize_generation(), re-optimize any
     blocked sibling probe whose jump range covers its address.

3. KUnit Regression Tests (Patch 5):
   - Add KUnit test cases to lib/tests/test_kprobes.c covering the recent
     kprobes, kretprobes, ftrace ops, and optprobe fixes.

Changes in v3:
[03/05]
 - Restore the __arm_kprobe() fix that was accidentally dropped in v2
   while omitting the reuse_unused_kprobe() hunk superseded by commit
   1ffd6efb1b58 ("kprobes: Arm reused optprobe on cooling list").
[05/05]
 - Check that the jump opcode at *opt_kp.addr is preserved across
   disable_kprobe() and enable_kprobe() in
   test_optprobe_rearm_unoptimizing().

Thank you,

---
base-commit: 1ffd6efb1b58e596ed70f660182803264a4c84d9

Masami Hiramatsu (Google) (5):
      kprobes: Switch ftrace ops when adding probe with post_handler
      kprobes: Fix UAF and queue cooling probes on forced unoptimization
      kprobes: Do not overwrite jump instruction when re-arming unoptimizing probe
      kprobes: Re-optimize blocked sibling probe when cooling finishes
      kprobes: Add KUnit test cases for kprobes regressions


 kernel/kprobes.c         |  116 +++++++++++++++++++++++++++-----------
 lib/tests/test_kprobes.c |  139 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 220 insertions(+), 35 deletions(-)

--
Masami Hiramatsu (Google) <mhiramat@kernel.org>

[PATCH v3 1/5] kprobes: Switch ftrace ops when adding probe with post_handler

[Masami Hiramatsu (Google)]

From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

When an ftrace-based kprobe is registered without a post_handler, it is
armed using kprobe_ftrace_ops. If a second probe with a post_handler is
subsequently registered at the same address, register_aggr_kprobe() sets
ap->post_handler to aggr_post_handler, but leaves the probe armed using
kprobe_ftrace_ops.

When the probe is later disarmed or disabled, disarm_kprobe_ftrace()
checks whether p->post_handler is non-NULL and attempts to disarm from
kprobe_ipmodify_ops. This fails with -ENOENT, decrements
kprobe_ipmodify_enabled below zero, leaves the probe armed on
kprobe_ftrace_ops, and aborts unregistration, leading to a use-after-free
if the probe module is unloaded.

Switch the armed probe from kprobe_ftrace_ops to kprobe_ipmodify_ops in
register_aggr_kprobe() outside of text_mutex and cpus_read_lock to avoid
lockdep and deadlock issues when modifying ftrace code. Defer adding the
new probe onto ap->list until the ftrace ops switch completes so that the
new probe's handlers are not published prematurely while ap is still on
kprobe_ftrace_ops, and mark ap disabled if rollback re-arming fails.
Also restore ftrace ops registration and counter if ftrace_set_filter_ip()
fails in __disarm_kprobe_ftrace().

Fixes: 5dd7caf0bdc5 ("kprobes: Skip clearing aggrprobe's post_handler in kprobe-on-ftrace case")
Cc: stable@vger.kernel.org
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
Changes in v2:
 - Defer add_new_kprobe() until after ftrace_switch_ops completes so
   that p is not visible on ap->list if switching ftrace ops fails.
 - Set KPROBE_FLAG_DISABLED on ap if rollback arm_kprobe_ftrace() fails.
---
 kernel/kprobes.c |   55 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 10 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index dddad1e01866..b6ff971df451 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1444,9 +1444,14 @@ static int __disarm_kprobe_ftrace(struct kprobe *p, struct ftrace_ops *ops,
 	(*cnt)--;
 
 	ret = ftrace_set_filter_ip(ops, (unsigned long)p->addr, 1, 0);
-	WARN_ONCE(ret < 0, "Failed to disarm kprobe-ftrace at %pS (error %d)\n",
-		  p->addr, ret);
-	return ret;
+	if (WARN_ONCE(ret < 0, "Failed to disarm kprobe-ftrace at %pS (error %d)\n",
+		      p->addr, ret)) {
+		if (*cnt == 0)
+			register_ftrace_function(ops);
+		(*cnt)++;
+		return ret;
+	}
+	return 0;
 }
 
 static int disarm_kprobe_ftrace(struct kprobe *p)
@@ -1575,17 +1580,13 @@ void kprobe_busy_end(void)
 }
 
 /* Add the new probe to 'ap->list'. */
-static int add_new_kprobe(struct kprobe *ap, struct kprobe *p)
+static void add_new_kprobe(struct kprobe *ap, struct kprobe *p)
 {
 	/* A reused cooling probe is already unoptimized */
 	if (p->post_handler && !kprobe_unused(ap))
 		unoptimize_kprobe(ap, true);	/* Fall back to normal kprobe */
 
 	list_add_rcu(&p->list, &ap->list);
-	if (p->post_handler && !ap->post_handler)
-		ap->post_handler = aggr_post_handler;
-
-	return 0;
 }
 
 /*
@@ -1618,6 +1619,8 @@ static int register_aggr_kprobe(struct kprobe *orig_p, struct kprobe *p)
 {
 	int ret = 0;
 	struct kprobe *ap = orig_p;
+	bool reset_post_handler = false;
+	bool ftrace_switch_ops = false;
 
 	scoped_guard(cpus_read_lock) {
 		/* For preparing optimization, jump_label_text_reserved() is called */
@@ -1664,18 +1667,50 @@ static int register_aggr_kprobe(struct kprobe *orig_p, struct kprobe *p)
 					| KPROBE_FLAG_DISABLED;
 		}
 
+		if (p->post_handler && !ap->post_handler) {
+			reset_post_handler = true;
+			if (kprobe_ftrace(ap) && !kprobes_all_disarmed &&
+			    !kprobe_disabled(ap))
+				ftrace_switch_ops = true;
+			else
+				ap->post_handler = aggr_post_handler;
+		}
+
 		/* Copy the insn slot of 'p' to 'ap'. */
 		copy_kprobe(ap, p);
-		ret = add_new_kprobe(ap, p);
+		if (!ftrace_switch_ops)
+			add_new_kprobe(ap, p);
+	}
+
+	if (ftrace_switch_ops) {
+		/*
+		 * Must switch ftrace ops outside of text_mutex and cpus_read_lock
+		 * to prevent deadlocks when ftrace code modification acquires text_mutex.
+		 */
+		ret = disarm_kprobe_ftrace(ap);
+		if (ret)
+			return ret;
+
+		ap->post_handler = aggr_post_handler;
+		ret = arm_kprobe_ftrace(ap);
+		if (ret) {
+			ap->post_handler = NULL;
+			if (arm_kprobe_ftrace(ap))
+				ap->flags |= KPROBE_FLAG_DISABLED;
+			return ret;
+		}
+		add_new_kprobe(ap, p);
 	}
 
-	if (ret == 0 && kprobe_disabled(ap) && !kprobe_disabled(p)) {
+	if (kprobe_disabled(ap) && !kprobe_disabled(p)) {
 		ap->flags &= ~KPROBE_FLAG_DISABLED;
 		if (!kprobes_all_disarmed) {
 			/* Arm the breakpoint again. */
 			ret = arm_kprobe(ap);
 			if (ret) {
 				ap->flags |= KPROBE_FLAG_DISABLED;
+				if (reset_post_handler)
+					ap->post_handler = NULL;
 				list_del_rcu(&p->list);
 				synchronize_rcu();
 			}

[PATCH v3 2/5] kprobes: Fix UAF and queue cooling probes on forced unoptimization

[Masami Hiramatsu (Google)]

From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

When unoptimize_kprobe() is called with force = true on an actively
optimized probe (whose op->list is empty), force_unoptimize_kprobe()
replaces the relative jump with a breakpoint, but leaves op->list empty
without queueing op to cur->cooling_list.

Because op->list remains empty, kprobe_disarmed() returns true
prematurely as soon as the probe is disabled. When unregister_kprobes()
runs, __unregister_kprobe_top() unhashes the probe immediately, and
unregister_kprobes() frees op and the optinsn slot after synchronize_rcu()
without waiting for Tasks RCU. Any task preempted inside the optinsn slot
then encounters a use-after-free when resuming execution.

Queue forcibly unoptimized probes onto cur->cooling_list and include
gen->cooling_list in optprobe_has_queued_probes() so that the generation
is dispatched and flushed properly. In do_unoptimize_kprobes(), collect
newly unoptimized probes into a separate local list so that probes
already placed on freeing_list are not processed a second time.

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
Changes in v2:
 - Drop the unreachable kprobe_unused(&op->kp) branch when force is true
   and list_empty(&op->list), since kprobe_disarmed() already checked for
   disabled probes with an empty op->list at the start of
   unoptimize_kprobe().
---
 kernel/kprobes.c |   21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index b6ff971df451..c0637302eeae 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -552,17 +552,13 @@ static unsigned long optimizer_passes;
 
 #define OPTIMIZE_DELAY 5
 
-/*
- * Note: gen->cooling_list is not checked here because it is strictly
- * an in-flight holding list for dispatched generations, so it is always
- * empty in optprobe_cur_gen.
- */
 static bool optprobe_has_queued_probes(void)
 {
 	struct optprobe_generation *gen = &optprobe_gens[optprobe_cur_gen];
 
 	return !list_empty(&gen->optimizing_list) ||
 	       !list_empty(&gen->unoptimizing_list) ||
+	       !list_empty(&gen->cooling_list) ||
 	       !list_empty(&gen->freeing_list);
 }
 
@@ -612,6 +608,7 @@ static void do_unoptimize_kprobes(struct list_head *unopt_list,
 				  struct list_head *cooling_list)
 {
 	struct optimized_kprobe *op, *tmp;
+	LIST_HEAD(unoptimized);
 
 	lockdep_assert_held(&text_mutex);
 	/*
@@ -626,11 +623,13 @@ static void do_unoptimize_kprobes(struct list_head *unopt_list,
 	 */
 	lockdep_assert_cpus_held();
 
-	if (!list_empty(unopt_list))
-		arch_unoptimize_kprobes(unopt_list, free_list);
+	if (list_empty(unopt_list))
+		return;
 
-	/* Loop on 'free_list' for disarming and removing from kprobe hash list */
-	list_for_each_entry_safe(op, tmp, free_list, list) {
+	arch_unoptimize_kprobes(unopt_list, &unoptimized);
+
+	/* Loop on newly unoptimized probes for disarming and hash list removal */
+	list_for_each_entry_safe(op, tmp, &unoptimized, list) {
 		/* Switching from detour code to origin */
 		op->kp.flags &= ~KPROBE_FLAG_OPTIMIZED;
 		/*
@@ -647,6 +646,7 @@ static void do_unoptimize_kprobes(struct list_head *unopt_list,
 			 * (reclaiming is done by do_free_cleaned_kprobes().)
 			 */
 			hlist_del_init_rcu(&op->kp.hlist);
+			list_move(&op->list, free_list);
 		} else {
 			/*
 			 * Keep on cooling_list until the quiescence period
@@ -1044,10 +1044,11 @@ static void unoptimize_kprobe(struct kprobe *p, bool force)
 	if (force) {
 		/* Forcibly update the code: this is a special case */
 		force_unoptimize_kprobe(op);
+		list_add(&op->list, &cur->cooling_list);
 	} else {
 		list_add(&op->list, &cur->unoptimizing_list);
-		kick_kprobe_optimizer();
 	}
+	kick_kprobe_optimizer();
 }
 
 /* Cancel unoptimizing for reusing */

[PATCH v3 3/5] kprobes: Do not overwrite jump instruction when re-arming unoptimizing probe

[Masami Hiramatsu (Google)]

From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

When an optimized probe is temporarily disarmed or disabled,
unoptimize_kprobe() places it on the unoptimizing list, but leaves
the relative jump in text until the optimizer thread runs.
__disarm_kprobe() explicitly skips arch_disarm_kprobe() when
kprobe_queued_unopt(p) is true.

However, if __arm_kprobe() is called before the unoptimization pass
executes, it unconditionally calls arch_arm_kprobe(p), writing a
breakpoint instruction over the first byte of the jump while leaving the
trailing bytes intact. optimize_kprobe(p) then dequeues the probe from
the unoptimizing queue, leaving the jump opcode overwritten with a
breakpoint in text while KPROBE_FLAG_OPTIMIZED remains set.

Skip calling arch_arm_kprobe() in __arm_kprobe() when
kprobe_queued_unopt(p) is true, keeping the existing relative jump intact
while unoptimization is cancelled. If optimize_kprobe() cannot optimize
the probe (e.g. due to a post_handler or disabled optimization), force
unoptimization immediately.

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
Changes in v3:
 - Restore the __arm_kprobe() fix that was accidentally dropped in v2
   while omitting the reuse_unused_kprobe() hunk superseded by commit
   1ffd6efb1b58 ("kprobes: Arm reused optprobe on cooling list").
---
 kernel/kprobes.c |   31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index c0637302eeae..f7cf77ca70de 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1263,6 +1263,19 @@ static void __init kprobe_sysctls_init(void)
 }
 #endif /* CONFIG_SYSCTL */
 
+/* Return true if the probe is queued on unoptimizing lists */
+static bool kprobe_queued_unopt(struct kprobe *p)
+{
+	struct optimized_kprobe *op;
+
+	if (kprobe_aggrprobe(p)) {
+		op = container_of(p, struct optimized_kprobe, kp);
+		if (optprobe_queued_unopt(op))
+			return true;
+	}
+	return false;
+}
+
 /* Put a breakpoint for a probe. */
 static void __arm_kprobe(struct kprobe *p)
 {
@@ -1276,21 +1289,11 @@ static void __arm_kprobe(struct kprobe *p)
 		/* Fallback to unoptimized kprobe */
 		unoptimize_kprobe(_p, true);
 
-	arch_arm_kprobe(p);
+	if (!kprobe_queued_unopt(p))
+		arch_arm_kprobe(p);
 	optimize_kprobe(p);	/* Try to optimize (add kprobe to a list) */
-}
-
-/* Return true if the probe is queued on unoptimizing lists */
-static bool kprobe_queued_unopt(struct kprobe *p)
-{
-	struct optimized_kprobe *op;
-
-	if (kprobe_aggrprobe(p)) {
-		op = container_of(p, struct optimized_kprobe, kp);
-		if (optprobe_queued_unopt(op))
-			return true;
-	}
-	return false;
+	if (unlikely(kprobe_queued_unopt(p)))
+		unoptimize_kprobe(p, true);
 }
 
 /* Remove the breakpoint of a probe. */

[PATCH v3 4/5] kprobes: Re-optimize blocked sibling probe when cooling finishes

[Masami Hiramatsu (Google)]

From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

When an optimized kprobe is disabled (for example, to arm a new kprobe
in its jump-optimized instruction range), its jump is replaced with a
breakpoint (or disarmed) and the probe is placed on the cooling list
until the Tasks RCU grace period completes. While the disabled probe is
still on the cooling list, any sibling probe whose jump range covers
this probe's address cannot be optimized because can_optimize() sees the
not-yet-disarmed probe via get_kprobe().

In optprobe_finalize_generation(), when a disabled probe on cooling_list
finishes cooling and transitions to safely disarmed (list_del_init()),
check whether an optimized sibling probe covers op->kp.addr via
get_optimized_kprobe() and re-queue it for optimization, matching the
freeing_list path in do_free_cleaned_kprobes().

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
Changes in v2:
 - Drop the unoptimize_kprobe() and reuse_unused_kprobe() hunks which
   were superseded by commits 0be7a931c5e4 and 1ffd6efb1b58 in
   probes/fixes.
 - Drop the kill_optimized_kprobe() cooling deferral since kill_kprobe()
   runs at MODULE_STATE_GOING when no tasks execute the module text.
---
 kernel/kprobes.c |    9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index f7cf77ca70de..f14f7ff45961 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -729,10 +729,17 @@ static void optprobe_finalize_generation(struct optprobe_generation *gen)
 			list_move(&op->list, &optprobe_gens[optprobe_cur_gen].freeing_list);
 			kick_kprobe_optimizer();
 		} else {
+			struct kprobe *_p;
+
 			/* Still in use; now safely disarmed */
 			list_del_init(&op->list);
-			if (!kprobe_disabled(&op->kp))
+			if (!kprobe_disabled(&op->kp)) {
 				optimize_kprobe(&op->kp);
+			} else {
+				_p = get_optimized_kprobe(op->kp.addr);
+				if (unlikely(_p))
+					optimize_kprobe(_p);
+			}
 		}
 	}
 

[PATCH v3 5/5] kprobes: Add KUnit test cases for kprobes regressions

[Masami Hiramatsu (Google)]

From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

Add KUnit test cases to lib/tests/test_kprobes.c to verify that the
recent regressions in kprobes registration, ftrace ops switching, and
optimization state transitions are resolved:

1. test_kretprobe_unregistered_cleanup:
   Verifies that unregistering an uninitialized or already-unregistered
   kretprobe does not cause a NULL pointer dereference or double free.

2. test_kretprobe_rereg_by_symbol:
   Verifies that attempting to re-register an active kretprobe by symbol
   name fails with -EINVAL and does not destroy the active probe's
   underlying rethook state.

3. test_kprobe_ftrace_post_handler_switch:
   Verifies that adding a second kprobe with a post_handler to a probe
   point armed with kprobe_ftrace_ops correctly switches to
   kprobe_ipmodify_ops and both probes cleanly unregister without
   disarming errors (-ENOENT).

4. test_optprobe_rearm_unoptimizing:
   Verifies that disabling an optimized kprobe on a notrace target and
   immediately re-enabling it before the optimizer thread runs works
   cleanly without overwriting the jump instruction or crashing.

Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
Changes in v3:
 - Check that the jump opcode at *opt_kp.addr is preserved across
   disable_kprobe() and enable_kprobe() in
   test_optprobe_rearm_unoptimizing().
Changes in v2:
 - Probe a notrace target function (kprobe_opt_target) in
   test_optprobe_rearm_unoptimizing() so that jump optimization is tested
   even when CONFIG_KPROBES_ON_FTRACE=y.
 - Guard test_optprobe_rearm_unoptimizing() with
   !IS_MODULE(CONFIG_KPROBES_SANITY_TEST) since
   wait_for_kprobe_optimizer() is not exported to modules.
 - Replace KUNIT_ASSERT_EQ() with KUNIT_EXPECT_EQ() while stack-allocated
   probes are registered to prevent stack use-after-free on test failure.
---
 lib/tests/test_kprobes.c |  139 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 139 insertions(+)

diff --git a/lib/tests/test_kprobes.c b/lib/tests/test_kprobes.c
index 06e729e4de05..df3f080b1852 100644
--- a/lib/tests/test_kprobes.c
+++ b/lib/tests/test_kprobes.c
@@ -358,8 +358,136 @@ static void test_stacktrace_on_nested_kretprobe(struct kunit *test)
 }
 #endif /* CONFIG_ARCH_CORRECT_STACKTRACE_ON_KRETPROBE */
 
+/* Test for: kprobes: Fix double free and NULL deref in unregister_kretprobes() */
+static void test_kretprobe_unregistered_cleanup(struct kunit *test)
+{
+	struct kretprobe unreg_rp = {
+		.kp.symbol_name = "kprobe_target",
+	};
+
+	current_test = test;
+
+	/* Unregistering a probe that was never registered must return cleanly */
+	unregister_kretprobe(&unreg_rp);
+
+	/* Double unregister of a valid registered probe */
+	KUNIT_ASSERT_EQ(test, 0, register_kretprobe(&unreg_rp));
+	unregister_kretprobe(&unreg_rp);
+	/* Second unregister must not cause double-free or NULL pointer dereference */
+	unregister_kretprobe(&unreg_rp);
+}
+
+/* Test for: kprobes: Validate re-registration earlier in register_kretprobe() */
+static void test_kretprobe_rereg_by_symbol(struct kunit *test)
+{
+	struct kretprobe test_rp = {
+		.kp.symbol_name = "kprobe_target",
+		.handler = return_handler,
+		.entry_handler = entry_handler,
+	};
+
+	current_test = test;
+	KUNIT_ASSERT_EQ(test, 0, register_kretprobe(&test_rp));
+
+	/* Verify the probe is live and functional */
+	krph_val = 0;
+	target(rand1);
+	KUNIT_EXPECT_EQ(test, krph_val, rand1);
+
+	/*
+	 * Attempt re-registration with symbol_name (with kp.addr cleared).
+	 * Must fail with -EINVAL and NOT destroy the live probe's rethook.
+	 */
+	test_rp.kp.addr = NULL;
+	KUNIT_EXPECT_EQ(test, -EINVAL, register_kretprobe(&test_rp));
+
+	/* The original kretprobe must still be functional */
+	krph_val = 0;
+	target(rand1);
+	KUNIT_EXPECT_EQ(test, krph_val, rand1);
+
+	unregister_kretprobe(&test_rp);
+}
+
 #endif /* CONFIG_KRETPROBES */
 
+#if defined(CONFIG_KPROBES_ON_FTRACE)
+/* Test for: kprobes: Switch ftrace ops when adding probe with post_handler */
+static void test_kprobe_ftrace_post_handler_switch(struct kunit *test)
+{
+	struct kprobe kp_entry = {
+		.symbol_name = "kprobe_target",
+		.pre_handler = kp_pre_handler,
+	};
+	struct kprobe kp_post = {
+		.symbol_name = "kprobe_target",
+		.post_handler = kp_post_handler,
+	};
+	int ret;
+
+	current_test = test;
+
+	/* 1. Register first probe without post_handler (armed with kprobe_ftrace_ops) */
+	KUNIT_ASSERT_EQ(test, 0, register_kprobe(&kp_entry));
+
+	/* 2. Register second probe with post_handler at same IP (must switch ops) */
+	ret = register_kprobe(&kp_post);
+	KUNIT_EXPECT_EQ(test, 0, ret);
+	if (!ret) {
+		/* 3. Execute target function to trigger both */
+		preh_val = 0;
+		posth_val = 0;
+		target(rand1);
+		KUNIT_EXPECT_NE(test, 0, preh_val);
+		KUNIT_EXPECT_NE(test, 0, posth_val);
+
+		/* 4. Unregistering must successfully disarm without -ENOENT */
+		unregister_kprobe(&kp_post);
+	}
+	unregister_kprobe(&kp_entry);
+}
+#endif
+
+#if defined(CONFIG_OPTPROBES) && !IS_MODULE(CONFIG_KPROBES_SANITY_TEST)
+static u32 (*opt_target)(u32 value);
+
+static noinline notrace u32 kprobe_opt_target(u32 value)
+{
+	return (value / div_factor) + 2;
+}
+
+/* Test for: kprobes: Do not overwrite jump instruction when re-arming unoptimizing probe */
+static void test_optprobe_rearm_unoptimizing(struct kunit *test)
+{
+	struct kprobe opt_kp = {
+		.symbol_name = "kprobe_opt_target",
+		.pre_handler = kp_pre_handler,
+	};
+	kprobe_opcode_t jmp_op;
+
+	current_test = test;
+	KUNIT_ASSERT_EQ(test, 0, register_kprobe(&opt_kp));
+
+	/* Allow background optimizer thread to jump-optimize the probe */
+	wait_for_kprobe_optimizer();
+	jmp_op = READ_ONCE(*opt_kp.addr);
+
+	/* Temporarily disable probe (queues on unoptimizing_list, leaves jump) */
+	KUNIT_EXPECT_EQ(test, 0, disable_kprobe(&opt_kp));
+
+	/* Immediately re-enable (re-arm) before optimizer thread unoptimizes */
+	KUNIT_EXPECT_EQ(test, 0, enable_kprobe(&opt_kp));
+
+	/* Verify jump opcode was not overwritten by a breakpoint on re-arm */
+	KUNIT_EXPECT_EQ(test, jmp_op, READ_ONCE(*opt_kp.addr));
+	preh_val = 0;
+	opt_target(rand1);
+	KUNIT_EXPECT_NE(test, 0, preh_val);
+
+	unregister_kprobe(&opt_kp);
+}
+#endif
+
 static int kprobes_test_init(struct kunit *test)
 {
 	KP_CLEAR(kp);
@@ -380,6 +508,9 @@ static int kprobes_test_init(struct kunit *test)
 	stacktrace_target = kprobe_stacktrace_target;
 	internal_target = kprobe_stacktrace_internal_target;
 	stacktrace_driver = kprobe_stacktrace_driver;
+#if defined(CONFIG_OPTPROBES) && !IS_MODULE(CONFIG_KPROBES_SANITY_TEST)
+	opt_target = kprobe_opt_target;
+#endif
 	rand1 = get_random_u32_above(div_factor);
 	return 0;
 }
@@ -391,10 +522,18 @@ static struct kunit_case kprobes_testcases[] = {
 #ifdef CONFIG_KRETPROBES
 	KUNIT_CASE(test_kretprobe),
 	KUNIT_CASE(test_kretprobes),
+	KUNIT_CASE(test_kretprobe_unregistered_cleanup),
+	KUNIT_CASE(test_kretprobe_rereg_by_symbol),
 #ifdef CONFIG_ARCH_CORRECT_STACKTRACE_ON_KRETPROBE
 	KUNIT_CASE(test_stacktrace_on_kretprobe),
 	KUNIT_CASE(test_stacktrace_on_nested_kretprobe),
 #endif
+#endif
+#if defined(CONFIG_KPROBES_ON_FTRACE)
+	KUNIT_CASE(test_kprobe_ftrace_post_handler_switch),
+#endif
+#if defined(CONFIG_OPTPROBES) && !IS_MODULE(CONFIG_KPROBES_SANITY_TEST)
+	KUNIT_CASE(test_optprobe_rearm_unoptimizing),
 #endif
 	{}
 };