[PATCH v3 11/19] iommu/riscv: Copy MSI IOVA table when replacing an iommufd domain

iommufd may replace one paging domain with another on a device while
leaving the device otherwise operational, e.g. when converting a
device's HWPT. If the old domain already had a populated MSI IOVA
table then the new domain must inherit that table rather than wait
to build its own, since the MSI IOVAs are only built on the next
irq_domain_alloc_irqs(), which does not happen on a domain
replacement.

Only copy from an old domain that is a genuine RISC-V paging domain;
iommu.c already has riscv_iommu_paging_domain_ops in scope to check
this, so do the check there and pass NULL down on a mismatch. This
keeps riscv_iommu_ir_attach_paging_domain() simple: a non-NULL old is
always safe to cast with iommu_domain_to_riscv().

Only iommufd cookie domains are copied from and to, since VFIO type1
and DMA API domains build their own tables from irq_domain_alloc_irqs()
before any device is attached, and are never replaced while a device is
live.

Signed-off-by: Andrew Jones <[email protected]>
---
 drivers/iommu/riscv/iommu-ir.c | 78 ++++++++++++++++++++++++++++++++--
 drivers/iommu/riscv/iommu.h    |  2 +
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/drivers/iommu/riscv/iommu-ir.c b/drivers/iommu/riscv/iommu-ir.c
index 0dff74aa9b18..7d9a1eaca92e 100644
--- a/drivers/iommu/riscv/iommu-ir.c
+++ b/drivers/iommu/riscv/iommu-ir.c
@@ -118,18 +118,26 @@ static int riscv_iommu_ir_irq_domain_alloc_irqs(struct irq_domain *irqdomain,
 	 * quiesced, including MSI teardown, before switching away from or freeing
 	 * the domain. iommu_dma_map_msi() requires the group mutex to be held;
 	 * take it around the domain lookup too so info->domain can't change
-	 * out from under the build.
+	 * out from under the build. Bump info->nr_msis here too, before
+	 * irq_domain_alloc_irqs_parent() runs unlocked below, so a concurrent
+	 * riscv_iommu_ir_attach_paging_domain() can never observe a count that
+	 * is lower than the number of MSIs actually in flight for this device.
 	 */
 	scoped_guard(iommu_group, info->dev) {
 		domain = rcu_dereference_protected(info->domain, true);
 		ret = domain ? riscv_iommu_ir_build_msi_iova(domain, info->dev) : 0;
+		if (!ret)
+			info->nr_msis += nr_irqs;
 	}
 	if (ret)
 		return ret;
 
 	ret = irq_domain_alloc_irqs_parent(irqdomain, irq_base, nr_irqs, arg);
-	if (ret)
+	if (ret) {
+		guard(iommu_group)(info->dev);
+		info->nr_msis -= nr_irqs;
 		return ret;
+	}
 
 	for (i = 0; i < nr_irqs; i++) {
 		data = irq_domain_get_irq_data(irqdomain, irq_base + i);
@@ -139,9 +147,25 @@ static int riscv_iommu_ir_irq_domain_alloc_irqs(struct irq_domain *irqdomain,
 	return 0;
 }
 
+static void riscv_iommu_ir_irq_domain_free_irqs(struct irq_domain *irqdomain,
+						unsigned int irq_base, unsigned int nr_irqs)
+{
+	struct riscv_iommu_info *info = irqdomain->host_data;
+
+	irq_domain_free_irqs_parent(irqdomain, irq_base, nr_irqs);
+
+	/*
+	 * Decrement only after the parent free completes, so a concurrent
+	 * riscv_iommu_ir_attach_paging_domain() never observes a count lower
+	 * than the number of MSIs that are actually still live.
+	 */
+	scoped_guard(iommu_group, info->dev)
+		info->nr_msis -= nr_irqs;
+}
+
 static const struct irq_domain_ops riscv_iommu_ir_irq_domain_ops = {
 	.alloc = riscv_iommu_ir_irq_domain_alloc_irqs,
-	.free = irq_domain_free_irqs_parent,
+	.free = riscv_iommu_ir_irq_domain_free_irqs,
 };
 
 static const struct msi_parent_ops riscv_iommu_ir_msi_parent_ops = {
@@ -213,6 +237,54 @@ void riscv_iommu_ir_irq_domain_remove(struct device *dev, struct riscv_iommu_inf
 int riscv_iommu_ir_attach_paging_domain(struct iommu_domain *iommu_domain, struct device *dev,
 					struct iommu_domain *old)
 {
+	struct riscv_iommu_domain *domain = iommu_domain_to_riscv(iommu_domain);
+	struct riscv_iommu_info *info = dev_iommu_priv_get(dev);
+	struct riscv_iommu_domain *old_domain = NULL;
+	dma_addr_t *msi_iova = NULL;
+
+	if (old && (old->type & __IOMMU_DOMAIN_PAGING))
+		old_domain = iommu_domain_to_riscv(old);
+
+	/*
+	 * Copying is only correct between two IOMMUFD domains: their MSI IOVAs
+	 * come from the fd-wide SW_MSI reservation, so they match across
+	 * domain instances. Every other cookie type derives its MSI IOVAs from
+	 * domain-local allocator state.
+	 */
+	if (old_domain && old_domain->domain.cookie_type == IOMMU_COOKIE_IOMMUFD &&
+	    iommu_domain->cookie_type == IOMMU_COOKIE_IOMMUFD) {
+		scoped_guard(mutex, &old_domain->mutex) {
+			if (old_domain->msi_iova) {
+				msi_iova = kmemdup(old_domain->msi_iova,
+						   riscv_iommu_ir_msi_iova_count() *
+						   sizeof(*msi_iova),
+						   GFP_KERNEL);
+				if (!msi_iova)
+					return -ENOMEM;
+			}
+		}
+
+		if (msi_iova) {
+			guard(mutex)(&domain->mutex);
+
+			if (domain->msi_iova)
+				kfree(msi_iova);
+			else
+				domain->msi_iova = msi_iova;
+
+			return 0;
+		}
+	}
+
+	/*
+	 * No table to copy: build one from scratch if this device has ever
+	 * allocated MSIs, since those MSIs may already be live and expecting
+	 * riscv_iommu_ir_compose_msi_msg() to find a populated table for
+	 * whatever domain is now attached.
+	 */
+	if (info->nr_msis)
+		return riscv_iommu_ir_build_msi_iova(domain, dev);
+
 	return 0;
 }
 
diff --git a/drivers/iommu/riscv/iommu.h b/drivers/iommu/riscv/iommu.h
index 8bd41dd63f0e..77328282a236 100644
--- a/drivers/iommu/riscv/iommu.h
+++ b/drivers/iommu/riscv/iommu.h
@@ -46,6 +46,8 @@ struct riscv_iommu_info {
 	struct riscv_iommu_domain __rcu *domain;
 	struct irq_domain *irqdomain;
 	struct device *dev;
+	/* Number of currently allocated MSIs; protected by the group mutex */
+	unsigned int nr_msis;
 };
 
 struct riscv_iommu_device;
-- 
2.43.0