[RFC PATCH v1 2/2] cgroup/misc: Add support for Arm CCA MECIDs

[Kohei Enju]

Arm CCA MECIDs are limited platform resources. Add a misc cgroup
resource to allow the system administrator to account for and limit
their use.

The cca_mec resource represents the number of Realm VMs using the
private MEC policy. Each such Realm consumes one MECID. Charge the
resource when the Realm is created, and uncharge it when the Realm is
successfully destroyed.

Realms using the shared MEC policy don't consume this resource.

Signed-off-by: Kohei Enju <enju.kohei at fujitsu.com>
---
 arch/arm64/include/asm/kvm_rmi.h |  2 ++
 arch/arm64/kvm/rmi.c             | 44 +++++++++++++++++++++++++++++++-
 drivers/firmware/arm_rmm/rmi.c   |  4 +++
 include/linux/misc_cgroup.h      |  4 +++
 kernel/cgroup/misc.c             |  4 +++
 5 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/kvm_rmi.h b/arch/arm64/include/asm/kvm_rmi.h
index 721033b132e1..182bdd22014f 100644
--- a/arch/arm64/include/asm/kvm_rmi.h
+++ b/arch/arm64/include/asm/kvm_rmi.h
@@ -61,6 +61,7 @@ enum realm_state {
  * @stage2_unmapped: The Realm stage-2 mappings have been removed
  * @rtts_destroyed: The non-root RTTs have been torn down
  * @mec_policy: MEC policy for the Realm VM
+ * @misc_cg: Misc cgroup charged while the Realm owns a private MECID
  */
 struct realm {
 	void *rd;
@@ -78,6 +79,7 @@ struct realm {
 	bool stage2_unmapped;
 	bool rtts_destroyed;
 	unsigned int mec_policy;
+	struct misc_cg *misc_cg;
 };
 
 /**
diff --git a/arch/arm64/kvm/rmi.c b/arch/arm64/kvm/rmi.c
index 629dea5e8c47..505fa80f635d 100644
--- a/arch/arm64/kvm/rmi.c
+++ b/arch/arm64/kvm/rmi.c
@@ -5,6 +5,7 @@
 
 #include <uapi/linux/psci.h>
 #include <linux/kvm_host.h>
+#include <linux/misc_cgroup.h>
 
 #include <asm/kvm_emulate.h>
 #include <asm/kvm_mmu.h>
@@ -525,6 +526,36 @@ static int realm_init_sve_param(struct kvm *kvm, struct realm_params *params)
 	return 0;
 }
 
+static int realm_mec_charge(struct realm *realm)
+{
+	struct misc_cg *misc_cg;
+	int ret;
+
+	if (realm->mec_policy != ARM_RMI_MEC_POLICY_PRIVATE)
+		return 0;
+
+	misc_cg = get_current_misc_cg();
+
+	ret = misc_cg_try_charge(MISC_CG_RES_CCA_MEC, misc_cg, 1);
+	if (ret) {
+		put_misc_cg(misc_cg);
+		return ret;
+	}
+
+	realm->misc_cg = misc_cg;
+	return 0;
+}
+
+static void realm_mec_uncharge(struct realm *realm)
+{
+	if (!realm->misc_cg)
+		return;
+
+	misc_cg_uncharge(MISC_CG_RES_CCA_MEC, realm->misc_cg, 1);
+	put_misc_cg(realm->misc_cg);
+	realm->misc_cg = NULL;
+}
+
 static int realm_create_rd(struct kvm *kvm)
 {
 	struct realm *realm = &kvm->arch.realm;
@@ -579,11 +610,15 @@ static int realm_create_rd(struct kvm *kvm)
 	if (r)
 		goto out_undelegate_tables;
 
+	r = realm_mec_charge(realm);
+	if (r)
+		goto out_undelegate_tables;
+
 	params_phys = virt_to_phys(params);
 
 	if (rmi_realm_create(rd_phys, params_phys, realm->sro)) {
 		r = -ENXIO;
-		goto out_undelegate_tables;
+		goto out_uncharge_misc_cg;
 	}
 
 	realm->rd = rd;
@@ -594,6 +629,8 @@ static int realm_create_rd(struct kvm *kvm)
 
 	return 0;
 
+out_uncharge_misc_cg:
+	realm_mec_uncharge(realm);
 out_undelegate_tables:
 	if (WARN_ON(rmi_undelegate_range(kvm->arch.mmu.pgd_phys, top_delegated - kvm->arch.mmu.pgd_phys))) {
 		/* Leak the pages if they cannot be returned */
@@ -1478,10 +1515,15 @@ void kvm_destroy_realm(struct kvm *kvm)
 		if (WARN_ON(realm_destroy_rtts(kvm)))
 			return;
 
+		/*
+		 * Keep the misc cgroup charge if RMI_REALM_DESTROY fails, as
+		 * the corresponding MECID may still be allocated by the RMM.
+		 */
 		if (WARN_ON(rmi_realm_destroy(rd_phys, realm->sro)))
 			return;
 		free_delegated_page(rd_phys);
 		realm->rd = NULL;
+		realm_mec_uncharge(realm);
 	}
 
 	if (WARN_ON(rmi_undelegate_range(kvm->arch.mmu.pgd_phys,
diff --git a/drivers/firmware/arm_rmm/rmi.c b/drivers/firmware/arm_rmm/rmi.c
index e9632c35e7db..85a7e6686d2b 100644
--- a/drivers/firmware/arm_rmm/rmi.c
+++ b/drivers/firmware/arm_rmm/rmi.c
@@ -6,6 +6,7 @@
 #include <linux/cpufeature.h>
 #include <linux/memblock.h>
 #include <linux/arm-rmi-cmds.h>
+#include <linux/misc_cgroup.h>
 #include <linux/processor.h>
 #include <linux/slab.h>
 
@@ -714,6 +715,9 @@ static int __init arm64_init_rmi(void)
 	if (ret)
 		return ret;
 
+	WARN_ON_ONCE(misc_cg_set_capacity(MISC_CG_RES_CCA_MEC,
+					  rmi_mec_count()));
+
 	arm64_rmi_is_available = true;
 	pr_info("RMI configured");
 
diff --git a/include/linux/misc_cgroup.h b/include/linux/misc_cgroup.h
index 0cb36a3ffc47..b542502dc52a 100644
--- a/include/linux/misc_cgroup.h
+++ b/include/linux/misc_cgroup.h
@@ -21,6 +21,10 @@ enum misc_res_type {
 #ifdef CONFIG_INTEL_TDX_HOST
 	/** @MISC_CG_RES_TDX: Intel TDX HKIDs resource */
 	MISC_CG_RES_TDX,
+#endif
+#ifdef CONFIG_ARM_RMM
+	/** @MISC_CG_RES_CCA_MEC: ARM64 CCA MECIDs resource */
+	MISC_CG_RES_CCA_MEC,
 #endif
 	/** @MISC_CG_RES_TYPES: count of enum misc_res_type constants */
 	MISC_CG_RES_TYPES
diff --git a/kernel/cgroup/misc.c b/kernel/cgroup/misc.c
index 4a9e2557141c..69259085115e 100644
--- a/kernel/cgroup/misc.c
+++ b/kernel/cgroup/misc.c
@@ -28,6 +28,10 @@ static const char *const misc_res_name[] = {
 	/* Intel TDX HKIDs resource */
 	"tdx",
 #endif
+#ifdef CONFIG_ARM_RMM
+	/* ARM64 CCA MECIDs resource */
+	"cca_mec",
+#endif
 };
 
 /* Root misc cgroup */
-- 
2.43.0