[RFC PATCH v1 8/9] iommu/arm-smmu-v3: Inherit the ASID/VMID of restored domains

[Pranjal Shrivastava <praan@google.com>]

The IOMMU core restores a preserved domain by allocating a new paging
domain around the preserved page table and attaching it to the restored
device. The new domain gets a new ASID/VMID though, while the device is
still doing DMA through the preserved CD/STE.

Similar to the Intel driver reclaiming the preserved Domain ID, record
the ASID/VMID in the attachment_id and let a restored domain inherit the
live ASID/VMID on its first attach. Check the stage, the page table root
and the ID against the live CD/STE, failing the attach on a mismatch.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 .../arm/arm-smmu-v3/arm-smmu-v3-liveupdate.c  | 115 ++++++++++++++++++
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c   |   6 +
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h   |   8 ++
 3 files changed, 129 insertions(+)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-liveupdate.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-liveupdate.c
index 68652123d0e1..c0772bcb8d3a 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-liveupdate.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-liveupdate.c
@@ -273,6 +273,11 @@ int arm_smmu_preserve_device(struct device *dev,
 	/* Link this master to the preserved IOMMU domain in the ABI */
 	device_ser->domain_iommu_ser.domain_phys = virt_to_phys(domain_ser);
 
+	/* Record the ASID/VMID for the incoming kernel to cross-check */
+	device_ser->domain_iommu_ser.attachment_id =
+		smmu_domain->stage == ARM_SMMU_DOMAIN_S1 ?
+			smmu_domain->cd.asid : smmu_domain->s2_cfg.vmid;
+
 	/* If it's not Stage-1, or the CD table isn't allocated, we're done */
 	if (smmu_domain->stage != ARM_SMMU_DOMAIN_S1 ||
 	    !arm_smmu_cdtab_allocated(&master->cd_table))
@@ -781,4 +786,114 @@ int arm_smmu_liveupdate_restore_cd_tables(struct arm_smmu_master *master)
 	return ret;
 }
 
+/* Hand over the reserved ASID/VMID, releasing the never programmed one */
+static int arm_smmu_liveupdate_inherit_id(struct arm_smmu_domain *smmu_domain,
+					  u32 id)
+{
+	struct arm_smmu_device *smmu = smmu_domain->smmu;
+	int ret;
+
+	if (smmu_domain->stage == ARM_SMMU_DOMAIN_S1) {
+		if (smmu_domain->cd.asid == id)
+			return 0;
+
+		/* Only a reserved entry, which loads as NULL, can be taken */
+		if (xa_load(&smmu->asid_map, id))
+			return -EBUSY;
+		ret = xa_err(xa_store(&smmu->asid_map, id, smmu_domain,
+				      GFP_KERNEL));
+		if (ret)
+			return ret;
+		xa_erase(&smmu->asid_map, smmu_domain->cd.asid);
+		smmu_domain->cd.asid = id;
+		return 0;
+	}
+
+	if (smmu_domain->s2_cfg.vmid == id)
+		return 0;
+
+	/* The vmid_map already holds @id, simply transfer its ownership */
+	ida_free(&smmu->vmid_map, smmu_domain->s2_cfg.vmid);
+	smmu_domain->s2_cfg.vmid = id;
+	return 0;
+}
+
+/*
+ * Inherit the live ASID/VMID of a restored master, after checking that
+ * @smmu_domain matches its live translation. Called before
+ * arm_smmu_attach_prepare() builds the invalidation array.
+ */
+int arm_smmu_liveupdate_attach_restored(struct arm_smmu_master *master,
+					struct arm_smmu_domain *smmu_domain)
+{
+	struct iommu_device_ser *dev_ser = dev_iommu_restored_state(master->dev);
+	struct arm_smmu_device *smmu = master->smmu;
+	struct pt_iommu_armv8_hw_info info;
+	struct arm_smmu_ste *ste;
+	u64 ste0, ttb, live_ttb;
+	u32 id;
+
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	if (!dev_ser || !iommu_domain_restored_state(&smmu_domain->domain))
+		return 0;
+
+	ste = arm_smmu_get_step_for_sid(smmu, master->streams[0].id);
+	ste0 = le64_to_cpu(ste->data[0]);
+	if (!(ste0 & STRTAB_STE_0_V))
+		return 0;
+
+	pt_iommu_armv8_hw_info(&smmu_domain->armv8pt, &info);
+
+	switch (FIELD_GET(STRTAB_STE_0_CFG, ste0)) {
+	case STRTAB_STE_0_CFG_S1_TRANS: {
+		struct arm_smmu_cd *cdptr;
+		u64 cd0;
+
+		if (smmu_domain->stage != ARM_SMMU_DOMAIN_S1)
+			goto err_mismatch;
+
+		cdptr = arm_smmu_get_cd_ptr(master, IOMMU_NO_PASID);
+		if (!cdptr)
+			goto err_mismatch;
+
+		cd0 = le64_to_cpu(cdptr->data[0]);
+		if (!(cd0 & CTXDESC_CD_0_V))
+			goto err_mismatch;
+
+		id = FIELD_GET(CTXDESC_CD_0_ASID, cd0);
+		live_ttb = le64_to_cpu(cdptr->data[1]) & CTXDESC_CD_1_TTB0_MASK;
+		ttb = info.ttb & CTXDESC_CD_1_TTB0_MASK;
+		break;
+	}
+	case STRTAB_STE_0_CFG_S2_TRANS:
+		if (smmu_domain->stage != ARM_SMMU_DOMAIN_S2)
+			goto err_mismatch;
+
+		id = FIELD_GET(STRTAB_STE_2_S2VMID, le64_to_cpu(ste->data[2]));
+		live_ttb = le64_to_cpu(ste->data[3]) & STRTAB_STE_3_S2TTB_MASK;
+		ttb = info.ttb & STRTAB_STE_3_S2TTB_MASK;
+		break;
+	default:
+		/* Nothing to inherit, nested STEs aren't restored (yet) */
+		return 0;
+	}
+
+	if (live_ttb != ttb || id != dev_ser->domain_iommu_ser.attachment_id)
+		goto err_mismatch;
+
+	/* A shared restored domain inherits its ID on the first attach */
+	if (!list_empty(&smmu_domain->devices) &&
+	    id != (smmu_domain->stage == ARM_SMMU_DOMAIN_S1 ?
+		   smmu_domain->cd.asid : smmu_domain->s2_cfg.vmid))
+		goto err_mismatch;
+
+	return arm_smmu_liveupdate_inherit_id(smmu_domain, id);
+
+err_mismatch:
+	dev_err(master->dev,
+		"restored domain doesn't match the live translation\n");
+	return -EINVAL;
+}
+
 #endif
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
index 54c97ebc15b2..b371ea6b6009 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -3838,6 +3838,12 @@ static int arm_smmu_attach_dev(struct iommu_domain *domain, struct device *dev,
 	 */
 	mutex_lock(&arm_smmu_asid_lock);
 
+	ret = arm_smmu_liveupdate_attach_restored(master, smmu_domain);
+	if (ret) {
+		mutex_unlock(&arm_smmu_asid_lock);
+		return ret;
+	}
+
 	ret = arm_smmu_attach_prepare(&state, domain);
 	if (ret) {
 		mutex_unlock(&arm_smmu_asid_lock);
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index 0a88c0ec66ca..1b33f713f3b7 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1214,6 +1214,8 @@ void arm_smmu_unpreserve(struct iommu_device *iommu,
 int arm_smmu_liveupdate_shutdown(struct arm_smmu_device *smmu);
 int arm_smmu_liveupdate_restore_strtab(struct arm_smmu_device *smmu);
 int arm_smmu_liveupdate_restore_cd_tables(struct arm_smmu_master *master);
+int arm_smmu_liveupdate_attach_restored(struct arm_smmu_master *master,
+					struct arm_smmu_domain *smmu_domain);
 #else
 static inline int arm_smmu_liveupdate_shutdown(struct arm_smmu_device *smmu)
 {
@@ -1227,6 +1229,12 @@ static inline int arm_smmu_liveupdate_restore_cd_tables(struct arm_smmu_master *
 {
 	return 0;
 }
+static inline int
+arm_smmu_liveupdate_attach_restored(struct arm_smmu_master *master,
+				    struct arm_smmu_domain *smmu_domain)
+{
+	return 0;
+}
 #endif
 
 static inline void arm_smmu_domain_free(struct arm_smmu_domain *smmu_domain)
-- 
2.56.0.rc1.315.gc6ed9934b7-goog