Introduce a new iommufd option to mark an iommufd as persistent. For now
this allocates it a unique persistent ID from an xarray index and keeps
a reference to the domain.

This will be used so that at serialisation time the open iommufds can be
iterated through and serialised.
---
 drivers/iommu/iommufd/iommufd_private.h |  1 +
 drivers/iommu/iommufd/main.c            | 47 +++++++++++++++++++++++++
 include/uapi/linux/iommufd.h            |  5 +++
 3 files changed, 53 insertions(+)

diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index 92efe30a8f0d..b23f7766066c 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -28,6 +28,7 @@ struct iommufd_ctx {
 	/* Compatibility with VFIO no iommu */
 	u8 no_iommu_mode;
 	struct iommufd_ioas *vfio_ioas;
+	unsigned long persistent_id;
 };
 
 /*
diff --git a/drivers/iommu/iommufd/main.c b/drivers/iommu/iommufd/main.c
index 83bbd7c5d160..6708ad629b1e 100644
--- a/drivers/iommu/iommufd/main.c
+++ b/drivers/iommu/iommufd/main.c
@@ -29,6 +29,8 @@ struct iommufd_object_ops {
 static const struct iommufd_object_ops iommufd_object_ops[];
 static struct miscdevice vfio_misc_dev;
 
+static DEFINE_XARRAY_ALLOC(persistent_iommufds);
+
 struct iommufd_object *_iommufd_object_alloc(struct iommufd_ctx *ictx,
 					     size_t size,
 					     enum iommufd_object_type type)
@@ -287,10 +289,52 @@ static int iommufd_fops_release(struct inode *inode, struct file *filp)
 			break;
 	}
 	WARN_ON(!xa_empty(&ictx->groups));
+
+	rcu_read_lock();
+	if (ictx->persistent_id)
+		xa_erase(&persistent_iommufds, ictx->persistent_id);
+	rcu_read_unlock();
 	kfree(ictx);
 	return 0;
 }
 
+static int iommufd_option_persistent(struct iommufd_ucmd *ucmd)
+{
+	unsigned int persistent_id;
+	int rc;
+	struct iommu_option *cmd = ucmd->cmd;
+	struct iommufd_ctx *ictx = ucmd->ictx;
+	struct xa_limit id_limit = XA_LIMIT(1, UINT_MAX);
+
+	if (cmd->op == IOMMU_OPTION_OP_GET) {
+		cmd->val64 = ictx->persistent_id;
+		return 0;
+	}
+
+	if (cmd->op == IOMMU_OPTION_OP_SET) {
+		/*
+		 * iommufds can only be marked persistent before they
+		 * have been used for DMA mappings. HWPTs must be known
+		 * to be persistent at creation time.
+		 */
+		if (!xa_empty(&ictx->objects)) {
+			pr_warn("iommufd can only be marked persistented when unused\n");
+			return -EFAULT;
+		}
+
+		rc = xa_alloc(&persistent_iommufds, &persistent_id, ictx, id_limit, GFP_KERNEL_ACCOUNT);
+		if (rc) {
+			pr_warn("Unable to keep track of iommufd object\n");
+			return rc;
+		}
+
+		ictx->persistent_id = persistent_id;
+		cmd->val64 = ictx->persistent_id;
+		return 0;
+	}
+	return -EOPNOTSUPP;
+}
+
 static int iommufd_option(struct iommufd_ucmd *ucmd)
 {
 	struct iommu_option *cmd = ucmd->cmd;
@@ -306,6 +350,9 @@ static int iommufd_option(struct iommufd_ucmd *ucmd)
 	case IOMMU_OPTION_HUGE_PAGES:
 		rc = iommufd_ioas_option(ucmd);
 		break;
+	case IOMMU_OPTION_PERSISTENT:
+		rc = iommufd_option_persistent(ucmd);
+		break;
 	default:
 		return -EOPNOTSUPP;
 	}
diff --git a/include/uapi/linux/iommufd.h b/include/uapi/linux/iommufd.h
index 4dde745cfb7e..7d8cb242e9b0 100644
--- a/include/uapi/linux/iommufd.h
+++ b/include/uapi/linux/iommufd.h
@@ -276,10 +276,15 @@ struct iommu_ioas_unmap {
  *    iommu mappings. Value 0 disables combining, everything is mapped to
  *    PAGE_SIZE. This can be useful for benchmarking.  This is a per-IOAS
  *    option, the object_id must be the IOAS ID.
+ * @IOMMU_OPTION_PERSISTENT
+ *    Value 1 sets this iommufd object as a persistent iommufd. Mappings will
+ *    survive across kexec. The returned value is the persistent ID which can
+ *    be used to restore the iommufd after kexec.
  */
 enum iommufd_option {
 	IOMMU_OPTION_RLIMIT_MODE = 0,
 	IOMMU_OPTION_HUGE_PAGES = 1,
+	IOMMU_OPTION_PERSISTENT = 2,
 };
 
 /**
-- 
2.34.1