diff mbox series

[v11,15/20] x86/virt/tdx: Configure global KeyID on all packages

Message ID 30358db4eff961c69783bbd4d9f3e50932a9a759.1685887183.git.kai.huang@intel.com (mailing list archive)
State New, archived
Headers show
Series TDX host kernel support | expand

Commit Message

Huang, Kai June 4, 2023, 2:27 p.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.

This SEAMCALL cannot run parallel on different cpus.  Loop all online
cpus and use smp_call_on_cpu() to call this SEAMCALL on the first cpu of
each package.

To keep things simple, this implementation takes no affirmative steps to
online cpus to make sure there's at least one cpu for each package.  The
callers (aka. KVM) can ensure success by ensuring that.

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 PAMTs.
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.

If TDH.SYS.KEY.CONFIG fails, the TDX module may already have used the
global KeyID to write the PAMTs.  Therefore, use WBINVD to flush cache
before returning the PAMTs back to the kernel.  Also convert all PAMTs
back to normal by using MOVDIR64B as suggested by the TDX module spec,
although on the platform without the "partial write machine check"
erratum it's OK to leave PAMTs as is.

Signed-off-by: Kai Huang <kai.huang@intel.com>
Reviewed-by: Isaku Yamahata <isaku.yamahata@intel.com>
---

v10 -> v11:
 - Convert PAMTs back to normal when module initialization fails.
 - Fixed an error in changelog

v9 -> v10:
 - Changed to use 'smp_call_on_cpu()' directly to do key configuration.

v8 -> v9:
 - Improved changelog (Dave).
 - Improved comments to explain the function to configure global KeyID
   "takes no affirmative action to online any cpu". (Dave).
 - Improved other comments suggested by Dave.

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 | 134 +++++++++++++++++++++++++++++++++++-
 arch/x86/virt/vmx/tdx/tdx.h |   1 +
 2 files changed, 132 insertions(+), 3 deletions(-)

Comments

Kirill A . Shutemov June 8, 2023, 11:53 p.m. UTC | #1 
On Mon, Jun 05, 2023 at 02:27:28AM +1200, 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.
> 
> This SEAMCALL cannot run parallel on different cpus.  Loop all online
> cpus and use smp_call_on_cpu() to call this SEAMCALL on the first cpu of
> each package.
> 
> To keep things simple, this implementation takes no affirmative steps to
> online cpus to make sure there's at least one cpu for each package.  The
> callers (aka. KVM) can ensure success by ensuring that.
> 
> 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 PAMTs.
> 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.
> 
> If TDH.SYS.KEY.CONFIG fails, the TDX module may already have used the
> global KeyID to write the PAMTs.  Therefore, use WBINVD to flush cache
> before returning the PAMTs back to the kernel.  Also convert all PAMTs
> back to normal by using MOVDIR64B as suggested by the TDX module spec,
> although on the platform without the "partial write machine check"
> erratum it's OK to leave PAMTs as is.
> 
> Signed-off-by: Kai Huang <kai.huang@intel.com>
> Reviewed-by: Isaku Yamahata <isaku.yamahata@intel.com>

Reviewed-by: Kirill A. Shutemov <kirill.shutemov@linux.intel.com>
Nikolay Borisov June 15, 2023, 8:12 a.m. UTC | #2 
On 4.06.23 г. 17:27 ч., 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.
> 
> This SEAMCALL cannot run parallel on different cpus.  Loop all online
> cpus and use smp_call_on_cpu() to call this SEAMCALL on the first cpu of
> each package.
> 
> To keep things simple, this implementation takes no affirmative steps to
> online cpus to make sure there's at least one cpu for each package.  The
> callers (aka. KVM) can ensure success by ensuring that.

The last sentence is a bit hard to read due to the multiple use of 
ensure/ensuring. OTOH I find the comment in the code somewhat more 
coherent:

 > + * 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.
 > + */

I'd suggest you just use those words. Or just saying "Callers (such as 
KVM) can ensure success by onlining at least 1 CPU per package."

<snip>


>   static int init_tdx_module(void)
>   {
>   	static DECLARE_PADDED_STRUCT(tdsysinfo_struct, tdsysinfo,
> @@ -980,15 +1073,47 @@ 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.
> +	 */
> +	wbinvd_on_all_cpus();
> +
> +	/* Config the key of global KeyID on all packages */
> +	ret = config_global_keyid();
> +	if (ret)
> +		goto out_reset_pamts;
> +
>   	/*
>   	 * TODO:
>   	 *
> -	 *  - Configure the global KeyID on all packages.
>   	 *  - Initialize all TDMRs.
>   	 *
>   	 *  Return error before all steps are done.
>   	 */
>   	ret = -EINVAL;
> +out_reset_pamts:
> +	if (ret) {

Again with those conditionals in the error paths. Just copy the 
put_mem_online(); ret 0. and this will decrease the indentation level 
and make the code linear. Here;s what the diff looks like:


diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 4aa41352edfc..49fda2a28f24 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -1131,6 +1131,8 @@ static int init_tdx_module(void)
         if (ret)
                 goto out_free_pamts;

+       pr_info("%lu KBs allocated for PAMT.\n",
+               tdmrs_count_pamt_pages(&tdx_tdmr_list) * 4);
         /*
          * Hardware doesn't guarantee cache coherency across different
          * KeyIDs.  The kernel needs to flush PAMT's dirty cachelines
@@ -1148,36 +1150,32 @@ static int init_tdx_module(void)

         /* Initialize TDMRs to complete the TDX module initialization */
         ret = init_tdmrs(&tdx_tdmr_list);
+
+       put_online_mems();
+
+       return 0;
  out_reset_pamts:
-       if (ret) {
-               /*
-                * Part of PAMTs may already have been initialized by the
-                * TDX module.  Flush cache before returning PAMTs back
-                * to the kernel.
-                */
-               wbinvd_on_all_cpus();
-               /*
-                * According to the TDX hardware spec, if the platform
-                * doesn't have the "partial write machine check"
-                * erratum, any kernel read/write will never cause #MC
-                * in kernel space, thus it's OK to not convert PAMTs
-                * back to normal.  But do the conversion anyway here
-                * as suggested by the TDX spec.
-                */
-               tdmrs_reset_pamt_all(&tdx_tdmr_list);
-       }
+       /*
+        * Part of PAMTs may already have been initialized by the
+        * TDX module.  Flush cache before returning PAMTs back
+        * to the kernel.
+        */
+       wbinvd_on_all_cpus();
+       /*
+        * According to the TDX hardware spec, if the platform
+        * doesn't have the "partial write machine check"
+        * erratum, any kernel read/write will never cause #MC
+        * in kernel space, thus it's OK to not convert PAMTs
+        * back to normal.  But do the conversion anyway here
+        * as suggested by the TDX spec.
+        */
+       tdmrs_reset_pamt_all(&tdx_tdmr_list);
  out_free_pamts:
-       if (ret)
-               tdmrs_free_pamt_all(&tdx_tdmr_list);
-       else
-               pr_info("%lu KBs allocated for PAMT.\n",
-                               tdmrs_count_pamt_pages(&tdx_tdmr_list) * 4);
+       tdmrs_free_pamt_all(&tdx_tdmr_list);
  out_free_tdmrs:
-       if (ret)
-               free_tdmr_list(&tdx_tdmr_list);
+       free_tdmr_list(&tdx_tdmr_list);
  out_free_tdx_mem:
-       if (ret)
-               free_tdx_memlist(&tdx_memlist);
+       free_tdx_memlist(&tdx_memlist);
  out:
         /*
          * @tdx_memlist is written here and read at memory hotplug time.


<snip>
Huang, Kai June 15, 2023, 10:24 p.m. UTC | #3 
On Thu, 2023-06-15 at 11:12 +0300, Nikolay Borisov wrote:
> 
> On 4.06.23 г. 17:27 ч., 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.
> > 
> > This SEAMCALL cannot run parallel on different cpus.  Loop all online
> > cpus and use smp_call_on_cpu() to call this SEAMCALL on the first cpu of
> > each package.
> > 
> > To keep things simple, this implementation takes no affirmative steps to
> > online cpus to make sure there's at least one cpu for each package.  The
> > callers (aka. KVM) can ensure success by ensuring that.
> 
> The last sentence is a bit hard to read due to the multiple use of 
> ensure/ensuring. OTOH I find the comment in the code somewhat more 
> coherent:
> 
>  > + * 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.
>  > + */
> 
> I'd suggest you just use those words. Or just saying "Callers (such as 
> KVM) can ensure success by onlining at least 1 CPU per package."
> 

OK will do.  Thanks.

> 
> 
> >   static int init_tdx_module(void)
> >   {
> >   	static DECLARE_PADDED_STRUCT(tdsysinfo_struct, tdsysinfo,
> > @@ -980,15 +1073,47 @@ 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.
> > +	 */
> > +	wbinvd_on_all_cpus();
> > +
> > +	/* Config the key of global KeyID on all packages */
> > +	ret = config_global_keyid();
> > +	if (ret)
> > +		goto out_reset_pamts;
> > +
> >   	/*
> >   	 * TODO:
> >   	 *
> > -	 *  - Configure the global KeyID on all packages.
> >   	 *  - Initialize all TDMRs.
> >   	 *
> >   	 *  Return error before all steps are done.
> >   	 */
> >   	ret = -EINVAL;
> > +out_reset_pamts:
> > +	if (ret) {
> 
> Again with those conditionals in the error paths. Just copy the 
> put_mem_online(); ret 0. and this will decrease the indentation level 
> and make the code linear. Here;s what the diff looks like:

(to your another reply too)

I noticed this too when preparing this series.  In old versions the TDMRs were
always freed no matter module initialization result, thus it's not good to do
what you suggested.  Now we can do what you suggested (assuming we don't change
back to always freeing TDMRs), I just wasn't so sure when I was preparing this
series.

I'll take a look at the yielding patches with your suggestion.  Thanks.

Hi Kirill/Dave,

Since I have received couple of tags from you, may I know which way do you
prefer?

> 
> 
> diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
> index 4aa41352edfc..49fda2a28f24 100644
> --- a/arch/x86/virt/vmx/tdx/tdx.c
> +++ b/arch/x86/virt/vmx/tdx/tdx.c
> @@ -1131,6 +1131,8 @@ static int init_tdx_module(void)
>          if (ret)
>                  goto out_free_pamts;
> 
> +       pr_info("%lu KBs allocated for PAMT.\n",
> +               tdmrs_count_pamt_pages(&tdx_tdmr_list) * 4);
>          /*
>           * Hardware doesn't guarantee cache coherency across different
>           * KeyIDs.  The kernel needs to flush PAMT's dirty cachelines
> @@ -1148,36 +1150,32 @@ static int init_tdx_module(void)
> 
>          /* Initialize TDMRs to complete the TDX module initialization */
>          ret = init_tdmrs(&tdx_tdmr_list);
> +
> +       put_online_mems();
> +
> +       return 0;
>   out_reset_pamts:
> -       if (ret) {
> -               /*
> -                * Part of PAMTs may already have been initialized by the
> -                * TDX module.  Flush cache before returning PAMTs back
> -                * to the kernel.
> -                */
> -               wbinvd_on_all_cpus();
> -               /*
> -                * According to the TDX hardware spec, if the platform
> -                * doesn't have the "partial write machine check"
> -                * erratum, any kernel read/write will never cause #MC
> -                * in kernel space, thus it's OK to not convert PAMTs
> -                * back to normal.  But do the conversion anyway here
> -                * as suggested by the TDX spec.
> -                */
> -               tdmrs_reset_pamt_all(&tdx_tdmr_list);
> -       }
> +       /*
> +        * Part of PAMTs may already have been initialized by the
> +        * TDX module.  Flush cache before returning PAMTs back
> +        * to the kernel.
> +        */
> +       wbinvd_on_all_cpus();
> +       /*
> +        * According to the TDX hardware spec, if the platform
> +        * doesn't have the "partial write machine check"
> +        * erratum, any kernel read/write will never cause #MC
> +        * in kernel space, thus it's OK to not convert PAMTs
> +        * back to normal.  But do the conversion anyway here
> +        * as suggested by the TDX spec.
> +        */
> +       tdmrs_reset_pamt_all(&tdx_tdmr_list);
>   out_free_pamts:
> -       if (ret)
> -               tdmrs_free_pamt_all(&tdx_tdmr_list);
> -       else
> -               pr_info("%lu KBs allocated for PAMT.\n",
> -                               tdmrs_count_pamt_pages(&tdx_tdmr_list) * 4);
> +       tdmrs_free_pamt_all(&tdx_tdmr_list);
>   out_free_tdmrs:
> -       if (ret)
> -               free_tdmr_list(&tdx_tdmr_list);
> +       free_tdmr_list(&tdx_tdmr_list);
>   out_free_tdx_mem:
> -       if (ret)
> -               free_tdx_memlist(&tdx_memlist);
> +       free_tdx_memlist(&tdx_memlist);
>   out:
>          /*
>           * @tdx_memlist is written here and read at memory hotplug time.
> 
> 
> <snip>
Kirill A . Shutemov June 19, 2023, 2:56 p.m. UTC | #4 
On Thu, Jun 15, 2023 at 10:24:17PM +0000, Huang, Kai wrote:
> Hi Kirill/Dave,
> 
> Since I have received couple of tags from you, may I know which way do you
> prefer?

I agree with Nikolay, removing these "if (ret)" helps readability.
Huang, Kai June 19, 2023, 11:38 p.m. UTC | #5 
On Mon, 2023-06-19 at 17:56 +0300, kirill.shutemov@linux.intel.com wrote:
> On Thu, Jun 15, 2023 at 10:24:17PM +0000, Huang, Kai wrote:
> > Hi Kirill/Dave,
> > 
> > Since I have received couple of tags from you, may I know which way do you
> > prefer?
> 
> I agree with Nikolay, removing these "if (ret)" helps readability.
> 

OK I'll change.  Thanks!
diff mbox series

Patch

diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 6222a1456ca4..4aa6cf27e1c2 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -31,6 +31,7 @@ 
 #include <asm/msr.h>
 #include <asm/archrandom.h>
 #include <asm/page.h>
+#include <asm/special_insns.h>
 #include <asm/tdx.h>
 #include "tdx.h"
 
@@ -612,7 +613,8 @@  static void tdmr_get_pamt(struct tdmr_info *tdmr, unsigned long *pamt_pfn,
 	*pamt_npages = pamt_sz >> PAGE_SHIFT;
 }
 
-static void tdmr_free_pamt(struct tdmr_info *tdmr)
+static void tdmr_do_pamt_func(struct tdmr_info *tdmr,
+		void (*pamt_func)(unsigned long pfn, unsigned long npages))
 {
 	unsigned long pamt_pfn, pamt_npages;
 
@@ -625,7 +627,12 @@  static void tdmr_free_pamt(struct tdmr_info *tdmr)
 	if (WARN_ON_ONCE(!pamt_pfn))
 		return;
 
-	free_contig_range(pamt_pfn, pamt_npages);
+	(*pamt_func)(pamt_pfn, pamt_npages);
+}
+
+static void tdmr_free_pamt(struct tdmr_info *tdmr)
+{
+	tdmr_do_pamt_func(tdmr, free_contig_range);
 }
 
 static void tdmrs_free_pamt_all(struct tdmr_info_list *tdmr_list)
@@ -656,6 +663,43 @@  static int tdmrs_set_up_pamt_all(struct tdmr_info_list *tdmr_list,
 	return ret;
 }
 
+/*
+ * Convert TDX private pages back to normal by using MOVDIR64B to
+ * clear these pages.  Note this function doesn't flush cache of
+ * these TDX private pages.  The caller should make sure of that.
+ */
+static void reset_tdx_pages(unsigned long pfn, unsigned long npages)
+{
+	const void *zero_page = (const void *)page_address(ZERO_PAGE(0));
+	unsigned long phys, end;
+
+	phys = pfn << PAGE_SHIFT;
+	end = phys + (npages << PAGE_SHIFT);
+
+	for (; phys < end; phys += 64)
+		movdir64b(__va(phys), zero_page);
+
+	/*
+	 * MOVDIR64B uses WC protocol.  Use memory barrier to
+	 * make sure any later user of these pages sees the
+	 * updated data.
+	 */
+	__mb();
+}
+
+static void tdmr_reset_pamt(struct tdmr_info *tdmr)
+{
+	tdmr_do_pamt_func(tdmr, reset_tdx_pages);
+}
+
+static void tdmrs_reset_pamt_all(struct tdmr_info_list *tdmr_list)
+{
+	int i;
+
+	for (i = 0; i < tdmr_list->nr_consumed_tdmrs; i++)
+		tdmr_reset_pamt(tdmr_entry(tdmr_list, i));
+}
+
 static unsigned long tdmrs_count_pamt_pages(struct tdmr_info_list *tdmr_list)
 {
 	unsigned long pamt_npages = 0;
@@ -936,6 +980,55 @@  static int config_tdx_module(struct tdmr_info_list *tdmr_list, u64 global_keyid)
 	return ret;
 }
 
+static int do_global_key_config(void *data)
+{
+	/*
+	 * 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.
+	 *
+	 * All '0's are just unused parameters.
+	 */
+	return seamcall(TDH_SYS_KEY_CONFIG, 0, 0, 0, 0, NULL, NULL);
+}
+
+/*
+ * 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 (TDMR initialization) 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.
+ */
+static int config_global_keyid(void)
+{
+	cpumask_var_t packages;
+	int cpu, ret = -EINVAL;
+
+	if (!zalloc_cpumask_var(&packages, GFP_KERNEL))
+		return -ENOMEM;
+
+	for_each_online_cpu(cpu) {
+		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_on_cpu(cpu, do_global_key_config, NULL, true);
+		if (ret)
+			break;
+	}
+
+	free_cpumask_var(packages);
+	return ret;
+}
+
 static int init_tdx_module(void)
 {
 	static DECLARE_PADDED_STRUCT(tdsysinfo_struct, tdsysinfo,
@@ -980,15 +1073,47 @@  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.
+	 */
+	wbinvd_on_all_cpus();
+
+	/* Config the key of global KeyID on all packages */
+	ret = config_global_keyid();
+	if (ret)
+		goto out_reset_pamts;
+
 	/*
 	 * TODO:
 	 *
-	 *  - Configure the global KeyID on all packages.
 	 *  - Initialize all TDMRs.
 	 *
 	 *  Return error before all steps are done.
 	 */
 	ret = -EINVAL;
+out_reset_pamts:
+	if (ret) {
+		/*
+		 * Part of PAMTs may already have been initialized by the
+		 * TDX module.  Flush cache before returning PAMTs back
+		 * to the kernel.
+		 */
+		wbinvd_on_all_cpus();
+		/*
+		 * According to the TDX hardware spec, if the platform
+		 * doesn't have the "partial write machine check"
+		 * erratum, any kernel read/write will never cause #MC
+		 * in kernel space, thus it's OK to not convert PAMTs
+		 * back to normal.  But do the conversion anyway here
+		 * as suggested by the TDX spec.
+		 */
+		tdmrs_reset_pamt_all(&tdx_tdmr_list);
+	}
 out_free_pamts:
 	if (ret)
 		tdmrs_free_pamt_all(&tdx_tdmr_list);
@@ -1035,6 +1160,9 @@  static int __tdx_enable(void)
  * lock to prevent any new cpu from becoming online; 2) done both VMXON
  * and tdx_cpu_enable() on all online cpus.
  *
+ * This function requires there's at least one online cpu for each CPU
+ * package to succeed.
+ *
  * This function can be called in parallel by multiple callers.
  *
  * Return 0 if TDX is enabled successfully, otherwise error.
diff --git a/arch/x86/virt/vmx/tdx/tdx.h b/arch/x86/virt/vmx/tdx/tdx.h
index 7baf0d662b50..15e7250e74c7 100644
--- a/arch/x86/virt/vmx/tdx/tdx.h
+++ b/arch/x86/virt/vmx/tdx/tdx.h
@@ -25,6 +25,7 @@ 
 #define TDH_SYS_LP_INIT		35
 #define TDH_SYS_INFO		32
 #define TDH_SYS_CONFIG		45
+#define TDH_SYS_KEY_CONFIG	31
 
 struct cmr_info {
 	u64	base;