ANDROID: KVM: arm64: iommu: Host stage-2 idmap callbacks

Add IOMMU callbacks for host stage-2 idmap changes.
'host_stage2_idmap_prepare' is called first and is expected to apply
the changes on the driver level, eg. update driver-specific page table
information. If successful, the generic code invokes
'host_stage2_idmap_apply' on each currently powered IOMMU device
associated with the driver to apply the changes.

Bug: 190463801
Change-Id: Ifcc063896f6e8967c332dbaa5b7e7f2ba138abbf
Signed-off-by: David Brazdil <dbrazdil@google.com>
(cherry picked from commit 4395ddff4b)
Signed-off-by: Mostafa Saleh <smostafa@google.com>
Signed-off-by: Quentin Perret <qperret@google.com>

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

@@ -25,6 +25,20 @@ struct pkvm_iommu_ops {
 	 */
 	int (*validate)(phys_addr_t base, size_t size);
 
+	/*
+	 * Callback to apply a host stage-2 mapping change at driver level.
+	 * Called before 'host_stage2_idmap_apply' with host lock held.
+	 */
+	void (*host_stage2_idmap_prepare)(phys_addr_t start, phys_addr_t end,
+					  enum kvm_pgtable_prot prot);
+
+	/*
+	 * Callback to apply a host stage-2 mapping change at device level.
+	 * Called after 'host_stage2_idmap_prepare' with host lock held.
+	 */
+	void (*host_stage2_idmap_apply)(struct pkvm_iommu *dev,
+					phys_addr_t start, phys_addr_t end);
+
 	/* Power management callbacks. Called with host lock held. */
 	int (*suspend)(struct pkvm_iommu *dev);
 	int (*resume)(struct pkvm_iommu *dev);
@@ -63,6 +77,8 @@ int pkvm_iommu_host_stage2_adjust_range(phys_addr_t addr, phys_addr_t *start,
 					phys_addr_t *end);
 bool pkvm_iommu_host_dabt_handler(struct kvm_cpu_context *host_ctxt, u32 esr,
 				  phys_addr_t fault_pa);
+void pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
+				  enum kvm_pgtable_prot prot);
 
 struct kvm_iommu_ops {
 	int (*init)(void);

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

@@ -380,3 +380,24 @@ bool pkvm_iommu_host_dabt_handler(struct kvm_cpu_context *host_ctxt, u32 esr,
 	}
 	return false;
 }
+
+void pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
+				  enum kvm_pgtable_prot prot)
+{
+	struct pkvm_iommu_driver *drv;
+	struct pkvm_iommu *dev;
+	size_t i;
+
+	assert_host_component_locked();
+
+	for (i = 0; i < ARRAY_SIZE(iommu_drivers); i++) {
+		drv = get_driver(i);
+		if (drv && is_driver_ready(drv) && drv->ops->host_stage2_idmap_prepare)
+			drv->ops->host_stage2_idmap_prepare(start, end, prot);
+	}
+
+	list_for_each_entry(dev, &iommu_list, list) {
+		if (dev->powered && dev->ops->host_stage2_idmap_apply)
+			dev->ops->host_stage2_idmap_apply(dev, start, end);
+	}
+}

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

@@ -416,9 +416,16 @@ int __pkvm_prot_finalize(void)
 
 int host_stage2_unmap_dev_locked(phys_addr_t start, u64 size)
 {
+	int ret;
+
 	hyp_assert_lock_held(&host_mmu.lock);
 
-	return kvm_pgtable_stage2_unmap(&host_mmu.pgt, start, size);
+	ret = kvm_pgtable_stage2_unmap(&host_mmu.pgt, start, size);
+	if (ret)
+		return ret;
+
+	pkvm_iommu_host_stage2_idmap(start, start + size, 0);
+	return 0;
 }
 
 static int host_stage2_unmap_dev_all(void)
@@ -508,8 +515,15 @@ static bool range_is_memory(u64 start, u64 end)
 static inline int __host_stage2_idmap(u64 start, u64 end,
 				      enum kvm_pgtable_prot prot)
 {
-	return kvm_pgtable_stage2_map(&host_mmu.pgt, start, end - start, start,
-				      prot, &host_s2_pool);
+	int ret;
+
+	ret = kvm_pgtable_stage2_map(&host_mmu.pgt, start, end - start, start,
+				     prot, &host_s2_pool);
+	if (ret)
+		return ret;
+
+	pkvm_iommu_host_stage2_idmap(start, end, prot);
+	return 0;
 }
 
 /*
@@ -585,6 +599,7 @@ static kvm_pte_t kvm_init_invalid_leaf_owner(u8 owner_id)
 int host_stage2_set_owner_locked(phys_addr_t addr, u64 size, u8 owner_id)
 {
 	kvm_pte_t annotation;
+	enum kvm_pgtable_prot prot;
 	int ret;
 
 	if (owner_id > KVM_MAX_OWNER_ID)
@@ -594,11 +609,12 @@ int host_stage2_set_owner_locked(phys_addr_t addr, u64 size, u8 owner_id)
 
 	ret = host_stage2_try(kvm_pgtable_stage2_annotate, &host_mmu.pgt,
 			      addr, size, &host_s2_pool, annotation);
+	if (ret)
+		return ret;
 
-	if (!ret && kvm_iommu_ops.host_stage2_set_owner)
-		kvm_iommu_ops.host_stage2_set_owner(addr, size, owner_id);
-
-	return ret;
+	prot = owner_id == PKVM_ID_HOST ? PKVM_HOST_MEM_PROT : 0;
+	pkvm_iommu_host_stage2_idmap(addr, addr + size, prot);
+	return 0;
 }
 
 static bool host_stage2_force_pte_cb(u64 addr, u64 end, enum kvm_pgtable_prot prot)