[PATCH 0/8] kprobes: Fixes for registration errors, ftrace ops and optprobe state machine

[PATCH 0/8] kprobes: Fixes for registration errors, ftrace ops and optprobe state machine

[Masami Hiramatsu (Google)]

Hi,

Here is a series of bug fixes for kprobes and optprobes, addressing error
recovery, ftrace ops management, and optprobe asynchronous state machine
edge cases:

1. Registration & Unregistration Error Handling (Patches 1-3):
   - In __register_kprobe(), when arm_kprobe() fails, the error code was
     ignored and the instruction slot was leaked.
   - In unregister_kretprobes(), failing __unregister_kprobe_top() still
     invoked rethook/rph freeing, causing a double free or NULL deref.
     Clearing rp->rh is also deferred until after synchronize_rcu() to
     prevent racing with in-flight pre_handler_kretprobe() invocations.
   - In register_kretprobe(), re-registration validation was skipped when
     registered by symbol_name without an explicit addr, leading to
     destruction of the live probe's rethook instance on failure.

2. Ftrace Ops Switching for Post-Handlers (Patch 4):
   - 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.

3. Optprobe State Machine & Tasks RCU Quiescence (Patches 5-7):
   - 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.
   - Re-arming an unoptimizing probe before the unoptimization pass ran
     overwrote the first byte of the relative jump with a breakpoint,
     corrupting the jump opcode in text.
   - Probes killed at module unload while child probes were still active
     were prematurely freed or improperly unhashed before Tasks RCU
     quiescence elapsed.

In addition to the bugfixes, the last patch adds KUnit test cases for
these bugs. (Patch 8)

Summary of patches in this series:

[01/08] kprobes: Fix error return and slot leak in __register_kprobe()
 - Release allocated instruction slot via arch_remove_kprobe(p) and
   return the arm_kprobe() error code upon failure.

[02/08] kprobes: Fix double free and NULL deref in unregister_kretprobes()
 - Skip rethook_free() and rph cleanup if __unregister_kprobe_top()
   fails, and clear rp->rh after synchronize_rcu().

[03/08] kprobes: Validate re-registration earlier in register_kretprobe()
 - Resolve target address via _kprobe_addr() up-front and check whether
   the probe is already registered via check_kprobe_rereg_at() before
   allocating rethook resources, dropping the WARN_ON_ONCE().

[04/08] kprobes: Switch ftrace ops when adding probe with post_handler
 - Switch from kprobe_ftrace_ops to kprobe_ipmodify_ops in
   register_aggr_kprobe() outside of text_mutex and cpus_read_lock.
 - Restore ftrace ops counter and registration if ftrace_set_filter_ip()
   fails during disarm.

[05/08] kprobes: Fix UAF and queue cooling probes on forced unoptimization
 - Queue forcibly unoptimized probes onto cur->cooling_list (or
   cur->freeing_list if unused), include cooling_list in
   optprobe_has_queued_probes(), and isolate newly unoptimized probes
   into a local list in do_unoptimize_kprobes() to prevent double disarm.

[06/08] kprobes: Do not overwrite jump instruction when re-arming
        unoptimizing probe
 - Skip arch_arm_kprobe() in __arm_kprobe() when kprobe_queued_unopt()
   is true, keeping the relative jump intact while unoptimization is
   cancelled.
 - Force unoptimization immediately if optimize_kprobe() cannot optimize
   the re-armed probe.
 - Avoid premature flag clearing or dequeuing in reuse_unused_kprobe().

[07/08] kprobes: Fix slot reclaim and sibling re-opt in optprobe state machine
 - In kill_optimized_kprobe(), queue in-use probes on unoptimizing_list
   to cooling_list rather than freeing_list to preserve kprobe_table
   lookup and ensure Tasks RCU elapses before slot reclamation.
 - Reclaim optinsn slots for kprobe_gone() probes upon cooling completion
   in optprobe_finalize_generation().
 - Re-optimize blocked sibling probes when a cooling probe finishes
   cooling to safely disarmed.

[08/08] kprobes: Add KUnit test cases for kprobes regressions
 - Add test cases to lib/tests/test_kprobes.c covering unregistered
   kretprobe cleanup, kretprobe re-registration validation by symbol,
   ftrace ops dynamic switching on post_handler addition, and optprobe
   re-arming while unoptimizing.

Thank you,

---
base-commit: 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4

Masami Hiramatsu (Google) (8):
      kprobes: Fix error return and slot leak in __register_kprobe()
      kprobes: Fix double free and NULL deref in unregister_kretprobes()
      kprobes: Validate re-registration earlier in register_kretprobe()
      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: Fix slot reclaim and sibling re-opt in optprobe state machine
      kprobes: Add KUnit test cases for kprobes regressions


 kernel/kprobes.c         |  226 ++++++++++++++++++++++++++++++----------------
 lib/tests/test_kprobes.c |  123 +++++++++++++++++++++++++
 2 files changed, 271 insertions(+), 78 deletions(-)

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

[PATCH 1/8] kprobes: Fix error return and slot leak in __register_kprobe()

[Masami Hiramatsu (Google)]

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

When arm_kprobe() fails in __register_kprobe(), the probe is removed
from kprobe_table and synchronize_rcu() is executed, but the error code
ret was ignored and try_to_optimize_kprobe(p) was executed, returning 0.
Additionally, the instruction slot allocated by prepare_kprobe() was
leaked.

Fix this by releasing the instruction slot via arch_remove_kprobe(p)
and returning ret when arm_kprobe() fails.

Fixes: 587e8e6d640b ("kprobes: Adopt guard() and scoped_guard()")
Cc: stable@vger.kernel.org
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |    2 ++
 1 file changed, 2 insertions(+)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 2441b9e39343..10946658bb35 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1946,6 +1946,8 @@ static int __register_kprobe(struct kprobe *p)
 		if (ret) {
 			hlist_del_rcu(&p->hlist);
 			synchronize_rcu();
+			arch_remove_kprobe(p);
+			return ret;
 		}
 	}
 

[PATCH 2/8] kprobes: Fix double free and NULL deref in unregister_kretprobes()

[Masami Hiramatsu (Google)]

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

In unregister_kretprobes(), if __unregister_kprobe_top() returns an error
(e.g., if the kretprobe was not registered or was already unregistered),
rethook_free(rps[i]->rh) or rcu_assign_pointer(rps[i]->rph->rp, NULL) is
still executed. For rethook-based kretprobes, calling rethook_free() on an
already-freed rethook causes a double-free / list corruption in RCU. For
legacy kretprobes, rps[i]->rph is NULL after previous cleanup, resulting
in a NULL pointer dereference.

Skip rethook_free() and rph cleanup if __unregister_kprobe_top() fails,
and clear rps[i]->rh after synchronize_rcu() has waited for in-flight
pre_handler_kretprobe() invocations to complete.

Fixes: 73f9b911faa7 ("kprobes: Use rethook for kretprobe if possible")
Cc: stable@vger.kernel.org
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |    8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 10946658bb35..7a3da809bac7 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -2605,8 +2605,10 @@ void unregister_kretprobes(struct kretprobe **rps, int num)
 	for (i = 0; i < num; i++) {
 		guard(mutex)(&kprobe_mutex);
 
-		if (__unregister_kprobe_top(&rps[i]->kp) < 0)
+		if (__unregister_kprobe_top(&rps[i]->kp) < 0) {
 			rps[i]->kp.addr = NULL;
+			continue;
+		}
 #ifdef CONFIG_KRETPROBE_ON_RETHOOK
 		rethook_free(rps[i]->rh);
 #else
@@ -2618,7 +2620,9 @@ void unregister_kretprobes(struct kretprobe **rps, int num)
 	for (i = 0; i < num; i++) {
 		if (rps[i]->kp.addr) {
 			__unregister_kprobe_bottom(&rps[i]->kp);
-#ifndef CONFIG_KRETPROBE_ON_RETHOOK
+#ifdef CONFIG_KRETPROBE_ON_RETHOOK
+			rps[i]->rh = NULL;
+#else
 			free_rp_inst(rps[i]);
 #endif
 		}

[PATCH 3/8] kprobes: Validate re-registration earlier in register_kretprobe()

[Masami Hiramatsu (Google)]

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

register_kretprobe() only checked for re-registration if rp->kp.addr was
specified, skipping the check when rp was registered with symbol_name
and rp->kp.addr was NULL. In that scenario, register_kretprobe()
proceeded to allocate a new rethook (or kretprobe holder), and when
register_kprobe(&rp->kp) subsequently failed due to re-registration, the
error path freed and NULL-ed rp->rh (or rp->rph) on the still-active
kretprobe, leading to NULL pointer dereferences when the active probe
was hit.

Resolve the target address via _kprobe_addr() up-front and check whether
the probe is already registered at that address via
check_kprobe_rereg_at() before allocating any resources, and drop the
WARN_ON_ONCE() in the re-registration check now that re-registration is
gracefully detected and rejected.

Fixes: 0188b87899ff ("kretprobe: Avoid re-registration of the same kretprobe earlier")
Cc: stable@vger.kernel.org
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |   72 ++++++++++++++++++++++++++++++------------------------
 1 file changed, 40 insertions(+), 32 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 7a3da809bac7..0b41d69540ae 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1793,52 +1793,61 @@ _kprobe_addr(kprobe_opcode_t *addr, const char *symbol_name,
 	return addr;
 }
 
-static kprobe_opcode_t *kprobe_addr(struct kprobe *p)
-{
-	bool on_func_entry;
-
-	return _kprobe_addr(p->addr, p->symbol_name, p->offset, &on_func_entry);
-}
-
 /*
  * Check the 'p' is valid and return the aggregator kprobe
  * at the same address.
  */
-static struct kprobe *__get_valid_kprobe(struct kprobe *p)
+static struct kprobe *__get_valid_kprobe_at(struct kprobe *p, kprobe_opcode_t *addr)
 {
+	struct hlist_head *head;
 	struct kprobe *ap, *list_p;
 
 	lockdep_assert_held(&kprobe_mutex);
 
-	ap = get_kprobe(p->addr);
-	if (unlikely(!ap))
+	if (unlikely(!addr))
 		return NULL;
 
-	if (p == ap)
-		return ap;
-
-	list_for_each_entry(list_p, &ap->list, list)
-		if (list_p == p)
-		/* kprobe p is a valid probe */
+	head = &kprobe_table[hash_ptr(addr, KPROBE_HASH_BITS)];
+	hlist_for_each_entry_rcu(ap, head, hlist,
+				 lockdep_is_held(&kprobe_mutex)) {
+		if (ap == p) {
+			p->addr = addr;
 			return ap;
+		}
+		if (ap->addr == addr && kprobe_aggrprobe(ap)) {
+			list_for_each_entry(list_p, &ap->list, list) {
+				if (list_p == p) {
+					p->addr = addr;
+					return ap;
+				}
+			}
+		}
+	}
 
 	return NULL;
 }
 
-/*
- * Warn and return error if the kprobe is being re-registered since
- * there must be a software bug.
- */
-static inline int warn_kprobe_rereg(struct kprobe *p)
+static inline struct kprobe *__get_valid_kprobe(struct kprobe *p)
+{
+	return __get_valid_kprobe_at(p, p->addr);
+}
+
+/* Return error if the kprobe is being re-registered */
+static inline int check_kprobe_rereg_at(struct kprobe *p, kprobe_opcode_t *addr)
 {
 	guard(mutex)(&kprobe_mutex);
 
-	if (WARN_ON_ONCE(__get_valid_kprobe(p)))
+	if (__get_valid_kprobe_at(p, addr))
 		return -EINVAL;
 
 	return 0;
 }
 
+static inline int check_kprobe_rereg(struct kprobe *p)
+{
+	return check_kprobe_rereg_at(p, p->addr);
+}
+
 static int check_ftrace_location(struct kprobe *p)
 {
 	unsigned long addr = (unsigned long)p->addr;
@@ -1969,7 +1978,7 @@ int register_kprobe(struct kprobe *p)
 		return PTR_ERR(addr);
 	p->addr = addr;
 
-	ret = warn_kprobe_rereg(p);
+	ret = check_kprobe_rereg(p);
 	if (ret)
 		return ret;
 
@@ -2505,20 +2514,19 @@ int register_kretprobe(struct kretprobe *rp)
 	int ret;
 	int i;
 	void *addr;
+	bool on_func_entry;
 
-	ret = kprobe_on_func_entry(rp->kp.addr, rp->kp.symbol_name, rp->kp.offset);
-	if (ret)
-		return ret;
+	addr = _kprobe_addr(rp->kp.addr, rp->kp.symbol_name, rp->kp.offset, &on_func_entry);
+	if (IS_ERR(addr))
+		return PTR_ERR(addr);
 
-	/* If only 'rp->kp.addr' is specified, check reregistering kprobes */
-	if (rp->kp.addr && warn_kprobe_rereg(&rp->kp))
+	if (!on_func_entry)
 		return -EINVAL;
 
-	if (kretprobe_blacklist_size) {
-		addr = kprobe_addr(&rp->kp);
-		if (IS_ERR(addr))
-			return PTR_ERR(addr);
+	if (check_kprobe_rereg_at(&rp->kp, addr))
+		return -EINVAL;
 
+	if (kretprobe_blacklist_size) {
 		for (i = 0; kretprobe_blacklist[i].name != NULL; i++) {
 			if (kretprobe_blacklist[i].addr == addr)
 				return -EINVAL;

[PATCH 4/8] 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. 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>
---
 kernel/kprobes.c |   56 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 46 insertions(+), 10 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 0b41d69540ae..135fec5d49cf 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1435,9 +1435,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)
@@ -1566,16 +1571,12 @@ 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)
 {
 	if (p->post_handler)
 		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;
 }
 
 /*
@@ -1608,6 +1609,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 */
@@ -1654,18 +1657,51 @@ 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);
+		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) {
+			list_del_rcu(&p->list);
+			synchronize_rcu();
+			return ret;
+		}
+		ap->post_handler = aggr_post_handler;
+		ret = arm_kprobe_ftrace(ap);
+		if (ret) {
+			ap->post_handler = NULL;
+			arm_kprobe_ftrace(ap);
+			list_del_rcu(&p->list);
+			synchronize_rcu();
+			return ret;
+		}
 	}
 
-	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 5/8] 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 or cur->freeing_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->freeing_list (if unused) or
cur->cooling_list (if still in use), 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 disarmed a second time.

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |   27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 135fec5d49cf..ceacd56245c2 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;
 		/* Disarm probes if marked disabled and not gone */
@@ -643,6 +642,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
@@ -1039,6 +1039,15 @@ static void unoptimize_kprobe(struct kprobe *p, bool force)
 	if (force) {
 		/* Forcibly update the code: this is a special case */
 		force_unoptimize_kprobe(op);
+		if (kprobe_unused(&op->kp)) {
+			if (!kprobe_gone(&op->kp))
+				arch_disarm_kprobe(&op->kp);
+			hlist_del_init_rcu(&op->kp.hlist);
+			list_add(&op->list, &cur->freeing_list);
+		} else {
+			list_add(&op->list, &cur->cooling_list);
+		}
+		kick_kprobe_optimizer();
 	} else {
 		list_add(&op->list, &cur->unoptimizing_list);
 		kick_kprobe_optimizer();

[PATCH 6/8] 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 called arch_arm_kprobe(p), writing a
breakpoint instruction over the first byte of the jump while leaving the
trailing 4 bytes intact. optimize_kprobe(p) then dequeued the probe from
the unoptimizing queue, leaving the jump opcode overwritten and corrupted
in text.

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. Also, in reuse_unused_kprobe(), do not
prematurely clear KPROBE_FLAG_DISABLED or call optimize_kprobe(), leaving
arming and optimization to register_aggr_kprobe().

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |   39 +++++++++++++++++----------------------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index ceacd56245c2..0ef3a8701eda 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -1059,19 +1059,11 @@ static int reuse_unused_kprobe(struct kprobe *ap)
 {
 	struct optimized_kprobe *op;
 
-	/*
-	 * Unused kprobe MUST be on the way of delayed unoptimizing (means
-	 * there is still a relative jump) and disabled.
-	 */
 	op = container_of(ap, struct optimized_kprobe, kp);
 	WARN_ON_ONCE(list_empty(&op->list));
-	/* Enable the probe again */
-	ap->flags &= ~KPROBE_FLAG_DISABLED;
-	/* Optimize it again. (remove from 'op->list') */
 	if (!kprobe_optready(ap))
 		return -EINVAL;
 
-	optimize_kprobe(ap);
 	return 0;
 }
 
@@ -1262,6 +1254,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)
 {
@@ -1275,21 +1280,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 7/8] kprobes: Fix slot reclaim and sibling re-opt in optprobe state machine

[Masami Hiramatsu (Google)]

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

Several edge cases in the optprobe state machine can prematurely release
instruction slots or stall blocked sibling probes:

1. In kill_optimized_kprobe(), when a probed module is unloaded while child
probes still exist on p->list, the probe is not unused (!kprobe_unused(p)).
If op was queued for unoptimization, move it to the cooling list rather
than immediately calling arch_remove_optimized_kprobe(), ensuring Tasks
RCU elapses before slot reclamation. In optprobe_finalize_generation(),
reclaim the optinsn slot once cooling finishes for kprobe_gone() probes.

2. In optprobe_finalize_generation(), when a disabled cooling probe
finishes cooling and transitions to safely disarmed, re-optimize any
blocked sibling probe at the same address.

Fixes: 06e0d3b89da1 ("kprobes: Make optprobe optimizer multi-generational and asynchronous")
Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 kernel/kprobes.c |   22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 0ef3a8701eda..53725c09ef9b 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -725,10 +725,19 @@ 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_gone(&op->kp)) {
+				arch_remove_optimized_kprobe(op);
+			} else if (!kprobe_disabled(&op->kp)) {
 				optimize_kprobe(&op->kp);
+			} else {
+				_p = get_optimized_kprobe(op->kp.addr);
+				if (unlikely(_p))
+					optimize_kprobe(_p);
+			}
 		}
 	}
 
@@ -1084,13 +1093,20 @@ static void kill_optimized_kprobe(struct kprobe *p)
 				list_del_init(&op->list);
 			}
 		} else {
-			list_del_init(&op->list);
+			if (optprobe_queued_unopt(op)) {
+				/* Move to cooling list so optinsn slot is cooled down */
+				list_move(&op->list,
+					  &optprobe_gens[optprobe_cur_gen].cooling_list);
+				kick_kprobe_optimizer();
+			} else if (!optprobe_queued_cooling(op) && !optprobe_queued_freeing(op)) {
+				list_del_init(&op->list);
+			}
 		}
 	}
 	op->kp.flags &= ~KPROBE_FLAG_OPTIMIZED;
 
 	/*
-	 * Don't remove the slot if it is queued for freeing or unoptimization;
+	 * Don't remove the slot if it is queued for freeing, unoptimization, or cooling;
 	 * the optimizer will reclaim it after the quiescence period.
 	 */
 	if (!optprobe_queued_freeing(op) && !optprobe_queued_unopt(op) &&

[PATCH 8/8] 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 and immediately
   re-enabling it before the optimizer thread runs does not overwrite the
   jump instruction opcode or crash on execution.

Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 lib/tests/test_kprobes.c |  123 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)

diff --git a/lib/tests/test_kprobes.c b/lib/tests/test_kprobes.c
index 06e729e4de05..c20f70cc63d7 100644
--- a/lib/tests/test_kprobes.c
+++ b/lib/tests/test_kprobes.c
@@ -358,8 +358,123 @@ 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,
+	};
+
+	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) */
+	KUNIT_ASSERT_EQ(test, 0, register_kprobe(&kp_post));
+
+	/* 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)
+/* 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_target",
+		.pre_handler = kp_pre_handler,
+	};
+
+	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();
+
+	/* Temporarily disable probe (queues on unoptimizing_list, leaves jump) */
+	KUNIT_ASSERT_EQ(test, 0, disable_kprobe(&opt_kp));
+
+	/* Immediately re-enable (re-arm) before optimizer thread unoptimizes */
+	KUNIT_ASSERT_EQ(test, 0, enable_kprobe(&opt_kp));
+
+	/* Verify jump opcode was not corrupted by an unaligned INT3 overwrite */
+	preh_val = 0;
+	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);
@@ -391,10 +506,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)
+	KUNIT_CASE(test_optprobe_rearm_unoptimizing),
 #endif
 	{}
 };