ANDROID: KVM: arm64: Add __pkvm_host_map_guest_mmio()

When a guest touches a page, the host will map it, similarly to normal
memory. However, there is a separate hypercall for MMIO to avoid
increasing the cost of memory mapping which is the hot path.

Upon MMIO mapping, the hypervisor will atomically check that parts
of the group the device belonging to is already donated, and tag
those devices as owned by the VM.

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

arch/arm64/include/asm/kvm_asm.h

@@ -122,6 +122,7 @@ enum __kvm_host_smccc_func {
 	__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,
+	__KVM_HOST_SMCCC_FUNC___pkvm_host_map_guest_mmio,
 
 	/*
 	 * Start of the dynamically registered hypercalls. Start a bit

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

@@ -72,6 +72,7 @@ u64 __pkvm_ptdump_get_config(pkvm_handle_t handle, enum pkvm_ptdump_ops op);
 u64 __pkvm_ptdump_walk_range(pkvm_handle_t handle, struct pkvm_ptdump_log_hdr *log_hva);
 
 int pkvm_hyp_donate_guest(struct pkvm_hyp_vcpu *vcpu, u64 pfn, u64 gfn);
+int hyp_check_range_owned(u64 addr, u64 size);
 
 bool addr_is_memory(phys_addr_t phys);
 int host_stage2_idmap_locked(phys_addr_t addr, u64 size,

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

@@ -184,5 +184,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);
+int pkvm_host_map_guest_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 pfn, u64 gfn);
 
 #endif /* __ARM64_KVM_NVHE_PKVM_H__ */

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

@@ -57,6 +57,24 @@ static struct pkvm_device *pkvm_get_device(u64 addr, size_t size)
 	return NULL;
 }
 
+static struct pkvm_device *pkvm_get_device_by_addr(u64 addr)
+{
+	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) && (addr < res->base + 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
@@ -128,3 +146,91 @@ out_unlock:
 	hyp_spin_unlock(&device_spinlock);
 	return ret;
 }
+
+static int __pkvm_device_assign(struct pkvm_device *dev, struct pkvm_hyp_vm *vm)
+{
+	int i;
+	struct pkvm_dev_resource *res;
+	int ret;
+
+	hyp_assert_lock_held(&device_spinlock);
+
+	for (i = 0 ; i < dev->nr_resources; ++i) {
+		res = &dev->resources[i];
+		ret = hyp_check_range_owned(res->base, res->size);
+		if (ret)
+			return ret;
+	}
+
+	dev->ctxt = vm;
+	return 0;
+}
+
+/*
+ * Atomically check that all the group is assigned to the hypervisor
+ * and tag the devices in the group as owned by the VM.
+ * This can't race with reclaim as it's protected by device_spinlock
+ */
+static int __pkvm_group_assign(u32 group_id, struct pkvm_hyp_vm *vm)
+{
+	int i;
+	int ret = 0;
+
+	hyp_assert_lock_held(&device_spinlock);
+
+	for (i = 0 ; i < registered_devices_nr ; ++i) {
+		struct pkvm_device *dev = &registered_devices[i];
+
+		if (dev->group_id != group_id)
+			continue;
+		if (dev->ctxt) {
+			ret = -EPERM;
+			break;
+		}
+		ret = __pkvm_device_assign(dev, vm);
+		if (ret)
+			break;
+	}
+
+	if (ret) {
+		while (i--) {
+			struct pkvm_device *dev = &registered_devices[i];
+
+			if (dev->group_id == group_id)
+				dev->ctxt = NULL;
+		}
+	}
+	return ret;
+}
+
+
+int pkvm_host_map_guest_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 pfn, u64 gfn)
+{
+	int ret = 0;
+	struct pkvm_device *dev = pkvm_get_device_by_addr(hyp_pfn_to_phys(pfn));
+	struct pkvm_hyp_vm *vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
+
+	if (!dev)
+		return -ENODEV;
+
+	hyp_spin_lock(&device_spinlock);
+
+	if (dev->ctxt == NULL) {
+		/*
+		 * First time the device is assigned to a guest, make sure the whole
+		 * group is assigned to the hypervisor.
+		 */
+		ret = __pkvm_group_assign(dev->group_id, vm);
+	} else if (dev->ctxt != vm) {
+		ret = -EBUSY;
+	}
+
+	if (ret)
+		goto out_ret;
+
+	ret = pkvm_hyp_donate_guest(hyp_vcpu, pfn, gfn);
+
+out_ret:
+	hyp_spin_unlock(&device_spinlock);
+	return ret;
+}

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

@@ -1772,6 +1772,34 @@ static void handle___pkvm_host_reclaim_hyp_mmio(struct kvm_cpu_context *host_ctx
 	cpu_reg(host_ctxt, 1) = pkvm_device_reclaim_mmio(pfn, nr_pages);
 }
 
+static void handle___pkvm_host_map_guest_mmio(struct kvm_cpu_context *host_ctxt)
+{
+	DECLARE_REG(u64, pfn, host_ctxt, 1);
+	DECLARE_REG(u64, gfn, host_ctxt, 2);
+	struct pkvm_hyp_vcpu *hyp_vcpu;
+	int ret = -EINVAL;
+
+	if (!is_protected_kvm_enabled())
+		goto out;
+
+	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
+	if (!hyp_vcpu)
+		goto out;
+
+	if (!pkvm_hyp_vcpu_is_protected(hyp_vcpu))
+		goto out;
+
+	/* Top-up our per-vcpu memcache from the host's */
+	ret = pkvm_refill_memcache(hyp_vcpu);
+	if (ret)
+		goto out;
+
+	ret = pkvm_host_map_guest_mmio(hyp_vcpu, pfn, gfn);
+
+out:
+	cpu_reg(host_ctxt, 1) = ret;
+}
+
 typedef void (*hcall_t)(struct kvm_cpu_context *);
 
 #define HANDLE_FUNC(x)	[__KVM_HOST_SMCCC_FUNC_##x] = (hcall_t)handle_##x
@@ -1844,6 +1872,7 @@ static const hcall_t host_hcall[] = {
 	HANDLE_FUNC(__pkvm_host_iommu_map_sg),
 	HANDLE_FUNC(__pkvm_host_donate_hyp_mmio),
 	HANDLE_FUNC(__pkvm_host_reclaim_hyp_mmio),
+	HANDLE_FUNC(__pkvm_host_map_guest_mmio),
 };
 
 static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)

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

@@ -1058,6 +1058,18 @@ static int __hyp_check_page_state_range(u64 addr, u64 size,
 	return check_page_state_range(&pkvm_pgtable, addr, size, &d);
 }
 
+int hyp_check_range_owned(u64 phys_addr, u64 size)
+{
+	int ret;
+
+	hyp_lock_component();
+	ret = __hyp_check_page_state_range((u64)hyp_phys_to_virt(phys_addr),
+					   size, PKVM_PAGE_OWNED);
+	hyp_unlock_component();
+
+	return ret;
+}
+
 static enum pkvm_page_state guest_get_page_state(kvm_pte_t pte, u64 addr)
 {
 	enum pkvm_page_state state = 0;