Provide dedicated helpers to (un)register virt hooks used during an
emergency crash/reboot, and WARN if there is an attempt to overwrite
the registered callback, or an attempt to do an unpaired unregister.

Opportunsitically use rcu_assign_pointer() instead of RCU_INIT_POINTER(),
mainly so that the set/unset paths are more symmetrical, but also because
any performance gains from using RCU_INIT_POINTER() are meaningless for
this code.

Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/include/asm/reboot.h |  5 +++--
 arch/x86/kernel/reboot.c      | 30 ++++++++++++++++++++++++------
 arch/x86/kvm/vmx/vmx.c        |  6 ++----
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/arch/x86/include/asm/reboot.h b/arch/x86/include/asm/reboot.h
index 33c8e911e0de..1d098a7d329a 100644
--- a/arch/x86/include/asm/reboot.h
+++ b/arch/x86/include/asm/reboot.h
@@ -26,8 +26,9 @@ void __noreturn machine_real_restart(unsigned int type);
 #define MRR_APM		1
 
 #if IS_ENABLED(CONFIG_KVM_INTEL)
-typedef void crash_vmclear_fn(void);
-extern crash_vmclear_fn __rcu *crash_vmclear_loaded_vmcss;
+typedef void (cpu_emergency_virt_cb)(void);
+void cpu_emergency_register_virt_callback(cpu_emergency_virt_cb *callback);
+void cpu_emergency_unregister_virt_callback(cpu_emergency_virt_cb *callback);
 #endif
 void cpu_emergency_disable_virtualization(void);
 
diff --git a/arch/x86/kernel/reboot.c b/arch/x86/kernel/reboot.c
index 6c0b1634b884..78182b2969db 100644
--- a/arch/x86/kernel/reboot.c
+++ b/arch/x86/kernel/reboot.c
@@ -795,17 +795,35 @@ void machine_crash_shutdown(struct pt_regs *regs)
  *
  * protected by rcu.
  */
-crash_vmclear_fn __rcu *crash_vmclear_loaded_vmcss;
-EXPORT_SYMBOL_GPL(crash_vmclear_loaded_vmcss);
+static cpu_emergency_virt_cb __rcu *cpu_emergency_virt_callback;
+
+void cpu_emergency_register_virt_callback(cpu_emergency_virt_cb *callback)
+{
+	if (WARN_ON_ONCE(rcu_access_pointer(cpu_emergency_virt_callback)))
+		return;
+
+	rcu_assign_pointer(cpu_emergency_virt_callback, callback);
+}
+EXPORT_SYMBOL_GPL(cpu_emergency_register_virt_callback);
+
+void cpu_emergency_unregister_virt_callback(cpu_emergency_virt_cb *callback)
+{
+	if (WARN_ON_ONCE(rcu_access_pointer(cpu_emergency_virt_callback) != callback))
+		return;
+
+	rcu_assign_pointer(cpu_emergency_virt_callback, NULL);
+	synchronize_rcu();
+}
+EXPORT_SYMBOL_GPL(cpu_emergency_unregister_virt_callback);
 
 static inline void cpu_crash_vmclear_loaded_vmcss(void)
 {
-	crash_vmclear_fn *do_vmclear_operation = NULL;
+	cpu_emergency_virt_cb *callback;
 
 	rcu_read_lock();
-	do_vmclear_operation = rcu_dereference(crash_vmclear_loaded_vmcss);
-	if (do_vmclear_operation)
-		do_vmclear_operation();
+	callback = rcu_dereference(cpu_emergency_virt_callback);
+	if (callback)
+		callback();
 	rcu_read_unlock();
 }
 #endif
diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
index 302086255be6..41095fc864f3 100644
--- a/arch/x86/kvm/vmx/vmx.c
+++ b/arch/x86/kvm/vmx/vmx.c
@@ -8551,8 +8551,7 @@ static void __vmx_exit(void)
 {
 	allow_smaller_maxphyaddr = false;
 
-	RCU_INIT_POINTER(crash_vmclear_loaded_vmcss, NULL);
-	synchronize_rcu();
+	cpu_emergency_unregister_virt_callback(crash_vmclear_local_loaded_vmcss);
 
 	vmx_cleanup_l1d_flush();
 }
@@ -8602,8 +8601,7 @@ static int __init vmx_init(void)
 		pi_init_cpu(cpu);
 	}
 
-	rcu_assign_pointer(crash_vmclear_loaded_vmcss,
-			   crash_vmclear_local_loaded_vmcss);
+	cpu_emergency_register_virt_callback(crash_vmclear_local_loaded_vmcss);
 
 	vmx_check_vmcs12_offsets();
 
-- 
2.40.0.rc1.284.g88254d51c5-goog

On Fri, Mar 10, 2023 at 01:42:17PM -0800, Sean Christopherson wrote:
>+void cpu_emergency_register_virt_callback(cpu_emergency_virt_cb *callback)
>+{
>+	if (WARN_ON_ONCE(rcu_access_pointer(cpu_emergency_virt_callback)))
>+		return;
>+
>+	rcu_assign_pointer(cpu_emergency_virt_callback, callback);

Was it intentional to not call synchronize_rcu() (in the original
code), different from the un-registration path?

>+}
>+EXPORT_SYMBOL_GPL(cpu_emergency_register_virt_callback);
>+
>+void cpu_emergency_unregister_virt_callback(cpu_emergency_virt_cb *callback)
>+{
>+	if (WARN_ON_ONCE(rcu_access_pointer(cpu_emergency_virt_callback) != callback))
>+		return;
>+
>+	rcu_assign_pointer(cpu_emergency_virt_callback, NULL);
>+	synchronize_rcu();
>+}
>+EXPORT_SYMBOL_GPL(cpu_emergency_unregister_virt_callback);
> 
> static inline void cpu_crash_vmclear_loaded_vmcss(void)
> {
>-	crash_vmclear_fn *do_vmclear_operation = NULL;
>+	cpu_emergency_virt_cb *callback;
> 
> 	rcu_read_lock();
>-	do_vmclear_operation = rcu_dereference(crash_vmclear_loaded_vmcss);
>-	if (do_vmclear_operation)
>-		do_vmclear_operation();
>+	callback = rcu_dereference(cpu_emergency_virt_callback);
>+	if (callback)
>+		callback();
> 	rcu_read_unlock();
> }
> #endif
>diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
>index 302086255be6..41095fc864f3 100644
>--- a/arch/x86/kvm/vmx/vmx.c
>+++ b/arch/x86/kvm/vmx/vmx.c
>@@ -8551,8 +8551,7 @@ static void __vmx_exit(void)
> {
> 	allow_smaller_maxphyaddr = false;
> 
>-	RCU_INIT_POINTER(crash_vmclear_loaded_vmcss, NULL);
>-	synchronize_rcu();
>+	cpu_emergency_unregister_virt_callback(crash_vmclear_local_loaded_vmcss);
> 
> 	vmx_cleanup_l1d_flush();
> }
>@@ -8602,8 +8601,7 @@ static int __init vmx_init(void)
> 		pi_init_cpu(cpu);
> 	}
> 
>-	rcu_assign_pointer(crash_vmclear_loaded_vmcss,
>-			   crash_vmclear_local_loaded_vmcss);
>+	cpu_emergency_register_virt_callback(crash_vmclear_local_loaded_vmcss);
> 
> 	vmx_check_vmcs12_offsets();
> 
>-- 
>2.40.0.rc1.284.g88254d51c5-goog
>

On Mon, Mar 13, 2023, Chao Gao wrote:
> On Fri, Mar 10, 2023 at 01:42:17PM -0800, Sean Christopherson wrote:
> >+void cpu_emergency_register_virt_callback(cpu_emergency_virt_cb *callback)
> >+{
> >+	if (WARN_ON_ONCE(rcu_access_pointer(cpu_emergency_virt_callback)))
> >+		return;
> >+
> >+	rcu_assign_pointer(cpu_emergency_virt_callback, callback);
> 
> Was it intentional to not call synchronize_rcu() (in the original
> code), different from the un-registration path?

Yes, synchronize_rcu() is needed when removing a callback to ensure any in-flight
invocations of the callback complete before KVM is unloaded, i.e. to prevent
use-after-free of the KVM module code.  Registering a callback doesn't have the
same concerns, and adding a synchronize_rcu() wouldn't do anything in terms of
ensuring virtualization isn't enabled after an emergency restart/shutdown is
initiated.