const struct tdx_sys_info *tdx_get_sysinfo(void);
 
 int tdx_guest_keyid_alloc(void);
+u32 tdx_get_nr_guest_keyids(void);
 void tdx_guest_keyid_free(unsigned int keyid);
 
 struct tdx_td {
 static inline void tdx_init(void) { }
 static inline int tdx_cpu_enable(void) { return -ENODEV; }
 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; }
 #endif /* CONFIG_INTEL_TDX_HOST */
 
 // SPDX-License-Identifier: GPL-2.0
 #include <linux/cpu.h>
 #include <asm/cpufeature.h>
+#include <linux/misc_cgroup.h>
 #include <asm/tdx.h>
 #include "capabilities.h"
 #include "mmu.h"
        tdx_guest_keyid_free(kvm_tdx->hkid);
        kvm_tdx->hkid = -1;
        atomic_dec(&nr_configured_hkid);
+       misc_cg_uncharge(MISC_CG_RES_TDX, kvm_tdx->misc_cg, 1);
+       put_misc_cg(kvm_tdx->misc_cg);
+       kvm_tdx->misc_cg = NULL;
 }
 
 static inline bool is_hkid_assigned(struct kvm_tdx *kvm_tdx)
        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);
+       if (ret)
+               goto free_hkid;
 
        ret = -ENOMEM;
 
                goto get_sysinfo_err;
        }
 
+       if (misc_cg_set_capacity(MISC_CG_RES_TDX, tdx_get_nr_guest_keyids())) {
+               r = -EINVAL;
+               goto get_sysinfo_err;
+       }
+
        /*
         * Leave hardware virtualization enabled after TDX is enabled
         * successfully.  TDX CPU hotplug depends on this.
 void tdx_cleanup(void)
 {
        if (enable_tdx) {
+               misc_cg_set_capacity(MISC_CG_RES_TDX, 0);
                __tdx_cleanup();
                kvm_disable_virtualization();
        }
 
 struct kvm_tdx {
        struct kvm kvm;
 
+       struct misc_cg *misc_cg;
        int hkid;
        enum kvm_tdx_state state;
 
 
 }
 EXPORT_SYMBOL_GPL(tdx_get_sysinfo);
 
+u32 tdx_get_nr_guest_keyids(void)
+{
+       return tdx_nr_guest_keyids;
+}
+EXPORT_SYMBOL_GPL(tdx_get_nr_guest_keyids);
+
 int tdx_guest_keyid_alloc(void)
 {
        return ida_alloc_range(&tdx_guest_keyid_pool, tdx_guest_keyid_start,
 
        MISC_CG_RES_SEV,
        /** @MISC_CG_RES_SEV_ES: AMD SEV-ES ASIDs resource */
        MISC_CG_RES_SEV_ES,
+#endif
+#ifdef CONFIG_INTEL_TDX_HOST
+       /* Intel TDX HKIDs resource */
+       MISC_CG_RES_TDX,
 #endif
        /** @MISC_CG_RES_TYPES: count of enum misc_res_type constants */
        MISC_CG_RES_TYPES
 
        /* AMD SEV-ES ASIDs resource */
        "sev_es",
 #endif
+#ifdef CONFIG_INTEL_TDX_HOST
+       /* Intel TDX HKIDs resource */
+       "tdx",
+#endif
 };
 
 /* Root misc cgroup */