[v8,13/16] x86/virt/tdx: Configure global KeyID on all packages

Commit Message

Huang, Kai Dec. 9, 2022, 6:52 a.m. UTC

After the list of TDMRs and the global KeyID are configured to the TDX
module, the kernel needs to configure the key of the global KeyID on all
packages using TDH.SYS.KEY.CONFIG.

TDH.SYS.KEY.CONFIG needs to be done on one (any) cpu for each package.
Also, it cannot run concurrently on different cpus, so just use
smp_call_function_single() to do it one by one.

Note to keep things simple, neither the function to configure the global
KeyID on all packages nor the tdx_enable() checks whether there's at
least one online cpu for each package.  Also, neither of them explicitly
prevents any cpu from going offline.  It is caller's responsibility to
guarantee this.

Intel hardware doesn't guarantee cache coherency across different
KeyIDs.  The kernel needs to flush PAMT's dirty cachelines (associated
with KeyID 0) before the TDX module uses the global KeyID to access the
PAMT.  Otherwise, those dirty cachelines can silently corrupt the TDX
module's metadata.  Note this breaks TDX from functionality point of
view but TDX's security remains intact.

Following the TDX module specification, flush cache before configuring
the global KeyID on all packages.  Given the PAMT size can be large
(~1/256th of system RAM), just use WBINVD on all CPUs to flush.

Note if any TDH.SYS.KEY.CONFIG fails, the TDX module may already have
used the global KeyID to write any PAMT.  Therefore, need to use WBINVD
to flush cache before freeing the PAMTs back to the kernel.

Reviewed-by: Isaku Yamahata <[e-mail]>
Signed-off-by: Kai Huang <[e-mail]>
---

v7 -> v8: (Dave)
 - Changelog changes:
  - Point out this is the step of "multi-steps" of init_tdx_module().
  - Removed MOVDIR64B part.
  - Other changes due to removing TDH.SYS.SHUTDOWN and TDH.SYS.LP.INIT.
 - Changed to loop over online cpus and use smp_call_function_single()
   directly as the patch to shut down TDX module has been removed.
 - Removed MOVDIR64B part in comment.

v6 -> v7:
 - Improved changelong and comment to explain why MOVDIR64B isn't used
   when returning PAMTs back to the kernel.

---
 arch/x86/virt/vmx/tdx/tdx.c | 97 +++++++++++++++++++++++++++++++++++--
 arch/x86/virt/vmx/tdx/tdx.h |  1 +
 2 files changed, 94 insertions(+), 4 deletions(-)

Comments

Dave Hansen Jan. 6, 2023, 10:49 p.m. UTC | #1

On 12/8/22 22:52, Kai Huang wrote:
> After the list of TDMRs and the global KeyID are configured to the TDX
> module, the kernel needs to configure the key of the global KeyID on all
> packages using TDH.SYS.KEY.CONFIG.
> 
> TDH.SYS.KEY.CONFIG needs to be done on one (any) cpu for each package.
> Also, it cannot run concurrently on different cpus, so just use
> smp_call_function_single() to do it one by one.
> 
> Note to keep things simple, neither the function to configure the global
> KeyID on all packages nor the tdx_enable() checks whether there's at
> least one online cpu for each package.  Also, neither of them explicitly
> prevents any cpu from going offline.  It is caller's responsibility to
> guarantee this.

OK, but does someone *actually* do this?

> Intel hardware doesn't guarantee cache coherency across different
> KeyIDs.  The kernel needs to flush PAMT's dirty cachelines (associated
> with KeyID 0) before the TDX module uses the global KeyID to access the
> PAMT.  Otherwise, those dirty cachelines can silently corrupt the TDX
> module's metadata.  Note this breaks TDX from functionality point of
> view but TDX's security remains intact.

	Intel hardware doesn't guarantee cache coherency across
	different KeyIDs.  The PAMTs are transitioning from being used
	by the kernel mapping (KeyId 0) to the TDX module's "global
	KeyID" mapping.

	This means that the kernel must flush any dirty KeyID-0 PAMT
	cachelines before the TDX module uses the global KeyID to access
	the PAMT.  Otherwise, if those dirty cachelines were written
	back, they would corrupt the TDX module's metadata.  Aside: This
	corruption would be detected by the memory integrity hardware on
	the next read of the memory with the global KeyID.  The result
	would likely be fatal to the system but would not impact TDX
	security.

> Following the TDX module specification, flush cache before configuring
> the global KeyID on all packages.  Given the PAMT size can be large
> (~1/256th of system RAM), just use WBINVD on all CPUs to flush.
> 
> Note if any TDH.SYS.KEY.CONFIG fails, the TDX module may already have
> used the global KeyID to write any PAMT.  Therefore, need to use WBINVD
> to flush cache before freeing the PAMTs back to the kernel.

						s/need to// ^


> diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
> index ab961443fed5..4c779e8412f1 100644
> --- a/arch/x86/virt/vmx/tdx/tdx.c
> +++ b/arch/x86/virt/vmx/tdx/tdx.c
> @@ -946,6 +946,66 @@ static int config_tdx_module(struct tdmr_info_list *tdmr_list, u64 global_keyid)
>  	return ret;
>  }
>  
> +static void do_global_key_config(void *data)
> +{
> +	int ret;
> +
> +	/*
> +	 * TDH.SYS.KEY.CONFIG may fail with entropy error (which is a
> +	 * recoverable error).  Assume this is exceedingly rare and
> +	 * just return error if encountered instead of retrying.
> +	 */
> +	ret = seamcall(TDH_SYS_KEY_CONFIG, 0, 0, 0, 0, NULL, NULL);
> +
> +	*(int *)data = ret;
> +}
> +
> +/*
> + * Configure the global KeyID on all packages by doing TDH.SYS.KEY.CONFIG
> + * on one online cpu for each package.  If any package doesn't have any
> + * online

This looks like it stopped mid-sentence.

> + * Note:
> + *
> + * This function neither checks whether there's at least one online cpu
> + * for each package, nor explicitly prevents any cpu from going offline.
> + * If any package doesn't have any online cpu then the SEAMCALL won't be
> + * done on that package and the later step of TDX module initialization
> + * will fail.  The caller needs to guarantee this.
> + */

*Does* the caller guarantee it?

You're basically saying, "this code needs $FOO to work", but you're not
saying who *provides* $FOO.

> +static int config_global_keyid(void)
> +{
> +	cpumask_var_t packages;
> +	int cpu, ret = 0;
> +
> +	if (!zalloc_cpumask_var(&packages, GFP_KERNEL))
> +		return -ENOMEM;
> +
> +	for_each_online_cpu(cpu) {
> +		int err;
> +
> +		if (cpumask_test_and_set_cpu(topology_physical_package_id(cpu),
> +					packages))
> +			continue;
> +
> +		/*
> +		 * TDH.SYS.KEY.CONFIG cannot run concurrently on
> +		 * different cpus, so just do it one by one.
> +		 */
> +		ret = smp_call_function_single(cpu, do_global_key_config, &err,
> +				true);
> +		if (ret)
> +			break;
> +		if (err) {
> +			ret = err;
> +			break;
> +		}
> +	}
> +
> +	free_cpumask_var(packages);
> +	return ret;
> +}
> +
>  static int init_tdx_module(void)
>  {
>  	/*
> @@ -998,19 +1058,46 @@ static int init_tdx_module(void)
>  	if (ret)
>  		goto out_free_pamts;
>  
> +	/*
> +	 * Hardware doesn't guarantee cache coherency across different
> +	 * KeyIDs.  The kernel needs to flush PAMT's dirty cachelines
> +	 * (associated with KeyID 0) before the TDX module can use the
> +	 * global KeyID to access the PAMT.  Given PAMTs are potentially
> +	 * large (~1/256th of system RAM), just use WBINVD on all cpus
> +	 * to flush the cache.
> +	 *
> +	 * Follow the TDX spec to flush cache before configuring the
> +	 * global KeyID on all packages.
> +	 */

I don't think this second paragraph adds very much clarity.

Huang, Kai Jan. 10, 2023, 10:15 a.m. UTC | #2

On Fri, 2023-01-06 at 14:49 -0800, Dave Hansen wrote:
> On 12/8/22 22:52, Kai Huang wrote:
> > After the list of TDMRs and the global KeyID are configured to the TDX
> > module, the kernel needs to configure the key of the global KeyID on all
> > packages using TDH.SYS.KEY.CONFIG.
> > 
> > TDH.SYS.KEY.CONFIG needs to be done on one (any) cpu for each package.
> > Also, it cannot run concurrently on different cpus, so just use
> > smp_call_function_single() to do it one by one.
> > 
> > Note to keep things simple, neither the function to configure the global
> > KeyID on all packages nor the tdx_enable() checks whether there's at
> > least one online cpu for each package.  Also, neither of them explicitly
> > prevents any cpu from going offline.  It is caller's responsibility to
> > guarantee this.
> 
> OK, but does someone *actually* do this?

Please see below reply around the code.

> 
> > Intel hardware doesn't guarantee cache coherency across different
> > KeyIDs.  The kernel needs to flush PAMT's dirty cachelines (associated
> > with KeyID 0) before the TDX module uses the global KeyID to access the
> > PAMT.  Otherwise, those dirty cachelines can silently corrupt the TDX
> > module's metadata.  Note this breaks TDX from functionality point of
> > view but TDX's security remains intact.
> 
> 	Intel hardware doesn't guarantee cache coherency across
> 	different KeyIDs.  The PAMTs are transitioning from being used
> 	by the kernel mapping (KeyId 0) to the TDX module's "global
> 	KeyID" mapping.
> 
> 	This means that the kernel must flush any dirty KeyID-0 PAMT
> 	cachelines before the TDX module uses the global KeyID to access
> 	the PAMT.  Otherwise, if those dirty cachelines were written
> 	back, they would corrupt the TDX module's metadata.  Aside: This
> 	corruption would be detected by the memory integrity hardware on
> 	the next read of the memory with the global KeyID.  The result
> 	would likely be fatal to the system but would not impact TDX
> 	security.

Thanks!

> 
> > Following the TDX module specification, flush cache before configuring
> > the global KeyID on all packages.  Given the PAMT size can be large
> > (~1/256th of system RAM), just use WBINVD on all CPUs to flush.
> > 
> > Note if any TDH.SYS.KEY.CONFIG fails, the TDX module may already have
> > used the global KeyID to write any PAMT.  Therefore, need to use WBINVD
> > to flush cache before freeing the PAMTs back to the kernel.
> 
> 						s/need to// ^

Will do.  Thanks.

> 
> 
> > diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
> > index ab961443fed5..4c779e8412f1 100644
> > --- a/arch/x86/virt/vmx/tdx/tdx.c
> > +++ b/arch/x86/virt/vmx/tdx/tdx.c
> > @@ -946,6 +946,66 @@ static int config_tdx_module(struct tdmr_info_list *tdmr_list, u64 global_keyid)
> >  	return ret;
> >  }
> >  
> > +static void do_global_key_config(void *data)
> > +{
> > +	int ret;
> > +
> > +	/*
> > +	 * TDH.SYS.KEY.CONFIG may fail with entropy error (which is a
> > +	 * recoverable error).  Assume this is exceedingly rare and
> > +	 * just return error if encountered instead of retrying.
> > +	 */
> > +	ret = seamcall(TDH_SYS_KEY_CONFIG, 0, 0, 0, 0, NULL, NULL);
> > +
> > +	*(int *)data = ret;
> > +}
> > +
> > +/*
> > + * Configure the global KeyID on all packages by doing TDH.SYS.KEY.CONFIG
> > + * on one online cpu for each package.  If any package doesn't have any
> > + * online
> 
> This looks like it stopped mid-sentence.

Oops I forgot to delete the broken sentence.

> 
> > + * Note:
> > + *
> > + * This function neither checks whether there's at least one online cpu
> > + * for each package, nor explicitly prevents any cpu from going offline.
> > + * If any package doesn't have any online cpu then the SEAMCALL won't be
> > + * done on that package and the later step of TDX module initialization
> > + * will fail.  The caller needs to guarantee this.
> > + */
> 
> *Does* the caller guarantee it?
> 
> You're basically saying, "this code needs $FOO to work", but you're not
> saying who *provides* $FOO.

In short, KVM can do something to guarantee but won't 100% guarantee this.

Specifically, KVM won't actively try to bring up cpu to guarantee this if
there's any package has no online cpu at all (see the first lore link below). 
But KVM can _check_ whether this condition has been met before calling
tdx_init() and speak out if not.  At the meantime, if the condition is met,
refuse to offline the last cpu for each package (or any cpu) during module
initialization.

And KVM needs similar handling anyway.  The reason is not only configuring the
global KeyID has such requirement, creating/destroying TD (which involves
programming/reclaiming one TDX KeyID) also require at least one online cpu for
each package.  

There were discussions around this on KVM how to handle.  IIUC the solution is
KVM will:
1) fail to create TD if any package has no online cpu.
2) refuse to offline the last cpu for each package when there's any _active_ TDX
guest running.

https://lore.kernel.org/lkml/[e-mail]/T/#m1ff338686cfcb7ba691cd969acc17b32ff194073
https://lore.kernel.org/lkml/de6b69[e-mail]/

Thus TDX module initialization in KVM can be handled in similar way.

Btw, in v7 (which has per-lp init requirement on all cpus), tdx_init() does
early check on whether all machine boot-time present cpu are online and simply 
returns error if condition is not met.  Here the difference is we don't have any
check but depend on SEAMCALL to fail.  To me there's no fundamental difference.

[snip]

> 
> >  static int init_tdx_module(void)
> >  {
> >  	/*
> > @@ -998,19 +1058,46 @@ static int init_tdx_module(void)
> >  	if (ret)
> >  		goto out_free_pamts;
> >  
> > +	/*
> > +	 * Hardware doesn't guarantee cache coherency across different
> > +	 * KeyIDs.  The kernel needs to flush PAMT's dirty cachelines
> > +	 * (associated with KeyID 0) before the TDX module can use the
> > +	 * global KeyID to access the PAMT.  Given PAMTs are potentially
> > +	 * large (~1/256th of system RAM), just use WBINVD on all cpus
> > +	 * to flush the cache.
> > +	 *
> > +	 * Follow the TDX spec to flush cache before configuring the
> > +	 * global KeyID on all packages.
> > +	 */
> 
> I don't think this second paragraph adds very much clarity.
> 
> 

OK will remove.

Dave Hansen Jan. 10, 2023, 4:53 p.m. UTC | #3

On 1/10/23 02:15, Huang, Kai wrote:
> On Fri, 2023-01-06 at 14:49 -0800, Dave Hansen wrote:
>> On 12/8/22 22:52, Kai Huang wrote:
...
>>> + * Note:
>>> + *
>>> + * This function neither checks whether there's at least one online cpu
>>> + * for each package, nor explicitly prevents any cpu from going offline.
>>> + * If any package doesn't have any online cpu then the SEAMCALL won't be
>>> + * done on that package and the later step of TDX module initialization
>>> + * will fail.  The caller needs to guarantee this.
>>> + */
>>
>> *Does* the caller guarantee it?
>>
>> You're basically saying, "this code needs $FOO to work", but you're not
>> saying who *provides* $FOO.
> 
> In short, KVM can do something to guarantee but won't 100% guarantee this.
> 
> Specifically, KVM won't actively try to bring up cpu to guarantee this if
> there's any package has no online cpu at all (see the first lore link below).
> But KVM can _check_ whether this condition has been met before calling
> tdx_init() and speak out if not.  At the meantime, if the condition is met,
> refuse to offline the last cpu for each package (or any cpu) during module
> initialization.
> 
> And KVM needs similar handling anyway.  The reason is not only configuring the
> global KeyID has such requirement, creating/destroying TD (which involves
> programming/reclaiming one TDX KeyID) also require at least one online cpu for
> each package.
> 
> There were discussions around this on KVM how to handle.  IIUC the solution is
> KVM will:
> 1) fail to create TD if any package has no online cpu.
> 2) refuse to offline the last cpu for each package when there's any _active_ TDX
> guest running.
> 
> https://lore.kernel.org/lkml/[e-mail]/T/#m1ff338686cfcb7ba691cd969acc17b32ff194073
> https://lore.kernel.org/lkml/de6b69[e-mail]/
> 
> Thus TDX module initialization in KVM can be handled in similar way.
> 
> Btw, in v7 (which has per-lp init requirement on all cpus), tdx_init() does
> early check on whether all machine boot-time present cpu are online and simply
> returns error if condition is not met.  Here the difference is we don't have any
> check but depend on SEAMCALL to fail.  To me there's no fundamental difference.

So, I'm going to call shenanigans here.

You say:

	The caller needs to guarantee this.

Then, you go and tell us how the *ONE* caller of this function doesn't
actually guarantee this.  Plus, you *KNOW* this.

Those are shenanigans.

Let's do something like this instead of asking for something impossible
and pretending that the callers are going to provide some fantasy solution.

/*
 * Attempt to configure the global KeyID on all physical packages.
 *
 * This requires running code on at least one CPU in each package.  If a
 * package has no online CPUs, that code will not run and TDX module
 * initialization (TDH.whatever) will fail.
 *
 * This code takes no affirmative steps to online CPUs.  Callers (aka.
 * KVM) can ensure success by ensuring sufficient CPUs are online for
 * this to succeed.
 */

Now, since this _is_ all imperfect, what will our users see if this
house of cards falls down?  Will they get a nice error message like:

     TDX: failed to configure module, no online CPUs in package 12

Or, will they see:

     TDX: Hurr, durr, I'm confused and you should be too

?

Huang, Kai Jan. 11, 2023, 12:06 a.m. UTC | #4

On Tue, 2023-01-10 at 08:53 -0800, Dave Hansen wrote:
> On 1/10/23 02:15, Huang, Kai wrote:
> > On Fri, 2023-01-06 at 14:49 -0800, Dave Hansen wrote:
> > > On 12/8/22 22:52, Kai Huang wrote:
> ...
> > > > + * Note:
> > > > + *
> > > > + * This function neither checks whether there's at least one online cpu
> > > > + * for each package, nor explicitly prevents any cpu from going offline.
> > > > + * If any package doesn't have any online cpu then the SEAMCALL won't be
> > > > + * done on that package and the later step of TDX module initialization
> > > > + * will fail.  The caller needs to guarantee this.
> > > > + */
> > > 
> > > *Does* the caller guarantee it?
> > > 
> > > You're basically saying, "this code needs $FOO to work", but you're not
> > > saying who *provides* $FOO.
> > 
> > In short, KVM can do something to guarantee but won't 100% guarantee this.
> > 
> > Specifically, KVM won't actively try to bring up cpu to guarantee this if
> > there's any package has no online cpu at all (see the first lore link below).
> > But KVM can _check_ whether this condition has been met before calling
> > tdx_init() and speak out if not.  At the meantime, if the condition is met,
> > refuse to offline the last cpu for each package (or any cpu) during module
> > initialization.
> > 
> > And KVM needs similar handling anyway.  The reason is not only configuring the
> > global KeyID has such requirement, creating/destroying TD (which involves
> > programming/reclaiming one TDX KeyID) also require at least one online cpu for
> > each package.
> > 
> > There were discussions around this on KVM how to handle.  IIUC the solution is
> > KVM will:
> > 1) fail to create TD if any package has no online cpu.
> > 2) refuse to offline the last cpu for each package when there's any _active_ TDX
> > guest running.
> > 
> > https://lore.kernel.org/lkml/[e-mail]/T/#m1ff338686cfcb7ba691cd969acc17b32ff194073
> > https://lore.kernel.org/lkml/de6b69[e-mail]/
> > 
> > Thus TDX module initialization in KVM can be handled in similar way.
> > 
> > Btw, in v7 (which has per-lp init requirement on all cpus), tdx_init() does
> > early check on whether all machine boot-time present cpu are online and simply
> > returns error if condition is not met.  Here the difference is we don't have any
> > check but depend on SEAMCALL to fail.  To me there's no fundamental difference.
> 
> So, I'm going to call shenanigans here.
> 
> You say:
> 
> 	The caller needs to guarantee this.
> 
> Then, you go and tell us how the *ONE* caller of this function doesn't
> actually guarantee this.  Plus, you *KNOW* this.
> 
> Those are shenanigans.

Agreed.

> 
> Let's do something like this instead of asking for something impossible
> and pretending that the callers are going to provide some fantasy solution.
> 
> /*
>  * Attempt to configure the global KeyID on all physical packages.
>  *
>  * This requires running code on at least one CPU in each package.  If a
>  * package has no online CPUs, that code will not run and TDX module
>  * initialization (TDH.whatever) will fail.
>  *
>  * This code takes no affirmative steps to online CPUs.  Callers (aka.
>  * KVM) can ensure success by ensuring sufficient CPUs are online for
>  * this to succeed.
>  */

Thanks.  Will update changelog accordingly.

> 
> Now, since this _is_ all imperfect, what will our users see if this
> house of cards falls down?  Will they get a nice error message like:
> 
>      TDX: failed to configure module, no online CPUs in package 12
> 
> Or, will they see:
> 
>      TDX: Hurr, durr, I'm confused and you should be too
> 
> ?

I am expecting the former.  I will work with Isaku to make sure of it.

Patch

diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index ab961443fed5..4c779e8412f1 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -946,6 +946,66 @@  static int config_tdx_module(struct tdmr_info_list *tdmr_list, u64 global_keyid)
 	return ret;
 }
 
+static void do_global_key_config(void *data)
+{
+	int ret;
+
+	/*
+	 * TDH.SYS.KEY.CONFIG may fail with entropy error (which is a
+	 * recoverable error).  Assume this is exceedingly rare and
+	 * just return error if encountered instead of retrying.
+	 */
+	ret = seamcall(TDH_SYS_KEY_CONFIG, 0, 0, 0, 0, NULL, NULL);
+
+	*(int *)data = ret;
+}
+
+/*
+ * Configure the global KeyID on all packages by doing TDH.SYS.KEY.CONFIG
+ * on one online cpu for each package.  If any package doesn't have any
+ * online
+ *
+ * Note:
+ *
+ * This function neither checks whether there's at least one online cpu
+ * for each package, nor explicitly prevents any cpu from going offline.
+ * If any package doesn't have any online cpu then the SEAMCALL won't be
+ * done on that package and the later step of TDX module initialization
+ * will fail.  The caller needs to guarantee this.
+ */
+static int config_global_keyid(void)
+{
+	cpumask_var_t packages;
+	int cpu, ret = 0;
+
+	if (!zalloc_cpumask_var(&packages, GFP_KERNEL))
+		return -ENOMEM;
+
+	for_each_online_cpu(cpu) {
+		int err;
+
+		if (cpumask_test_and_set_cpu(topology_physical_package_id(cpu),
+					packages))
+			continue;
+
+		/*
+		 * TDH.SYS.KEY.CONFIG cannot run concurrently on
+		 * different cpus, so just do it one by one.
+		 */
+		ret = smp_call_function_single(cpu, do_global_key_config, &err,
+				true);
+		if (ret)
+			break;
+		if (err) {
+			ret = err;
+			break;
+		}
+	}
+
+	free_cpumask_var(packages);
+	return ret;
+}
+
 static int init_tdx_module(void)
 {
 	/*
@@ -998,19 +1058,46 @@  static int init_tdx_module(void)
 	if (ret)
 		goto out_free_pamts;
 
+	/*
+	 * Hardware doesn't guarantee cache coherency across different
+	 * KeyIDs.  The kernel needs to flush PAMT's dirty cachelines
+	 * (associated with KeyID 0) before the TDX module can use the
+	 * global KeyID to access the PAMT.  Given PAMTs are potentially
+	 * large (~1/256th of system RAM), just use WBINVD on all cpus
+	 * to flush the cache.
+	 *
+	 * Follow the TDX spec to flush cache before configuring the
+	 * global KeyID on all packages.
+	 */
+	wbinvd_on_all_cpus();
+
+	/* Config the key of global KeyID on all packages */
+	ret = config_global_keyid();
+	if (ret)
+		goto out_free_pamts;
+
 	/*
 	 * TODO:
 	 *
-	 *  - Configure the global KeyID on all packages.
 	 *  - Initialize all TDMRs.
 	 *
 	 *  Return error before all steps are done.
 	 */
 	ret = -EINVAL;
 out_free_pamts:
-	if (ret)
+	if (ret) {
+		/*
+		 * Part of PAMT may already have been initialized by the
+		 * TDX module.  Flush cache before returning PAMT back
+		 * to the kernel.
+		 *
+		 * No need to worry about integrity checks here.  KeyID
+		 * 0 has integrity checking disabled.
+		 */
+		wbinvd_on_all_cpus();
+
 		tdmrs_free_pamt_all(&tdmr_list);
-	else
+	} else
 		pr_info("%lu pages allocated for PAMT.\n",
 				tdmrs_count_pamt_pages(&tdmr_list));
 out_free_tdmrs:
@@ -1057,7 +1144,9 @@  static int __tdx_enable(void)
  * tdx_enable - Enable TDX by initializing the TDX module
  *
  * The caller must make sure all online cpus are in VMX operation before
- * calling this function.
+ * calling this function.  Also, the caller must make sure there is at
+ * least one online cpu for each package, and to prevent any cpu from
+ * going offline during this function.
  *
  * This function can be called in parallel by multiple callers.
  *
diff --git a/arch/x86/virt/vmx/tdx/tdx.h b/arch/x86/virt/vmx/tdx/tdx.h
index 4d2edd477480..f5c12a2543d4 100644
--- a/arch/x86/virt/vmx/tdx/tdx.h
+++ b/arch/x86/virt/vmx/tdx/tdx.h
@@ -19,6 +19,7 @@ 
 /*
  * TDX module SEAMCALL leaf functions
  */
+#define TDH_SYS_KEY_CONFIG	31
 #define TDH_SYS_INFO		32
 #define TDH_SYS_CONFIG		45