ANDROID: KVM: arm64: iommu: Snapshot host stage-2

Snapshot the host stage-2 in the idmapped domain, this is done
differently from the Android14:

- Instead of assuming that the page table maps all the memory and
  the snapshot function would only unmap the donated pages, now the
  polarity is the opposite, where the page table is assumed to be
  empty and the snapshot function would map all host memory, so this
  slower, but done only once at boot, and that won't require the
  module to know the memory map.

- Instead of calling snapshot at init, now it is exported to modules,
  as some IOMMUs as SMMUv3 doens't know the format of the page table
  until the device is attached.

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

arch/arm64/include/asm/kvm_pkvm_module.h

@@ -206,6 +206,7 @@ struct pkvm_module_ops {
 	typeof(__list_add_valid_or_report) *list_add_valid_or_report;
 	typeof(__list_del_entry_valid_or_report) *list_del_entry_valid_or_report;
 #endif
+	int (*iommu_snapshot_host_stage2)(struct kvm_hyp_iommu_domain *domain);
 	ANDROID_KABI_RESERVE(1);
 	ANDROID_KABI_RESERVE(2);
 	ANDROID_KABI_RESERVE(3);

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

@@ -35,6 +35,7 @@ void kvm_iommu_reclaim_pages(void *p, u8 order);
 
 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);
 
 struct kvm_iommu_ops {
 	int (*init)(void);
@@ -55,6 +56,8 @@ struct kvm_iommu_ops {
 	void (*iotlb_sync)(struct kvm_hyp_iommu_domain *domain,
 			   struct iommu_iotlb_gather *gather);
 	bool (*dabt_handler)(struct kvm_cpu_context *host_ctxt, u64 esr, u64 addr);
+	void (*host_stage2_idmap)(struct kvm_hyp_iommu_domain *domain,
+				  phys_addr_t start, phys_addr_t end, int prot);
 };
 
 int kvm_iommu_init(struct kvm_iommu_ops *ops);

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

@@ -28,6 +28,18 @@ DECLARE_PER_CPU(struct kvm_hyp_req, host_hyp_reqs);
 /* Protects domains in kvm_hyp_iommu_domains */
 static DEFINE_HYP_SPINLOCK(kvm_iommu_domain_lock);
 
+static atomic_t kvm_iommu_idmap_initialized;
+
+static inline void kvm_iommu_idmap_init_done(void)
+{
+	atomic_set_release(&kvm_iommu_idmap_initialized, 1);
+}
+
+static inline bool kvm_iommu_is_ready(void)
+{
+	return atomic_read_acquire(&kvm_iommu_idmap_initialized) == 1;
+}
+
 static int kvm_iommu_refill(struct kvm_hyp_memcache *host_mc)
 {
 	if (!kvm_iommu_ops)
@@ -100,7 +112,7 @@ void kvm_iommu_reclaim_pages(void *p, u8 order)
 }
 
 static struct kvm_hyp_iommu_domain *
-handle_to_domain(pkvm_handle_t domain_id)
+__handle_to_domain(pkvm_handle_t domain_id, bool alloc)
 {
 	int idx;
 	struct kvm_hyp_iommu_domain *domains;
@@ -112,6 +124,8 @@ handle_to_domain(pkvm_handle_t domain_id)
 	idx = domain_id / KVM_IOMMU_DOMAINS_PER_PAGE;
 	domains = (struct kvm_hyp_iommu_domain *)READ_ONCE(kvm_hyp_iommu_domains[idx]);
 	if (!domains) {
+		if (!alloc)
+			return NULL;
 		domains = kvm_iommu_donate_page();
 		if (!domains)
 			return NULL;
@@ -131,6 +145,12 @@ handle_to_domain(pkvm_handle_t domain_id)
 	return &domains[domain_id % KVM_IOMMU_DOMAINS_PER_PAGE];
 }
 
+static struct kvm_hyp_iommu_domain *
+handle_to_domain(pkvm_handle_t domain_id)
+{
+	return __handle_to_domain(domain_id, true);
+}
+
 static int domain_get(struct kvm_hyp_iommu_domain *domain)
 {
 	int old = atomic_fetch_inc_acquire(&domain->refs);
@@ -433,7 +453,71 @@ int kvm_iommu_init_device(struct kvm_hyp_iommu *iommu)
 	return pkvm_init_power_domain(&iommu->power_domain, &iommu_power_ops);
 }
 
+static inline int pkvm_to_iommu_prot(int prot)
+{
+	switch (prot) {
+	case PKVM_HOST_MEM_PROT:
+		return IOMMU_READ | IOMMU_WRITE;
+	case PKVM_HOST_MMIO_PROT:
+		return IOMMU_READ | IOMMU_WRITE | IOMMU_MMIO;
+	case 0:
+		return 0;
+	default:
+		/* We don't understand that, it might cause corruption, so panic. */
+		BUG();
+	}
+
+	return 0;
+}
+
 void kvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
 				 enum kvm_pgtable_prot prot)
 {
+	struct kvm_hyp_iommu_domain *domain;
+
+	if (!kvm_iommu_is_ready())
+		return;
+
+	domain = __handle_to_domain(KVM_IOMMU_DOMAIN_IDMAP_ID, false);
+
+	kvm_iommu_ops->host_stage2_idmap(domain, start, end, pkvm_to_iommu_prot(prot));
+}
+
+static int __snapshot_host_stage2(const struct kvm_pgtable_visit_ctx *ctx,
+				  enum kvm_pgtable_walk_flags visit)
+{
+	u64 start = ctx->addr;
+	kvm_pte_t pte = *ctx->ptep;
+	u32 level = ctx->level;
+	struct kvm_hyp_iommu_domain *domain = ctx->arg;
+	u64 end = start + kvm_granule_size(level);
+	int prot = IOMMU_READ | IOMMU_WRITE;
+
+	if (!addr_is_memory(start))
+		prot |= IOMMU_MMIO;
+
+	if (!pte || kvm_pte_valid(pte))
+		kvm_iommu_ops->host_stage2_idmap(domain, start, end, prot);
+
+	return 0;
+}
+
+int kvm_iommu_snapshot_host_stage2(struct kvm_hyp_iommu_domain *domain)
+{
+	int ret;
+	struct kvm_pgtable_walker walker = {
+		.cb	= __snapshot_host_stage2,
+		.flags	= KVM_PGTABLE_WALK_LEAF,
+		.arg = domain,
+	};
+	struct kvm_pgtable *pgt = &host_mmu.pgt;
+
+	hyp_spin_lock(&host_mmu.lock);
+	ret = kvm_pgtable_walk(pgt, 0, BIT(pgt->ia_bits), &walker);
+	/* Start receiving calls to host_stage2_idmap. */
+	if (!ret)
+		kvm_iommu_idmap_init_done();
+	hyp_spin_unlock(&host_mmu.lock);
+
+	return ret;
 }

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

@@ -165,6 +165,7 @@ const struct pkvm_module_ops module_ops = {
 	.list_add_valid_or_report = __list_add_valid_or_report,
 	.list_del_entry_valid_or_report = __list_del_entry_valid_or_report,
 #endif
+	.iommu_snapshot_host_stage2 = kvm_iommu_snapshot_host_stage2,
 };
 
 int __pkvm_init_module(void *module_init)

drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c

@@ -229,7 +229,7 @@ static int kvm_arm_smmu_domain_finalize(struct kvm_arm_smmu_domain *kvm_smmu_dom
 	 * unique, and it has to be in range of this smmu, which can be
 	 * either 8 or 16 bits.
 	 */
-	ret = ida_alloc_range(&kvm_arm_smmu_domain_ida, 0,
+	ret = ida_alloc_range(&kvm_arm_smmu_domain_ida, KVM_IOMMU_DOMAIN_NR_START,
 			      min(KVM_IOMMU_MAX_DOMAINS, max_domains), GFP_KERNEL);
 	if (ret < 0)
 		return ret;

include/kvm/iommu.h

@@ -9,6 +9,17 @@
 #include <nvhe/spinlock.h>
 #endif
 
+/*
+ * Domain ID for identity mapped domain that the host can attach
+ * to get the same mapping available to the CPU page table.
+ */
+#define KVM_IOMMU_DOMAIN_IDMAP_ID		0
+
+/* Used in alloc_domain type argument. */
+#define KVM_IOMMU_DOMAIN_IDMAP_TYPE		0
+
+#define KVM_IOMMU_DOMAIN_NR_START		(KVM_IOMMU_DOMAIN_IDMAP_ID + 1)
+
 struct kvm_hyp_iommu_domain {
 	atomic_t		refs;
 	pkvm_handle_t		domain_id;