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

Reclaim of MMIO happen in two cases:
- VM is dying, in that case MMIO would be eagerly reclaimed to the host
  from VM teardown context without host intervention.
- The VM was not launched or died before claiming the device, and it's is
  still considered as host device, but the MMIO was already donated to
  the hypervisor preparing for the VM to access it, in that case the host
  will use this function from an HVC to reclaim the MMIO from KVM/VFIO
  file release context or incase of failure at initialization.

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

arch/arm64/include/asm/kvm_asm.h

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

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

@@ -183,5 +183,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);
+int pkvm_device_reclaim_mmio(u64 pfn, u64 nr_pages);
 
 #endif /* __ARM64_KVM_NVHE_PKVM_H__ */

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

@@ -94,3 +94,37 @@ out_unlock:
 	hyp_spin_unlock(&device_spinlock);
 	return ret;
 }
+
+/*
+ * Reclaim of MMIO can happen in two cases:
+ * - VM is dying, in that case MMIO would be eagerly reclaimed to the host
+ *   from VM teardown context without host intervention.
+ * - The VM was not launched or died before claiming the device, and it's is
+ *   still considered as host device, but the MMIO was already donated to
+ *   the hypervisor preparing for the VM to access it, in that case the host
+ *   will use this function from an HVC to reclaim the MMIO from KVM/VFIO
+ *   file release context or incase of failure at initialization.
+ */
+int pkvm_device_reclaim_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);
+	if (dev->ctxt) {
+		ret = -EBUSY;
+		goto out_unlock;
+	}
+
+	ret = __pkvm_hyp_donate_host(pfn, nr_pages);
+
+out_unlock:
+	hyp_spin_unlock(&device_spinlock);
+	return ret;
+}

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

@@ -1761,6 +1761,17 @@ static void handle___pkvm_host_donate_hyp_mmio(struct kvm_cpu_context *host_ctxt
 	cpu_reg(host_ctxt, 1) = pkvm_device_hyp_assign_mmio(pfn, nr_pages);
 }
 
+static void handle___pkvm_host_reclaim_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_reclaim_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
@@ -1832,6 +1843,7 @@ static const hcall_t host_hcall[] = {
 	HANDLE_FUNC(__pkvm_ptdump),
 	HANDLE_FUNC(__pkvm_host_iommu_map_sg),
 	HANDLE_FUNC(__pkvm_host_donate_hyp_mmio),
+	HANDLE_FUNC(__pkvm_host_reclaim_hyp_mmio),
 };
 
 static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)