ANDROID: KVM: arm64: s2mpu: Add driver initializer

Create 'struct pkvm_iommu_ops' for the S2MPU and a new driver ID to the
list of IOMMU drivers. Implement the 'init' callback, accepting donated
memory from the host to back SMPTs. If the donation is successful,
the SMPTs are assigned to 'host_mpt'.

Export 'pkvm_iommu_s2mpu_register' for a kernel module to call to
register an S2MPU device. First call to this function will also
run the global S2MPU driver initializer.

Bug: 190463801
Change-Id: Icad06379e5cf695fba4f3a18a0773e302f3ead06
Signed-off-by: David Brazdil <dbrazdil@google.com>
(cherry picked from commit 41707102f4)
Signed-off-by: Mostafa Saleh <smostafa@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -387,6 +387,7 @@ enum kvm_iommu_driver {
 };
 
 enum pkvm_iommu_driver_id {
+	PKVM_IOMMU_DRIVER_S2MPU,
 	PKVM_IOMMU_NR_DRIVERS,
 };
 
@@ -401,6 +402,8 @@ int pkvm_iommu_register(struct device *dev, enum pkvm_iommu_driver_id drv_id,
 int pkvm_iommu_suspend(struct device *dev);
 int pkvm_iommu_resume(struct device *dev);
 
+int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t pa);
+
 struct vcpu_reset_state {
 	unsigned long	pc;
 	unsigned long	r0;

arch/arm64/include/asm/kvm_s2mpu.h

@@ -123,8 +123,9 @@ static_assert(SMPT_GRAN <= PAGE_SIZE);
 #define SMPT_ELEMS_PER_WORD			(SMPT_WORD_SIZE * SMPT_ELEMS_PER_BYTE)
 #define SMPT_WORD_BYTE_RANGE			(SMPT_GRAN * SMPT_ELEMS_PER_WORD)
 #define SMPT_NUM_ELEMS				(SZ_1G / SMPT_GRAN)
-#define SMPT_NUM_WORDS				(SMPT_SIZE / SMPT_WORD_SIZE)
 #define SMPT_SIZE				(SMPT_NUM_ELEMS / SMPT_ELEMS_PER_BYTE)
+#define SMPT_NUM_WORDS				(SMPT_SIZE / SMPT_WORD_SIZE)
+#define SMPT_NUM_PAGES				(SMPT_SIZE / PAGE_SIZE)
 #define SMPT_ORDER				get_order(SMPT_SIZE)
 
 /*

arch/arm64/kvm/hyp/include/nvhe/iommu.h

@@ -80,6 +80,8 @@ bool pkvm_iommu_host_dabt_handler(struct kvm_cpu_context *host_ctxt, u32 esr,
 void pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
 				  enum kvm_pgtable_prot prot);
 
+extern const struct pkvm_iommu_ops pkvm_s2mpu_ops;
+
 struct kvm_iommu_ops {
 	int (*init)(void);
 	bool (*host_smc_handler)(struct kvm_cpu_context *host_ctxt);

arch/arm64/kvm/hyp/nvhe/iommu.c

@@ -63,6 +63,8 @@ static inline struct pkvm_iommu_driver *get_driver(enum pkvm_iommu_driver_id id)
 static const struct pkvm_iommu_ops *get_driver_ops(enum pkvm_iommu_driver_id id)
 {
 	switch (id) {
+	case PKVM_IOMMU_DRIVER_S2MPU:
+		return IS_ENABLED(CONFIG_KVM_S2MPU) ? &pkvm_s2mpu_ops : NULL;
 	default:
 		return NULL;
 	}

arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c

@@ -470,50 +470,72 @@ static bool s2mpu_host_mmio_dabt_handler(struct kvm_cpu_context *host_ctxt,
 	return true;
 }
 
-static int s2mpu_init(void)
+static int s2mpu_init(void *data, size_t size)
 {
-	struct pkvm_iommu *dev;
+	struct mpt in_mpt;
+	u32 *smpt;
+	phys_addr_t pa;
 	unsigned int gb;
-	int ret;
+	int ret = 0;
 
-	/* Map data structures in EL2 stage-1. */
-	ret = pkvm_create_mappings(s2mpus,
-				   s2mpus + nr_s2mpus,
-				   PAGE_HYP);
-	if (ret)
-		return ret;
+	if (size != sizeof(in_mpt))
+		return -EINVAL;
 
+	/* The host can concurrently modify 'data'. Copy it to avoid TOCTOU. */
+	memcpy(&in_mpt, data, sizeof(in_mpt));
+
+	/* Take ownership of all SMPT buffers. This will also map them in. */
 	for_each_gb(gb) {
-		ret = pkvm_create_mappings(
-			host_mpt.fmpt[gb].smpt,
-			host_mpt.fmpt[gb].smpt + SMPT_NUM_WORDS,
-			PAGE_HYP);
+		smpt = kern_hyp_va(in_mpt.fmpt[gb].smpt);
+		pa = __hyp_pa(smpt);
+
+		if (!IS_ALIGNED(pa, SMPT_SIZE)) {
+			ret = -EINVAL;
+			break;
+		}
+
+		ret = __pkvm_host_donate_hyp(pa >> PAGE_SHIFT, SMPT_NUM_PAGES);
 		if (ret)
-			return ret;
+			break;
+
+		host_mpt.fmpt[gb] = (struct fmpt){
+			.smpt = smpt,
+			.gran_1g = true,
+			.prot = MPT_PROT_NONE,
+		};
 	}
 
-	/* Map S2MPU MMIO regions in EL2 stage-1. */
-	for_each_s2mpu(dev) {
-		ret = __pkvm_create_private_mapping(
-			dev->pa, S2MPU_MMIO_SIZE, PAGE_HYP_DEVICE,(unsigned long *)(&dev->va));
-		if (ret)
-			return ret;
+	/* Try to return memory back if there was an error. */
+	if (ret) {
+		for_each_gb(gb) {
+			smpt = host_mpt.fmpt[gb].smpt;
+			if (!smpt)
+				break;
+
+			WARN_ON(__pkvm_hyp_donate_host(__hyp_pa(smpt) >> PAGE_SHIFT,
+						       SMPT_NUM_PAGES));
+		}
+		memset(&host_mpt, 0, sizeof(host_mpt));
 	}
 
-	/*
-	 * Program all S2MPUs powered on at boot. Note that they may not be in
-	 * the blocking reset state as the bootloader may have programmed them.
-	 */
-	for_each_powered_s2mpu(dev) {
-		ret = initialize_with_mpt(dev, &host_mpt);
-		if (ret)
-			return ret;
-	}
+	return ret;
+}
+
+static int s2mpu_validate(phys_addr_t pa, size_t size)
+{
+	if (size != S2MPU_MMIO_SIZE)
+		return -EINVAL;
+
 	return 0;
 }
 
-const struct kvm_iommu_ops kvm_s2mpu_ops = (struct kvm_iommu_ops){
+const struct pkvm_iommu_ops pkvm_s2mpu_ops = (struct pkvm_iommu_ops){
 	.init = s2mpu_init,
+	.validate = s2mpu_validate,
+	.data_size = sizeof(struct s2mpu_drv_data),
+};
+
+const struct kvm_iommu_ops kvm_s2mpu_ops = (struct kvm_iommu_ops){
 	.host_smc_handler = s2mpu_host_smc_handler,
 	.host_mmio_dabt_handler = s2mpu_host_mmio_dabt_handler,
 	.host_stage2_set_owner = s2mpu_host_stage2_set_owner,

arch/arm64/kvm/iommu/s2mpu.c

@@ -6,3 +6,81 @@
 
 #include <linux/kvm_host.h>
 #include <asm/kvm_s2mpu.h>
+
+static int init_s2mpu_driver(void)
+{
+	static DEFINE_MUTEX(lock);
+	static bool init_done;
+
+	struct mpt *mpt;
+	unsigned int gb;
+	unsigned long addr;
+	u64 pfn;
+	int ret = 0;
+
+	mutex_lock(&lock);
+	if (init_done)
+		goto out;
+
+	/* Allocate a page for driver data. Must fit MPT descriptor. */
+	BUILD_BUG_ON(sizeof(*mpt) > PAGE_SIZE);
+	addr = __get_free_page(GFP_KERNEL);
+	if (!addr) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	mpt = (struct mpt *)addr;
+
+	/* Allocate SMPT buffers. */
+	for_each_gb(gb) {
+		addr = __get_free_pages(GFP_KERNEL, SMPT_ORDER);
+		if (!addr) {
+			ret = -ENOMEM;
+			goto out_free;
+		}
+		mpt->fmpt[gb].smpt = (u32 *)addr;
+	}
+
+	/* Share MPT descriptor with hyp. */
+	pfn = __pa(mpt) >> PAGE_SHIFT;
+	ret = kvm_call_hyp_nvhe(__pkvm_host_share_hyp, pfn);
+	if (ret)
+		goto out_free;
+
+	/* Hypercall to initialize EL2 driver. */
+	ret = pkvm_iommu_driver_init(PKVM_IOMMU_DRIVER_S2MPU, mpt, sizeof(*mpt));
+	if (ret)
+		goto out_unshare;
+
+	init_done = true;
+
+out_unshare:
+	WARN_ON(kvm_call_hyp_nvhe(__pkvm_host_unshare_hyp, pfn));
+out_free:
+	/* TODO - will driver return the memory? */
+	if (ret) {
+		for_each_gb(gb)
+			free_pages((unsigned long)mpt->fmpt[gb].smpt, SMPT_ORDER);
+		free_page((unsigned long)mpt);
+	}
+out:
+	mutex_unlock(&lock);
+	return ret;
+}
+
+int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t addr)
+{
+	int ret;
+
+	if (!is_protected_kvm_enabled())
+		return -ENODEV;
+
+	ret = init_s2mpu_driver();
+	if (ret)
+		return ret;
+
+	return pkvm_iommu_register(dev, PKVM_IOMMU_DRIVER_S2MPU,
+				   addr, S2MPU_MMIO_SIZE);
+}
+EXPORT_SYMBOL_GPL(pkvm_iommu_s2mpu_register);