ANDROID: KVM: arm64: Add HVC to donate assignable MMIO

Devices assigned to guest has to transition first to hypervisor,
this guarantees that there is a point of time that the device is
neither accessible from the host or the guest, so the hypervisor
can reset it and block it's IOMMU.
The host will donate the whole device first to the hypervisor
before the guest touches or requests any part of the device
and upon the first request or access the hypervisor will ensure
that the device is fully donated and then will execute the transition
sequence before assigning to the guest.

Bug: 357781595
Bug: 348382247
Change-Id: I4ae5fc888e2eb4cd98a036d42d090c84d865f6bd
Signed-off-by: Mostafa Saleh <smostafa@google.com>

arch/arm64/include/asm/kvm_asm.h

@@ -120,6 +120,7 @@ enum __kvm_host_smccc_func {
 	__KVM_HOST_SMCCC_FUNC___pkvm_host_hvc_pd,
 	__KVM_HOST_SMCCC_FUNC___pkvm_ptdump,
 	__KVM_HOST_SMCCC_FUNC___pkvm_host_iommu_map_sg,
+	__KVM_HOST_SMCCC_FUNC___pkvm_host_donate_hyp_mmio,
 
 	/*
 	 * Start of the dynamically registered hypercalls. Start a bit

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

@@ -182,5 +182,6 @@ static inline int pkvm_init_power_domain(struct kvm_power_domain *pd,
 }
 
 int pkvm_init_devices(void);
+int pkvm_device_hyp_assign_mmio(u64 pfn, u64 nr_pages);
 
 #endif /* __ARM64_KVM_NVHE_PKVM_H__ */

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

@@ -12,6 +12,16 @@
 struct pkvm_device *registered_devices;
 unsigned long registered_devices_nr;
 
+/*
+ * This lock protects all devices in registered_devices when ctxt changes,
+ * this is overlocking and can be improved. However, the device context
+ * only changes at boot time and at teardown and in theory there shouldn't
+ * be congestion on that path.
+ * All changes/checks to MMIO state or IOMMU must be atomic with the ctxt
+ * of the device.
+ */
+static DEFINE_HYP_SPINLOCK(device_spinlock);
+
 int pkvm_init_devices(void)
 {
 	size_t dev_sz;
@@ -27,3 +37,60 @@ int pkvm_init_devices(void)
 		registered_devices_nr = 0;
 	return ret;
 }
+
+/* return device from a resource, addr and size must match. */
+static struct pkvm_device *pkvm_get_device(u64 addr, size_t size)
+{
+	struct pkvm_device *dev;
+	struct pkvm_dev_resource *res;
+	int i, j;
+
+	for (i = 0 ; i < registered_devices_nr ; ++i) {
+		dev = &registered_devices[i];
+		for (j = 0 ; j < dev->nr_resources; ++j) {
+			res = &dev->resources[j];
+			if ((addr == res->base) && (size == res->size))
+				return dev;
+		}
+	}
+
+	return NULL;
+}
+
+/*
+ * Devices assigned to guest has to transition first to hypervisor,
+ * this guarantees that there is a point of time that the device is
+ * neither accessible from the host or the guest, so the hypervisor
+ * can reset it and block it's IOOMU.
+ * The host will donate the whole device first to the hypervisor
+ * before the guest touches or requests any part of the device
+ * and upon the first request or access the hypervisor will ensure
+ * that the device is fully donated first.
+ */
+int pkvm_device_hyp_assign_mmio(u64 pfn, u64 nr_pages)
+{
+	struct pkvm_device *dev;
+	int ret;
+	size_t size = nr_pages << PAGE_SHIFT;
+	u64 phys = pfn << PAGE_SHIFT;
+
+	dev = pkvm_get_device(phys, size);
+	if (!dev)
+		return -ENODEV;
+
+	hyp_spin_lock(&device_spinlock);
+	/* A VM already have this device, no take backs. */
+	if (dev->ctxt) {
+		ret = -EBUSY;
+		goto out_unlock;
+	}
+
+	ret = ___pkvm_host_donate_hyp_prot(pfn, nr_pages, true, PAGE_HYP_DEVICE);
+	/* Hyp have device mapping, while host may have issue cacheable writes.*/
+	if (!ret)
+		kvm_flush_dcache_to_poc(__hyp_va(phys), PAGE_SIZE);
+
+out_unlock:
+	hyp_spin_unlock(&device_spinlock);
+	return ret;
+}

arch/arm64/kvm/hyp/nvhe/hyp-main.c

@@ -1750,6 +1750,17 @@ static void handle___pkvm_host_iommu_map_sg(struct kvm_cpu_context *host_ctxt)
 	hyp_reqs_smccc_encode(ret, host_ctxt, this_cpu_ptr(&host_hyp_reqs));
 }
 
+static void handle___pkvm_host_donate_hyp_mmio(struct kvm_cpu_context *host_ctxt)
+{
+	DECLARE_REG(u64, pfn, host_ctxt, 1);
+	DECLARE_REG(u64, nr_pages, host_ctxt, 2);
+
+	if (!is_protected_kvm_enabled())
+		return;
+
+	cpu_reg(host_ctxt, 1) = pkvm_device_hyp_assign_mmio(pfn, nr_pages);
+}
+
 typedef void (*hcall_t)(struct kvm_cpu_context *);
 
 #define HANDLE_FUNC(x)	[__KVM_HOST_SMCCC_FUNC_##x] = (hcall_t)handle_##x
@@ -1820,6 +1831,7 @@ static const hcall_t host_hcall[] = {
 	HANDLE_FUNC(__pkvm_host_hvc_pd),
 	HANDLE_FUNC(__pkvm_ptdump),
 	HANDLE_FUNC(__pkvm_host_iommu_map_sg),
+	HANDLE_FUNC(__pkvm_host_donate_hyp_mmio),
 };
 
 static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)