[PATCH v2 4/4] x86/kmsan: Don't call instrumented code from kmsan_virt_addr_valid()

[Viorel Cernateanu via B4 Relay <devnull+vrilutza.gmail.com@kernel.org>]

From: Viorel Cernateanu <vrilutza@gmail.com>

With CONFIG_DEBUG_PREEMPT enabled, a KMSAN kernel hangs during boot
immediately after KMSAN is enabled.

kmsan_virt_addr_valid() calls preempt_disable() and pfn_valid(), which
in turn calls rcu_read_lock_sched(). With CONFIG_DEBUG_PREEMPT, these
paths call the instrumented preempt_count_add(). Its instrumentation
requests metadata and re-enters kmsan_virt_addr_valid(), causing
recursion.

Use preempt_disable_notrace() and __pfn_valid() to validate the PFN
without entering those instrumented helpers. Disabling preemption
provides the RCU-sched read-side protection required by __pfn_valid().

Use preempt_enable_no_resched_notrace() on exit to avoid both
instrumentation and a reschedule from the metadata lookup.

Fixes: f6564fce256a ("mm, kmsan: fix infinite recursion due to RCU critical section")
Link: https://lore.kernel.org/20240308043448.masllzeqwht45d4j@M910t
Suggested-by: Alexander Potapenko <glider@google.com>
Signed-off-by: Viorel Cernateanu <vrilutza@gmail.com>
---
 arch/x86/include/asm/kmsan.h | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/arch/x86/include/asm/kmsan.h b/arch/x86/include/asm/kmsan.h
index c6715b2db0dc..e36df95439df 100644
--- a/arch/x86/include/asm/kmsan.h
+++ b/arch/x86/include/asm/kmsan.h
@@ -85,18 +85,21 @@ static inline bool kmsan_virt_addr_valid(void *addr)
 	}
 
 	/*
-	 * pfn_valid() relies on RCU, and may call into the scheduler on exiting
-	 * the critical section. However, this would result in recursion with
-	 * KMSAN. Therefore, disable preemption here, and re-enable preemption
-	 * below while suppressing reschedules to avoid recursion.
+	 * pfn_valid() takes rcu_read_lock_sched(), which can call instrumented
+	 * code (preempt_count_add() with CONFIG_DEBUG_PREEMPT, lock_acquire()
+	 * with CONFIG_DEBUG_LOCK_ALLOC). That code looks up metadata itself and
+	 * would recurse here. Disable preemption with the notrace helpers
+	 * instead, which keeps this an RCU-sched read-side critical section,
+	 * and re-enable it without a reschedule to avoid entering the scheduler
+	 * from the metadata lookup.
 	 *
 	 * Note, this sacrifices occasionally breaking scheduling guarantees.
 	 * Although, a kernel compiled with KMSAN has already given up on any
 	 * performance guarantees due to being heavily instrumented.
 	 */
-	preempt_disable();
-	ret = pfn_valid(x >> PAGE_SHIFT);
-	preempt_enable_no_resched();
+	preempt_disable_notrace();
+	ret = __pfn_valid(x >> PAGE_SHIFT);
+	preempt_enable_no_resched_notrace();
 
 	return ret;
 }

-- 
2.53.0