ANDROID: KVM: arm64: devices: Block IOMMU before and after assignment

DMA should be blocked when a device ownership changes, and all
the exisiting mapping should be removed to prevent attacks during
device initialization.

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

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

@@ -8,6 +8,7 @@
 #include <kvm/iommu.h>
 
 #include <nvhe/alloc_mgt.h>
+#include <nvhe/pkvm.h>
 
 /* alloc/free from atomic pool. */
 void *kvm_iommu_donate_pages_atomic(u8 order);
@@ -43,6 +44,8 @@ void kvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
 				 enum kvm_pgtable_prot prot);
 int kvm_iommu_snapshot_host_stage2(struct kvm_hyp_iommu_domain *domain);
 
+int kvm_iommu_dev_block_dma(pkvm_handle_t iommu_id, u32 endpoint_id, bool host_to_guest);
+
 struct kvm_iommu_ops {
 	int (*init)(void);
 	int (*alloc_domain)(struct kvm_hyp_iommu_domain *domain, int type);
@@ -66,6 +69,8 @@ struct kvm_iommu_ops {
 				  phys_addr_t start, phys_addr_t end, int prot);
 	int (*suspend)(struct kvm_hyp_iommu *iommu);
 	int (*resume)(struct kvm_hyp_iommu *iommu);
+	int (*dev_block_dma)(struct kvm_hyp_iommu *iommu, u32 endpoint_id,
+			     bool is_host_to_guest);
 	ANDROID_KABI_RESERVE(1);
 	ANDROID_KABI_RESERVE(2);
 	ANDROID_KABI_RESERVE(3);

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

@@ -3,6 +3,8 @@
  * Copyright (C) 2023 Google LLC
  * Author: Mostafa Saleh <smostafa@google.com>
  */
+
+#include <nvhe/iommu.h>
 #include <nvhe/mem_protect.h>
 #include <nvhe/mm.h>
 #include <nvhe/pkvm.h>
@@ -150,13 +152,27 @@ out_unlock:
 
 static int pkvm_device_reset(struct pkvm_device *dev, bool host_to_guest)
 {
+	struct pkvm_dev_iommu *iommu;
+	int ret;
+	int i;
+
 	hyp_assert_lock_held(&device_spinlock);
 
 	/* Reset is mandatory. */
 	if (!dev->reset_handler)
 		return -ENODEV;
 
-	return dev->reset_handler(dev->cookie, host_to_guest);
+	ret = dev->reset_handler(dev->cookie, host_to_guest);
+	if (ret)
+		return ret;
+
+	for (i = 0 ; i < dev->nr_iommus ; ++i) {
+		iommu = &dev->iommus[i];
+		ret = kvm_iommu_dev_block_dma(iommu->id, iommu->endpoint, host_to_guest);
+		if (WARN_ON(ret))
+			return ret;
+	}
+	return 0;
 }
 
 static int __pkvm_device_assign(struct pkvm_device *dev, struct pkvm_hyp_vm *vm)

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

@@ -538,6 +538,20 @@ out_put_domain:
 	return total_mapped;
 }
 
+int kvm_iommu_dev_block_dma(pkvm_handle_t iommu_id, u32 endpoint_id, bool host_to_guest)
+{
+	struct kvm_hyp_iommu *iommu;
+
+	if (!kvm_iommu_ops || !kvm_iommu_ops->dev_block_dma)
+		return -ENODEV;
+
+	iommu = kvm_iommu_ops->get_iommu_by_id(iommu_id);
+	if (!iommu)
+		return -ENOENT;
+
+	return kvm_iommu_ops->dev_block_dma(iommu, endpoint_id, host_to_guest);
+}
+
 static int iommu_power_on(struct kvm_power_domain *pd)
 {
 	struct kvm_hyp_iommu *iommu = container_of(pd, struct kvm_hyp_iommu,