[PATCH] KVM: TDX: Charge misc cgroup before allocating HKID

Add a tdx_hkid_alloc() helper that charges the misc cgroup before
allocating an HKID, and unwind the charge if HKID allocation fails.

__tdx_td_init() currently allocates an HKID before charging the misc
cgroup. If the charge fails, the error path calls tdx_hkid_free(), which
uncharges a resource that was never successfully charged. This can make
the misc-cgroup usage negative.

Charge the cgroup before allocating the HKID.  Wrapping both steps in
tdx_hkid_alloc() makes it the exact counterpart of tdx_hkid_free(), i.e.
keeps resource allocation and release symmetric, and lets __tdx_td_init()
simply bail on failure instead of open coding the unwind.

Reported-by: [email protected]
Closes: https://lore.kernel.org/all/[email protected]
Closes: https://lore.kernel.org/all/[email protected]
Fixes: 7c035bea9407 ("KVM: TDX: Register TDX host key IDs to cgroup misc controller")
Signed-off-by: Binbin Wu <[email protected]>
---
 arch/x86/kvm/vmx/tdx.c | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index b272c20586a7..d1fe279b5bc3 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -260,6 +260,31 @@ static __always_inline void tdvmcall_set_return_val(struct kvm_vcpu *vcpu,
 	to_tdx(vcpu)->vp_enter_args.r11 = val;
 }
 
+static int tdx_hkid_alloc(struct kvm_tdx *kvm_tdx)
+{
+	int ret;
+
+	kvm_tdx->misc_cg = get_current_misc_cg();
+	ret = misc_cg_try_charge(MISC_CG_RES_TDX, kvm_tdx->misc_cg, 1);
+	if (ret)
+		goto put_misc_cg;
+
+	ret = tdx_guest_keyid_alloc();
+	if (ret < 0)
+		goto uncharge_misc_cg;
+
+	kvm_tdx->hkid = ret;
+	return 0;
+
+uncharge_misc_cg:
+	misc_cg_uncharge(MISC_CG_RES_TDX, kvm_tdx->misc_cg, 1);
+put_misc_cg:
+	put_misc_cg(kvm_tdx->misc_cg);
+	kvm_tdx->misc_cg = NULL;
+
+	return ret;
+}
+
 static inline void tdx_hkid_free(struct kvm_tdx *kvm_tdx)
 {
 	tdx_guest_keyid_free(kvm_tdx->hkid);
@@ -2445,14 +2470,9 @@ static int __tdx_td_init(struct kvm *kvm, struct td_params *td_params,
 	u64 err, rcx;
 
 	*seamcall_err = 0;
-	ret = tdx_guest_keyid_alloc();
-	if (ret < 0)
-		return ret;
-	kvm_tdx->hkid = ret;
-	kvm_tdx->misc_cg = get_current_misc_cg();
-	ret = misc_cg_try_charge(MISC_CG_RES_TDX, kvm_tdx->misc_cg, 1);
+	ret = tdx_hkid_alloc(kvm_tdx);
 	if (ret)
-		goto free_hkid;
+		return ret;
 
 	ret = -ENOMEM;
 

base-commit: 1b731e5ded480bd1e5546aed35584238661ce72e
-- 
2.46.0