ANDROID: KVM: arm64: devices: Check host ownership for IOMMU calls

As currently we support pVMs, host shouldn't be able to attach or
detach any device owned by a pVM.

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

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

@@ -164,6 +164,8 @@ int pkvm_init_scmi_pd(struct kvm_power_domain *pd,
 
 bool pkvm_device_request_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 *exit_code);
 void pkvm_devices_teardown(struct pkvm_hyp_vm *vm);
+int pkvm_devices_get_context(u64 iommu_id, u32 endpoint_id);
+void pkvm_devices_put_context(u64 iommu_id, u32 endpoint_id);
 
 /*
  * Register a power domain. When the hypervisor catches power requests from the

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

@@ -99,7 +99,7 @@ int pkvm_device_hyp_assign_mmio(u64 pfn, u64 nr_pages)
 
 	hyp_spin_lock(&device_spinlock);
 	/* A VM already have this device, no take backs. */
-	if (dev->ctxt) {
+	if (dev->ctxt || dev->refcount) {
 		ret = -EBUSY;
 		goto out_unlock;
 	}
@@ -184,7 +184,7 @@ static int __pkvm_group_assign(u32 group_id, struct pkvm_hyp_vm *vm)
 
 		if (dev->group_id != group_id)
 			continue;
-		if (dev->ctxt) {
+		if (dev->ctxt || dev->refcount) {
 			ret = -EPERM;
 			break;
 		}
@@ -331,3 +331,52 @@ void pkvm_devices_teardown(struct pkvm_hyp_vm *vm)
 	}
 	hyp_spin_unlock(&device_spinlock);
 }
+
+static struct pkvm_device *pkvm_get_device_by_iommu(u64 id, u32 endpoint_id)
+{
+	struct pkvm_device *dev = NULL;
+	struct pkvm_dev_iommu *iommu;
+	int i, j;
+
+	for (i = 0 ; i < registered_devices_nr ; ++i) {
+		dev = &registered_devices[i];
+		for (j = 0 ; j < dev->nr_iommus; ++j) {
+			iommu = &dev->iommus[j];
+			if ((id == iommu->id) && (endpoint_id == iommu->endpoint))
+				return dev;
+		}
+	}
+
+	return NULL;
+}
+
+int pkvm_devices_get_context(u64 iommu_id, u32 endpoint_id)
+{
+	struct pkvm_device *dev = pkvm_get_device_by_iommu(iommu_id, endpoint_id);
+	int ret = 0;
+
+	if (!dev)
+		return 0;
+
+	hyp_spin_lock(&device_spinlock);
+	if (dev->ctxt)
+		ret = -EPERM;
+	else
+		hyp_refcount_inc(dev->refcount);
+	hyp_spin_unlock(&device_spinlock);
+	return ret;
+}
+
+void pkvm_devices_put_context(u64 iommu_id, u32 endpoint_id)
+{
+	struct pkvm_device *dev = pkvm_get_device_by_iommu(iommu_id, endpoint_id);
+
+	if (!dev)
+		return;
+
+	hyp_spin_lock(&device_spinlock);
+	BUG_ON(dev->ctxt);
+	hyp_refcount_dec(dev->refcount);
+
+	hyp_spin_unlock(&device_spinlock);
+}

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

@@ -10,6 +10,7 @@
 #include <hyp/adjust_pc.h>
 
 #include <kvm/iommu.h>
+#include <kvm/device.h>
 
 #include <nvhe/iommu.h>
 #include <nvhe/mem_protect.h>
@@ -298,13 +299,26 @@ int kvm_iommu_attach_dev(pkvm_handle_t iommu_id, pkvm_handle_t domain_id,
 	if (!iommu)
 		return -EINVAL;
 
+	/*
+	 * At the moment the IOMMU in EL2 is not aware of guests and pvIOMMU
+	 * doesn't exist yet, so all attaches come from host, this should change soon.
+	 */
+	ret = pkvm_devices_get_context(iommu_id, endpoint_id);
+	if (ret)
+		return ret;
+
 	domain = handle_to_domain(domain_id);
-	if (!domain || domain_get(domain))
-		return -EINVAL;
+	if (!domain || domain_get(domain)) {
+		ret = -EINVAL;
+		goto out_unlock;
+	}
 
 	ret = kvm_iommu_ops->attach_dev(iommu, domain, endpoint_id, pasid, pasid_bits);
 	if (ret)
 		domain_put(domain);
+
+out_unlock:
+	pkvm_devices_put_context(iommu_id, endpoint_id);
 	return ret;
 }
 
@@ -319,14 +333,25 @@ int kvm_iommu_detach_dev(pkvm_handle_t iommu_id, pkvm_handle_t domain_id,
 	if (!iommu)
 		return -EINVAL;
 
+	/* See kvm_iommu_attach_dev(). */
+	ret = pkvm_devices_get_context(iommu_id, endpoint_id);
+	if (ret)
+		return ret;
+
 	domain = handle_to_domain(domain_id);
-	if (!domain || atomic_read(&domain->refs) <= 1)
-		return -EINVAL;
+	if (!domain || atomic_read(&domain->refs) <= 1) {
+		ret = -EINVAL;
+		goto out_unlock;
+	}
 
 	ret = kvm_iommu_ops->detach_dev(iommu, domain, endpoint_id, pasid);
 	if (ret)
-		return ret;
+		goto out_unlock;
+
 	domain_put(domain);
+
+out_unlock:
+	pkvm_devices_put_context(iommu_id, endpoint_id);
 	return ret;
 }
 

include/kvm/device.h

@@ -37,6 +37,7 @@ struct pkvm_device {
 	u32 nr_iommus;
 	u32 group_id;
 	void *ctxt; /* Current context of the device*/
+	unsigned short refcount;
 };
 
 #endif /* #ifndef __KVM_DEVICE_H */