[v19,130/130] RFC: KVM: x86, TDX: Add check for KVM_SET_CPUID2

Commit Message

Isaku Yamahata Feb. 26, 2024, 8:27 a.m. UTC

From: Isaku Yamahata <[e-mail]>

Implement a hook of KVM_SET_CPUID2 for additional consistency check.

Intel TDX or AMD SEV has a restriction on the value of cpuid.  For example,
some values must be the same between all vcpus.  Check if the new values
are consistent with the old values.  The check is light because the cpuid
consistency is very model specific and complicated.  The user space VMM
should set cpuid and MSRs consistently.

Suggested-by: Sean Christopherson <[e-mail]>
Link: https://lore.kernel.org/lkml/[e-mail]/
Signed-off-by: Isaku Yamahata <[e-mail]>
---
v18:
- Use TDH.SYS.RD() instead of struct tdsysinfo_struct
---
 arch/x86/kvm/vmx/main.c    | 10 ++++++
 arch/x86/kvm/vmx/tdx.c     | 66 ++++++++++++++++++++++++++++++++++----
 arch/x86/kvm/vmx/tdx.h     |  7 ++++
 arch/x86/kvm/vmx/x86_ops.h |  4 +++
 4 files changed, 81 insertions(+), 6 deletions(-)

Comments

Edgecombe, Rick P March 21, 2024, 11:12 p.m. UTC | #1

On Mon, 2024-02-26 at 00:27 -0800, [e-mail] wrote:
> Implement a hook of KVM_SET_CPUID2 for additional consistency check.
> 
> Intel TDX or AMD SEV has a restriction on the value of cpuid.  For
> example,
> some values must be the same between all vcpus.  Check if the new
> values
> are consistent with the old values.  The check is light because the
> cpuid
> consistency is very model specific and complicated.  The user space
> VMM
> should set cpuid and MSRs consistently.

I see that this was suggested by Sean, but can you explain the problem
that this is working around? From the linked thread, it seems like the
problem is what to do when userspace also calls SET_CPUID after already
configuring CPUID to the TDX module in the special way. The choices
discussed included:
1. Reject the call
2. Check the consistency between the first CPUID configuration and the
second one.

1 is a lot simpler, but the reasoning for 2 is because "some KVM code
paths rely on guest CPUID configuration" it seems. Is this a
hypothetical or real issue? Which code paths are problematic for
TDX/SNP?

Just trying to assess what we should do with these two patches.

Kai Huang March 22, 2024, 7:10 a.m. UTC | #2

On Thu, 2024-03-21 at 23:12 +0000, Edgecombe, Rick P wrote:
> On Mon, 2024-02-26 at 00:27 -0800, [e-mail] wrote:
> > Implement a hook of KVM_SET_CPUID2 for additional consistency check.
> > 
> > Intel TDX or AMD SEV has a restriction on the value of cpuid.  For
> > example,
> > some values must be the same between all vcpus.  Check if the new
> > values
> > are consistent with the old values.  The check is light because the
> > cpuid
> > consistency is very model specific and complicated.  The user space
> > VMM
> > should set cpuid and MSRs consistently.
> 
> I see that this was suggested by Sean, but can you explain the problem
> that this is working around? From the linked thread, it seems like the
> problem is what to do when userspace also calls SET_CPUID after already
> configuring CPUID to the TDX module in the special way. The choices
> discussed included:
> 1. Reject the call
> 2. Check the consistency between the first CPUID configuration and the
> second one.
> 
> 1 is a lot simpler, but the reasoning for 2 is because "some KVM code
> paths rely on guest CPUID configuration" it seems. Is this a
> hypothetical or real issue? Which code paths are problematic for
> TDX/SNP?

There might be use case that TDX guest wants to use some CPUID which
isn't handled by the TDX module but purely by KVM.  These (PV) CPUIDs need to be
provided via KVM_SET_CPUID2.


Btw, Isaku, I don't understand why you tag the last two patches as RFC and put
them at last.  I think I've expressed this before.  Per the discussion with
Sean, my understanding is this isn't something optional but the right thing we
should do?

https://lore.kernel.org/lkml/[e-mail]/

Edgecombe, Rick P March 22, 2024, 4:06 p.m. UTC | #3

On Fri, 2024-03-22 at 07:10 +0000, Huang, Kai wrote:
> > I see that this was suggested by Sean, but can you explain the
> > problem
> > that this is working around? From the linked thread, it seems like
> > the
> > problem is what to do when userspace also calls SET_CPUID after
> > already
> > configuring CPUID to the TDX module in the special way. The choices
> > discussed included:
> > 1. Reject the call
> > 2. Check the consistency between the first CPUID configuration and
> > the
> > second one.
> > 
> > 1 is a lot simpler, but the reasoning for 2 is because "some KVM
> > code
> > paths rely on guest CPUID configuration" it seems. Is this a
> > hypothetical or real issue? Which code paths are problematic for
> > TDX/SNP?
> 
> There might be use case that TDX guest wants to use some CPUID which
> isn't handled by the TDX module but purely by KVM.  These (PV) CPUIDs
> need to be
> provided via KVM_SET_CPUID2.

Right, but are there any needed today? I read that Sean's point was
that KVM_SET_CPUID2 can't accept anything today what we would want to
block later, otherwise it would introduce a regression. This was the
major constraint IIUC, and means the base series requires *something*
here.

If we want to support only the most basic support first, we don't need
to support PV CPUIDs on day 1, right?

So I'm wondering, if we could shrink the base series by going with
option 1 to start, and then expanding it with this solution later to
enable more features. Do you see a problem or conflict with Sean's
comments?

Isaku Yamahata March 23, 2024, 1:54 a.m. UTC | #4

On Fri, Mar 22, 2024 at 07:10:42AM +0000,
"Huang, Kai" <[e-mail]> wrote:

> On Thu, 2024-03-21 at 23:12 +0000, Edgecombe, Rick P wrote:
> > On Mon, 2024-02-26 at 00:27 -0800, [e-mail] wrote:
> > > Implement a hook of KVM_SET_CPUID2 for additional consistency check.
> > > 
> > > Intel TDX or AMD SEV has a restriction on the value of cpuid.  For
> > > example,
> > > some values must be the same between all vcpus.  Check if the new
> > > values
> > > are consistent with the old values.  The check is light because the
> > > cpuid
> > > consistency is very model specific and complicated.  The user space
> > > VMM
> > > should set cpuid and MSRs consistently.
> > 
> > I see that this was suggested by Sean, but can you explain the problem
> > that this is working around? From the linked thread, it seems like the
> > problem is what to do when userspace also calls SET_CPUID after already
> > configuring CPUID to the TDX module in the special way. The choices
> > discussed included:
> > 1. Reject the call
> > 2. Check the consistency between the first CPUID configuration and the
> > second one.
> > 
> > 1 is a lot simpler, but the reasoning for 2 is because "some KVM code
> > paths rely on guest CPUID configuration" it seems. Is this a
> > hypothetical or real issue? Which code paths are problematic for
> > TDX/SNP?
> 
> There might be use case that TDX guest wants to use some CPUID which
> isn't handled by the TDX module but purely by KVM.  These (PV) CPUIDs need to be
> provided via KVM_SET_CPUID2.
> 
> 
> Btw, Isaku, I don't understand why you tag the last two patches as RFC and put
> them at last.  I think I've expressed this before.  Per the discussion with
> Sean, my understanding is this isn't something optional but the right thing we
> should do?
> 
> https://lore.kernel.org/lkml/[e-mail]/

Ok, let's remove RFC and reorder this patches.  Do you see any issue of the
cpuid check logic itself?

Kai Huang March 25, 2024, 11:14 a.m. UTC | #5

On Fri, 2024-03-22 at 16:06 +0000, Edgecombe, Rick P wrote:
> On Fri, 2024-03-22 at 07:10 +0000, Huang, Kai wrote:
> > > I see that this was suggested by Sean, but can you explain the
> > > problem
> > > that this is working around? From the linked thread, it seems like
> > > the
> > > problem is what to do when userspace also calls SET_CPUID after
> > > already
> > > configuring CPUID to the TDX module in the special way. The choices
> > > discussed included:
> > > 1. Reject the call
> > > 2. Check the consistency between the first CPUID configuration and
> > > the
> > > second one.
> > > 
> > > 1 is a lot simpler, but the reasoning for 2 is because "some KVM
> > > code
> > > paths rely on guest CPUID configuration" it seems. Is this a
> > > hypothetical or real issue? Which code paths are problematic for
> > > TDX/SNP?
> > 
> > There might be use case that TDX guest wants to use some CPUID which
> > isn't handled by the TDX module but purely by KVM.  These (PV) CPUIDs
> > need to be
> > provided via KVM_SET_CPUID2.
> 
> Right, but are there any needed today? 
> 

I am not sure.  Isaku may know better?

> I read that Sean's point was
> that KVM_SET_CPUID2 can't accept anything today what we would want to
> block later, otherwise it would introduce a regression. This was the
> major constraint IIUC, and means the base series requires *something*
> here.
> 
> If we want to support only the most basic support first, we don't need
> to support PV CPUIDs on day 1, right?
> 
> So I'm wondering, if we could shrink the base series by going with
> option 1 to start, and then expanding it with this solution later to
> enable more features. Do you see a problem or conflict with Sean's
> comments?
> 
> 

To confirm, I mean you want to simply make KVM_SET_CPUID2 return error for TDX
guest?

It is acceptable to me, and I don't see any conflict with Sean's comments.

But I don't know Sean's perference.  As he said, I think  the consistency
checking is quite straight-forward:

"
It's not complicated at all.  Walk through the leafs defined during
TDH.MNG.INIT, reject KVM_SET_CPUID if a leaf isn't present or doesn't match
exactly.
"

So to me it's not a big deal. 

Either way, we need a patch to handle SET_CPUID2:

1) if we go option 1) -- that is reject SET_CPUID2 completely -- we need to make
vcpu's CPUID point to KVM's saved CPUID during TDH.MNG.INIT.

2) if we do consistency check, we do a for loop and reject when in-consistency
found.

I'll leave to you to judge :-)

Edgecombe, Rick P March 25, 2024, 3:32 p.m. UTC | #6

On Mon, 2024-03-25 at 11:14 +0000, Huang, Kai wrote:
> To confirm, I mean you want to simply make KVM_SET_CPUID2 return error for TDX
> guest?
> 
> It is acceptable to me, and I don't see any conflict with Sean's comments.
> 
> But I don't know Sean's perference.  As he said, I think  the consistency
> checking is quite straight-forward:
> 
> "
> It's not complicated at all.  Walk through the leafs defined during
> TDH.MNG.INIT, reject KVM_SET_CPUID if a leaf isn't present or doesn't match
> exactly.
> "
> 
Yea, I'm just thinking if we could take two patches down to one small one it might be a way to
essentially break off this work to another series without affecting the ability to boot a TD. It
*seems* to be the way things are going.

> So to me it's not a big deal. 
> 
> Either way, we need a patch to handle SET_CPUID2:
> 
> 1) if we go option 1) -- that is reject SET_CPUID2 completely -- we need to make
> vcpu's CPUID point to KVM's saved CPUID during TDH.MNG.INIT.

Ah, I missed this part. Can you elaborate? By dropping these two patches it doesn't prevent a TD
boot. If we then reject SET_CPUID, this will break things unless we make other changes? And they are
not small?

Isaku Yamahata March 25, 2024, 9:10 p.m. UTC | #7

On Mon, Mar 25, 2024 at 11:14:21AM +0000,
"Huang, Kai" <[e-mail]> wrote:

> On Fri, 2024-03-22 at 16:06 +0000, Edgecombe, Rick P wrote:
> > On Fri, 2024-03-22 at 07:10 +0000, Huang, Kai wrote:
> > > > I see that this was suggested by Sean, but can you explain the
> > > > problem
> > > > that this is working around? From the linked thread, it seems like
> > > > the
> > > > problem is what to do when userspace also calls SET_CPUID after
> > > > already
> > > > configuring CPUID to the TDX module in the special way. The choices
> > > > discussed included:
> > > > 1. Reject the call
> > > > 2. Check the consistency between the first CPUID configuration and
> > > > the
> > > > second one.
> > > > 
> > > > 1 is a lot simpler, but the reasoning for 2 is because "some KVM
> > > > code
> > > > paths rely on guest CPUID configuration" it seems. Is this a
> > > > hypothetical or real issue? Which code paths are problematic for
> > > > TDX/SNP?
> > > 
> > > There might be use case that TDX guest wants to use some CPUID which
> > > isn't handled by the TDX module but purely by KVM.  These (PV) CPUIDs
> > > need to be
> > > provided via KVM_SET_CPUID2.
> > 
> > Right, but are there any needed today? 
> > 
> 
> I am not sure.  Isaku may know better?

It's not needed to boot TD.  The check is safe guard.  The multiple of source of
cpuids can be inconsistent.

Isaku Yamahata March 25, 2024, 9:17 p.m. UTC | #8

On Mon, Mar 25, 2024 at 03:32:59PM +0000,
"Edgecombe, Rick P" <[e-mail]> wrote:

> On Mon, 2024-03-25 at 11:14 +0000, Huang, Kai wrote:
> > To confirm, I mean you want to simply make KVM_SET_CPUID2 return error for TDX
> > guest?
> > 
> > It is acceptable to me, and I don't see any conflict with Sean's comments.
> > 
> > But I don't know Sean's perference.  As he said, I think  the consistency
> > checking is quite straight-forward:
> > 
> > "
> > It's not complicated at all.  Walk through the leafs defined during
> > TDH.MNG.INIT, reject KVM_SET_CPUID if a leaf isn't present or doesn't match
> > exactly.
> > "
> > 
> Yea, I'm just thinking if we could take two patches down to one small one it might be a way to
> essentially break off this work to another series without affecting the ability to boot a TD. It
> *seems* to be the way things are going.
> 
> > So to me it's not a big deal. 
> > 
> > Either way, we need a patch to handle SET_CPUID2:
> > 
> > 1) if we go option 1) -- that is reject SET_CPUID2 completely -- we need to make
> > vcpu's CPUID point to KVM's saved CPUID during TDH.MNG.INIT.
> 
> Ah, I missed this part. Can you elaborate? By dropping these two patches it doesn't prevent a TD
> boot. If we then reject SET_CPUID, this will break things unless we make other changes? And they are
> not small?

If we go forthis, the extended topology enumeration (cpuid[0xb or 0x1f]) would
need special handling because it's per-vcpu. not TD wide.

Kai Huang March 25, 2024, 10:31 p.m. UTC | #9

> On Mon, 2024-03-25 at 11:14 +0000, Huang, Kai wrote:
> > To confirm, I mean you want to simply make KVM_SET_CPUID2 return error
> > for TDX guest?
> >
> > It is acceptable to me, and I don't see any conflict with Sean's comments.
> >
> > But I don't know Sean's perference.  As he said, I think  the
> > consistency checking is quite straight-forward:
> >
> > "
> > It's not complicated at all.  Walk through the leafs defined during
> > TDH.MNG.INIT, reject KVM_SET_CPUID if a leaf isn't present or doesn't
> > match exactly.
> > "
> >
> Yea, I'm just thinking if we could take two patches down to one small one it
> might be a way to essentially break off this work to another series without
> affecting the ability to boot a TD. It
> *seems* to be the way things are going.
> 
> > So to me it's not a big deal.
> >
> > Either way, we need a patch to handle SET_CPUID2:
> >
> > 1) if we go option 1) -- that is reject SET_CPUID2 completely -- we
> > need to make vcpu's CPUID point to KVM's saved CPUID during
> TDH.MNG.INIT.
> 
> Ah, I missed this part. Can you elaborate? By dropping these two patches it
> doesn't prevent a TD boot. If we then reject SET_CPUID, this will break things
> unless we make other changes? And they are not small?
> 

(sorry replying from outlook due to some issue to my linux box environment)

It booted because Qemu does sane thing, i.e., it always passes the correct CPUIDs in KVM_SET_CPUID2.

Per-Sean's comments, KVM should guarantee the consistency between CPUIDs done in TDH.MNG.INIT and KVM_SET_CPUID2, otherwise if Qemu passes in-consistent CPUIDs KVM can easily fail to work with TD.

To guarantee the consistency, KVM could do two options as we discussed:

1) reject KVM_SET_CPUID2 completely.
2) Still allow KVM_SET_CPUID2 but manually check the CPUID consistency between the one done in TDH.MNG.INIT and the one passed in KVM_SET_CPUID2.

1) can obviously guarantee consistency.  But KVM maintains CPUIDs in 'vcpu', so to make the existing KVM code continue to work, we need to manually set 'vcpu->cpuid' to the one that is done in TDH.MNG.INIT. 

2) you need to check the consistency and reject KVM_SET_CPUID2 if in-consistency found.  But other than that, KVM doesn't need to anything more because if we allow KVM_SET_CPUID2, the 'vcpu' will have its own CPUIDs populated anyway.

Edgecombe, Rick P March 25, 2024, 10:37 p.m. UTC | #10

On Mon, 2024-03-25 at 22:31 +0000, Huang, Kai wrote:
> (sorry replying from outlook due to some issue to my linux box environment)
> 
> It booted because Qemu does sane thing, i.e., it always passes the correct CPUIDs in
> KVM_SET_CPUID2.
> 
> Per-Sean's comments, KVM should guarantee the consistency between CPUIDs done in TDH.MNG.INIT and
> KVM_SET_CPUID2, otherwise if Qemu passes in-consistent CPUIDs KVM can easily fail to work with TD.
> 
> To guarantee the consistency, KVM could do two options as we discussed:
> 
> 1) reject KVM_SET_CPUID2 completely.
> 2) Still allow KVM_SET_CPUID2 but manually check the CPUID consistency between the one done in
> TDH.MNG.INIT and the one passed in KVM_SET_CPUID2.
> 
> 1) can obviously guarantee consistency.  But KVM maintains CPUIDs in 'vcpu', so to make the
> existing KVM code continue to work, we need to manually set 'vcpu->cpuid' to the one that is done
> in TDH.MNG.INIT. 
> 
> 2) you need to check the consistency and reject KVM_SET_CPUID2 if in-consistency found.  But other
> than that, KVM doesn't need to anything more because if we allow KVM_SET_CPUID2, the 'vcpu' will
> have its own CPUIDs populated anyway.

Ah, thanks for explaining. So 1 is not that simple, it is a maybe slightly smaller separate
solution. Now I see why the discussion was to just do the consistency checking up front.

Patch

diff --git a/arch/x86/kvm/vmx/main.c b/arch/x86/kvm/vmx/main.c
index 2cd404fd7176..1a979ec644d0 100644
--- a/arch/x86/kvm/vmx/main.c
+++ b/arch/x86/kvm/vmx/main.c
@@ -432,6 +432,15 @@  static void vt_vcpu_deliver_init(struct kvm_vcpu *vcpu)
 	kvm_vcpu_deliver_init(vcpu);
 }
 
+static int vt_vcpu_check_cpuid(struct kvm_vcpu *vcpu,
+			       struct kvm_cpuid_entry2 *e2, int nent)
+{
+	if (is_td_vcpu(vcpu))
+		return tdx_vcpu_check_cpuid(vcpu, e2, nent);
+
+	return 0;
+}
+
 static void vt_vcpu_after_set_cpuid(struct kvm_vcpu *vcpu)
 {
 	if (is_td_vcpu(vcpu))
@@ -1125,6 +1134,7 @@  struct kvm_x86_ops vt_x86_ops __initdata = {
 
 	.get_exit_info = vt_get_exit_info,
 
+	.vcpu_check_cpuid = vt_vcpu_check_cpuid,
 	.vcpu_after_set_cpuid = vt_vcpu_after_set_cpuid,
 
 	.has_wbinvd_exit = cpu_has_vmx_wbinvd_exit,
diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 7be1be161dc2..a71093f7c3e3 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -547,6 +547,9 @@  void tdx_vm_free(struct kvm *kvm)
 
 	free_page((unsigned long)__va(kvm_tdx->tdr_pa));
 	kvm_tdx->tdr_pa = 0;
+
+	kfree(kvm_tdx->cpuid);
+	kvm_tdx->cpuid = NULL;
 }
 
 static int tdx_do_tdh_mng_key_config(void *param)
@@ -661,6 +664,39 @@  int tdx_vcpu_create(struct kvm_vcpu *vcpu)
 	return 0;
 }
 
+int tdx_vcpu_check_cpuid(struct kvm_vcpu *vcpu, struct kvm_cpuid_entry2 *e2, int nent)
+{
+	struct kvm_tdx *kvm_tdx = to_kvm_tdx(vcpu->kvm);
+	int i;
+
+	/*
+	 * Simple check that new cpuid is consistent with created one.
+	 * For simplicity, only trivial check.  Don't try comprehensive checks
+	 * with the cpuid virtualization table in the TDX module spec.
+	 */
+	for (i = 0; i < tdx_info->num_cpuid_config; i++) {
+		const struct kvm_tdx_cpuid_config *c = &tdx_info->cpuid_configs[i];
+		u32 index = c->sub_leaf == KVM_TDX_CPUID_NO_SUBLEAF ? 0 : c->sub_leaf;
+		const struct kvm_cpuid_entry2 *old =
+			kvm_find_cpuid_entry2(kvm_tdx->cpuid, kvm_tdx->cpuid_nent,
+					      c->leaf, index);
+		const struct kvm_cpuid_entry2 *new = kvm_find_cpuid_entry2(e2, nent,
+									   c->leaf, index);
+
+		if (!!old != !!new)
+			return -EINVAL;
+		if (!old && !new)
+			continue;
+
+		if ((old->eax ^ new->eax) & c->eax ||
+		    (old->ebx ^ new->ebx) & c->ebx ||
+		    (old->ecx ^ new->ecx) & c->ecx ||
+		    (old->edx ^ new->edx) & c->edx)
+			return -EINVAL;
+	}
+	return 0;
+}
+
 void tdx_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
 {
 	struct vcpu_tdx *tdx = to_tdx(vcpu);
@@ -2205,9 +2241,10 @@  static int setup_tdparams_eptp_controls(struct kvm_cpuid2 *cpuid,
 	return 0;
 }
 
-static void setup_tdparams_cpuids(struct kvm_cpuid2 *cpuid,
+static void setup_tdparams_cpuids(struct kvm *kvm, struct kvm_cpuid2 *cpuid,
 				  struct td_params *td_params)
 {
+	struct kvm_tdx *kvm_tdx = to_kvm_tdx(kvm);
 	int i;
 
 	/*
@@ -2215,6 +2252,7 @@  static void setup_tdparams_cpuids(struct kvm_cpuid2 *cpuid,
 	 * be same to the one of struct tdsysinfo.{num_cpuid_config, cpuid_configs}
 	 * It's assumed that td_params was zeroed.
 	 */
+	kvm_tdx->cpuid_nent = 0;
 	for (i = 0; i < tdx_info->num_cpuid_config; i++) {
 		const struct kvm_tdx_cpuid_config *c = &tdx_info->cpuid_configs[i];
 		/* KVM_TDX_CPUID_NO_SUBLEAF means index = 0. */
@@ -2237,6 +2275,10 @@  static void setup_tdparams_cpuids(struct kvm_cpuid2 *cpuid,
 		value->ebx = entry->ebx & c->ebx;
 		value->ecx = entry->ecx & c->ecx;
 		value->edx = entry->edx & c->edx;
+
+		/* Remember the setting to check for KVM_SET_CPUID2. */
+		kvm_tdx->cpuid[kvm_tdx->cpuid_nent] = *entry;
+		kvm_tdx->cpuid_nent++;
 	}
 }
 
@@ -2324,7 +2366,7 @@  static int setup_tdparams(struct kvm *kvm, struct td_params *td_params,
 	ret = setup_tdparams_eptp_controls(cpuid, td_params);
 	if (ret)
 		return ret;
-	setup_tdparams_cpuids(cpuid, td_params);
+	setup_tdparams_cpuids(kvm, cpuid, td_params);
 	ret = setup_tdparams_xfam(cpuid, td_params);
 	if (ret)
 		return ret;
@@ -2548,11 +2590,18 @@  static int tdx_td_init(struct kvm *kvm, struct kvm_tdx_cmd *cmd)
 	if (cmd->flags)
 		return -EINVAL;
 
-	init_vm = kzalloc(sizeof(*init_vm) +
-			  sizeof(init_vm->cpuid.entries[0]) * KVM_MAX_CPUID_ENTRIES,
-			  GFP_KERNEL);
-	if (!init_vm)
+	WARN_ON_ONCE(kvm_tdx->cpuid);
+	kvm_tdx->cpuid = kzalloc(flex_array_size(init_vm, cpuid.entries, KVM_MAX_CPUID_ENTRIES),
+				 GFP_KERNEL);
+	if (!kvm_tdx->cpuid)
 		return -ENOMEM;
+
+	init_vm = kzalloc(struct_size(init_vm, cpuid.entries, KVM_MAX_CPUID_ENTRIES),
+			  GFP_KERNEL);
+	if (!init_vm) {
+		ret = -ENOMEM;
+		goto out;
+	}
 	if (copy_from_user(init_vm, (void __user *)cmd->data, sizeof(*init_vm))) {
 		ret = -EFAULT;
 		goto out;
@@ -2602,6 +2651,11 @@  static int tdx_td_init(struct kvm *kvm, struct kvm_tdx_cmd *cmd)
 
 out:
 	/* kfree() accepts NULL. */
+	if (ret) {
+		kfree(kvm_tdx->cpuid);
+		kvm_tdx->cpuid = NULL;
+		kvm_tdx->cpuid_nent = 0;
+	}
 	kfree(init_vm);
 	kfree(td_params);
 	return ret;
diff --git a/arch/x86/kvm/vmx/tdx.h b/arch/x86/kvm/vmx/tdx.h
index 11c74c34555f..af3a2b8afee8 100644
--- a/arch/x86/kvm/vmx/tdx.h
+++ b/arch/x86/kvm/vmx/tdx.h
@@ -31,6 +31,13 @@  struct kvm_tdx {
 
 	u64 tsc_offset;
 
+	/*
+	 * For KVM_SET_CPUID to check consistency. Remember the one passed to
+	 * TDH.MNG_INIT
+	 */
+	int cpuid_nent;
+	struct kvm_cpuid_entry2 *cpuid;
+
 	/* For KVM_MEMORY_MAPPING */
 	struct mutex source_lock;
 	struct page *source_page;
diff --git a/arch/x86/kvm/vmx/x86_ops.h b/arch/x86/kvm/vmx/x86_ops.h
index c507f1513dac..6b067842a67f 100644
--- a/arch/x86/kvm/vmx/x86_ops.h
+++ b/arch/x86/kvm/vmx/x86_ops.h
@@ -162,6 +162,8 @@  u8 tdx_get_mt_mask(struct kvm_vcpu *vcpu, gfn_t gfn, bool is_mmio);
 
 void tdx_deliver_interrupt(struct kvm_lapic *apic, int delivery_mode,
 			   int trig_mode, int vector);
+int tdx_vcpu_check_cpuid(struct kvm_vcpu *vcpu, struct kvm_cpuid_entry2 *e2,
+			 int nent);
 void tdx_inject_nmi(struct kvm_vcpu *vcpu);
 void tdx_get_exit_info(struct kvm_vcpu *vcpu, u32 *reason,
 		u64 *info1, u64 *info2, u32 *intr_info, u32 *error_code);
@@ -221,6 +223,8 @@  static inline u8 tdx_get_mt_mask(struct kvm_vcpu *vcpu, gfn_t gfn, bool is_mmio)
 
 static inline void tdx_deliver_interrupt(struct kvm_lapic *apic, int delivery_mode,
 					 int trig_mode, int vector) {}
+static inline int tdx_vcpu_check_cpuid(struct kvm_vcpu *vcpu, struct kvm_cpuid_entry2 *e2,
+				       int nent) { return -EOPNOTSUPP; }
 static inline void tdx_inject_nmi(struct kvm_vcpu *vcpu) {}
 static inline void tdx_get_exit_info(struct kvm_vcpu *vcpu, u32 *reason, u64 *info1,
 				     u64 *info2, u32 *intr_info, u32 *error_code) {}