Re: [PATCH v3 6/6] KVM: TDX: Explicitly do WBINVD upon reboot notifier

["Edgecombe, Rick P" <rick.p.edgecombe@intel.com>]

On Thu, 2025-06-26 at 22:48 +1200, Kai Huang wrote:
> On TDX platforms, during kexec, the kernel needs to make sure there's no
> dirty cachelines of TDX private memory before booting to the new kernel
> to avoid silent memory corruption to the new kernel.
> 
> During kexec, the kexec-ing CPU firstly invokes native_stop_other_cpus()
> to stop all remote CPUs before booting to the new kernel.  The remote
> CPUs will then execute stop_this_cpu() to stop themselves.
> 
> The kernel has a percpu boolean to indicate whether the cache of a CPU
> may be in incoherent state.  In stop_this_cpu(), the kernel does WBINVD
> if that percpu boolean is true.
> 
> TDX turns on that percpu boolean on a CPU when the kernel does SEAMCALL.
> This makes sure the caches will be flushed during kexec.
> 
> However, the native_stop_other_cpus() and stop_this_cpu() have a "race"
> which is extremely rare to happen but could cause system to hang.
> 
> Specifically, the native_stop_other_cpus() firstly sends normal reboot
> IPI to remote CPUs and wait one second for them to stop.  If that times
> out, native_stop_other_cpus() then sends NMIs to remote CPUs to stop
> them.
> 
> The aforementioned race happens when NMIs are sent.  Doing WBINVD in
> stop_this_cpu() makes each CPU take longer time to stop and increases
> the chance of the race to happen.
> 
> Register reboot notifier in KVM to explicitly flush caches upon
> receiving reboot notifier (e.g., during kexec) for TDX.  This moves the
> WBINVD to an earlier stage than stop_this_cpus(), avoiding a possibly
> lengthy operation at a time where it could cause this race.
> 

I think this will reduce the chance of the race, but it feels like the wrong way
to address the race. But I don't know how to properly fix it either. Maybe this
is just the nature of x86 NMIs, to have code like this.

Functionally it looks good, but a few nits to take or leave below.

> Signed-off-by: Kai Huang <kai.huang@intel.com>
> Acked-by: Paolo Bonzini <pbonzini@redhat.com>
> Tested-by: Farrah Chen <farrah.chen@intel.com>
> ---
> 
> v2 -> v3:
>  - Update changelog to address Paolo's comments and Add Paolo's Ack:
>    https://lore.kernel.org/lkml/3a7c0856-6e7b-4d3d-b966-6f17f1aca42e@redhat.com/
> 
> ---
>  arch/x86/include/asm/tdx.h  |  3 +++
>  arch/x86/kvm/vmx/tdx.c      | 45 +++++++++++++++++++++++++++++++++++++
>  arch/x86/virt/vmx/tdx/tdx.c |  9 ++++++++
>  3 files changed, 57 insertions(+)
> 
> diff --git a/arch/x86/include/asm/tdx.h b/arch/x86/include/asm/tdx.h
> index d4c624c69d7f..e6b11982c6c6 100644
> --- a/arch/x86/include/asm/tdx.h
> +++ b/arch/x86/include/asm/tdx.h
> @@ -221,6 +221,8 @@ u64 tdh_mem_page_remove(struct tdx_td *td, u64 gpa, u64 level, u64 *ext_err1, u6
>  u64 tdh_phymem_cache_wb(bool resume);
>  u64 tdh_phymem_page_wbinvd_tdr(struct tdx_td *td);
>  u64 tdh_phymem_page_wbinvd_hkid(u64 hkid, struct page *page);
> +

Nit: There is a new line here, but not below. I guess it's ok.

> +void tdx_cpu_flush_cache(void);
>  #else
>  static inline void tdx_init(void) { }
>  static inline int tdx_cpu_enable(void) { return -ENODEV; }
> @@ -228,6 +230,7 @@ static inline int tdx_enable(void)  { return -ENODEV; }
>  static inline u32 tdx_get_nr_guest_keyids(void) { return 0; }
>  static inline const char *tdx_dump_mce_info(struct mce *m) { return NULL; }
>  static inline const struct tdx_sys_info *tdx_get_sysinfo(void) { return NULL; }
> +static inline void tdx_cpu_flush_cache(void) { }
>  #endif	/* CONFIG_INTEL_TDX_HOST */
>  
>  #endif /* !__ASSEMBLER__ */
> diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
> index 1ad20c273f3b..c567a64a6cb0 100644
> --- a/arch/x86/kvm/vmx/tdx.c
> +++ b/arch/x86/kvm/vmx/tdx.c
> @@ -5,7 +5,9 @@
>  #include <asm/fpu/xcr.h>
>  #include <linux/misc_cgroup.h>
>  #include <linux/mmu_context.h>
> +#include <linux/reboot.h>
>  #include <asm/tdx.h>
> +#include <asm/processor.h>
>  #include "capabilities.h"
>  #include "mmu.h"
>  #include "x86_ops.h"
> @@ -3347,6 +3349,33 @@ static int tdx_offline_cpu(unsigned int cpu)
>  	return -EBUSY;
>  }
>  
> +static void smp_func_cpu_flush_cache(void *unused)
> +{
> +	tdx_cpu_flush_cache();
> +}
> +
> +static int tdx_reboot_notify(struct notifier_block *nb, unsigned long code,
> +			     void *unused)
> +{
> +	/*
> +	 * Flush cache for all CPUs upon the reboot notifier.  This
> +	 * avoids having to do WBINVD in stop_this_cpu() during kexec.
> +	 *
> +	 * Kexec calls native_stop_other_cpus() to stop remote CPUs
> +	 * before booting to new kernel, but that code has a "race"
> +	 * when the normal REBOOT IPI timesout and NMIs are sent to
> +	 * remote CPUs to stop them.  Doing WBINVD in stop_this_cpu()
> +	 * could potentially increase the posibility of the "race".
> +	 */
> +	if (code == SYS_RESTART)
> +		on_each_cpu(smp_func_cpu_flush_cache, NULL, 1);
> +	return NOTIFY_DONE;
> +}
> +
> +static struct notifier_block tdx_reboot_nb = {
> +	.notifier_call = tdx_reboot_notify,
> +};
> +
>  static void __do_tdx_cleanup(void)
>  {
>  	/*
> @@ -3504,6 +3533,11 @@ void tdx_cleanup(void)
>  {
>  	if (enable_tdx) {
>  		misc_cg_set_capacity(MISC_CG_RES_TDX, 0);
> +		/*
> +		 * Ignore the return value.  See the comment in
> +		 * tdx_bringup().
> +		 */
> +		unregister_reboot_notifier(&tdx_reboot_nb);
>  		__tdx_cleanup();
>  		kvm_disable_virtualization();
>  	}
> @@ -3587,6 +3621,17 @@ int __init tdx_bringup(void)
>  		enable_tdx = 0;
>  	}
>  
> +	if (enable_tdx)
> +		/*
> +		 * Ignore the return value.  @tdx_reboot_nb is used to flush
> +		 * cache for all CPUs upon rebooting to avoid having to do
> +		 * WBINVD in kexec while the kexec-ing CPU stops all remote
> +		 * CPUs.  Failure to register isn't fatal, because if KVM
> +		 * doesn't flush cache explicitly upon rebooting the kexec
> +		 * will do it anyway.
> +		 */
> +		register_reboot_notifier(&tdx_reboot_nb);
> +

The comment should be inside a {}.

>  	return r;
>  
>  success_disable_tdx:
> diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
> index c7a9a087ccaf..73425e9bee39 100644
> --- a/arch/x86/virt/vmx/tdx/tdx.c
> +++ b/arch/x86/virt/vmx/tdx/tdx.c
> @@ -1870,3 +1870,12 @@ u64 tdh_phymem_page_wbinvd_hkid(u64 hkid, struct page *page)
>  	return seamcall(TDH_PHYMEM_PAGE_WBINVD, &args);
>  }
>  EXPORT_SYMBOL_GPL(tdh_phymem_page_wbinvd_hkid);
> +
> +void tdx_cpu_flush_cache(void)
> +{
> +	lockdep_assert_preemption_disabled();
> +
> +	wbinvd();
> +	this_cpu_write(cache_state_incoherent, false);
> +}
> +EXPORT_SYMBOL_GPL(tdx_cpu_flush_cache);

Does this need to be here? Why not in KVM?