[PATCH v7 17/18] vfio/pci: Preserve the iommufd state of the vfio cdev

[Samiullah Khawaja <skhawaja@google.com>]

If the vfio cdev is attached to an iommufd, preserve the state of the
attached iommufd also. Basically preserve the iommu specific state of
the device and also the attach iommu HW unit.

Once the device and its iommufd attachment is preserved, it cannot be
detached or attached to another IOAS until it is unpreserved.

VFIO PCI live update currently supports only devices in noiommu mode.
Lift that restriction for IOMMU-backed devices that are preserved with
their iommufd state.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
Reviewed-by: Pranjal Shrivastava <praan@google.com>
Reviewed-by: Lu Baolu <baolu.lu@linux.intel.com>
---
 drivers/vfio/device_cdev.c             |  10 ++
 drivers/vfio/pci/Kconfig               |  13 ++-
 drivers/vfio/pci/vfio_pci_liveupdate.c | 129 +++++++++++++++++++------
 drivers/vfio/pci/vfio_pci_priv.h       |  12 ++-
 include/linux/kho/abi/vfio_pci.h       |  15 ++-
 5 files changed, 141 insertions(+), 38 deletions(-)

diff --git a/drivers/vfio/device_cdev.c b/drivers/vfio/device_cdev.c
index a53922adc388..edcec2824e18 100644
--- a/drivers/vfio/device_cdev.c
+++ b/drivers/vfio/device_cdev.c
@@ -316,6 +316,11 @@ int vfio_df_ioctl_attach_pt(struct vfio_device_file *df,
 	}
 
 	mutex_lock(&device->dev_set->lock);
+	if (iommufd_device_is_preserved(device->iommufd_device)) {
+		ret = -EBUSY;
+		goto out_unlock;
+	}
+
 	if (attach.flags & VFIO_DEVICE_ATTACH_PASID)
 		ret = device->ops->pasid_attach_ioas(device,
 						     attach.pasid,
@@ -374,6 +379,11 @@ int vfio_df_ioctl_detach_pt(struct vfio_device_file *df,
 	}
 
 	mutex_lock(&device->dev_set->lock);
+	if (iommufd_device_is_preserved(device->iommufd_device)) {
+		mutex_unlock(&device->dev_set->lock);
+		return -EBUSY;
+	}
+
 	if (detach.flags & VFIO_DEVICE_DETACH_PASID)
 		device->ops->pasid_detach_ioas(device, detach.pasid);
 	else
diff --git a/drivers/vfio/pci/Kconfig b/drivers/vfio/pci/Kconfig
index 71289d8908f6..a9884d3ec864 100644
--- a/drivers/vfio/pci/Kconfig
+++ b/drivers/vfio/pci/Kconfig
@@ -44,12 +44,15 @@ config VFIO_PCI_IGD
 	  To enable Intel IGD assignment through vfio-pci, say Y.
 
 config VFIO_PCI_LIVEUPDATE
-	bool "VFIO PCI support for Live Update in no-iommu mode"
-	depends on PCI_LIVEUPDATE && VFIO_NOIOMMU && VFIO_DEVICE_CDEV
+	bool "VFIO PCI support for Live Update"
+	depends on PCI_LIVEUPDATE && VFIO_DEVICE_CDEV
+	depends on VFIO_NOIOMMU || IOMMU_LIVEUPDATE
 	help
-	  Support for preserving devices bound to vfio-pci in no-iommu mode
-	  across a Live Update. It enables a device to continue working (DMA)
-	  while host kernel is going through a kexec live update.
+	  Support for preserving devices bound to vfio-pci across a Live
+	  Update. It enables a device to continue working (DMA) while host
+	  kernel is going through a kexec live update. Devices in no-iommu
+	  mode need VFIO_NOIOMMU, and IOMMU-backed devices need
+	  IOMMU_LIVEUPDATE to preserve their IOMMU state.
 
 	  If you don't know what to do here, say N.
 
diff --git a/drivers/vfio/pci/vfio_pci_liveupdate.c b/drivers/vfio/pci/vfio_pci_liveupdate.c
index 00c008117259..86f532e4c9f1 100644
--- a/drivers/vfio/pci/vfio_pci_liveupdate.c
+++ b/drivers/vfio/pci/vfio_pci_liveupdate.c
@@ -9,15 +9,16 @@
 /**
  * DOC: VFIO PCI Preservation via LUO
  *
- * VFIO PCI devices in noiommu mode can be preserved over a kexec using the Live
+ * VFIO PCI devices can be preserved over a kexec using the Live
  * Update Orchestrator (LUO) file preservation while keeping the device running
  * and performing DMA across the kexec transition. This allows userspace (such
  * as a VMM or userspace driver) to transfer an in-use device to the next kernel
  * without resetting it.
  *
  * .. note::
- *    This support is enabled by ``CONFIG_VFIO_PCI_LIVEUPDATE`` and is limited
- *    to devices in noiommu mode (see Restrictions below).
+ *    This support is enabled by ``CONFIG_VFIO_PCI_LIVEUPDATE``. Devices in
+ *    noiommu mode are supported. IOMMU-backed devices are supported when
+ *    ``CONFIG_IOMMU_LIVEUPDATE`` is enabled (see Restrictions below).
  *
  * Usage Example
  * =============
@@ -54,15 +55,22 @@
  *
  * The kernel imposes the following restrictions when preserving VFIO devices:
  *
- *  * The device must be bound to the ``vfio-pci`` driver (not a variant driver)
- *    and registered in **noiommu** mode (``vfio.enable_unsafe_noiommu_mode=1``
- *    with no IOMMU group attached to the device). Attempting to preserve an
- *    IOMMU-backed device fails with ``-EOPNOTSUPP``.
+ *  * The device must be bound to the ``vfio-pci`` driver (not a variant driver).
+ *
+ *  * A device in **noiommu** mode (``vfio.enable_unsafe_noiommu_mode=1`` with
+ *    no IOMMU group attached to the device) is preserved on its own.
+ *
+ *  * An IOMMU-backed device is preserved together with its iommufd state. The
+ *    device must be attached to a HWPT marked for preservation with
+ *    ``IOMMU_HWPT_LIVEUPDATE_MARK_PRESERVE``, and the iommufd must be preserved
+ *    in the same session before the device. Otherwise preserving the device
+ *    fails. After kexec, the device can only be retrieved if the iommufd it
+ *    was preserved with can be retrieved.
  *
  *  * The device file descriptor must be obtained by opening the VFIO character
- *    device (``/dev/vfio/devices/noiommu-vfioX``) and not via
- *    ``VFIO_GROUP_GET_DEVICE_FD``, and ``VFIO_DEVICE_BIND_IOMMUFD`` must have
- *    completed on that file descriptor before it is preserved.
+ *    device (``/dev/vfio/devices/vfioX`` or ``/dev/vfio/devices/noiommu-vfioX``)
+ *    and not via ``VFIO_GROUP_GET_DEVICE_FD``, and ``VFIO_DEVICE_BIND_IOMMUFD``
+ *    must have completed on that file descriptor before it is preserved.
  *
  *  * ``CONFIG_VFIO_PCI_ZDEV_KVM`` must not be enabled. This may be relaxed in
  *    the future.
@@ -189,6 +197,7 @@
 #include <linux/cleanup.h>
 #include <linux/errno.h>
 #include <linux/file.h>
+#include <linux/iommufd.h>
 #include <linux/kexec_handover.h>
 #include <linux/kho/abi/vfio_pci.h>
 #include <linux/liveupdate.h>
@@ -202,6 +211,8 @@
 
 #include "vfio_pci_priv.h"
 
+MODULE_IMPORT_NS("IOMMUFD");
+
 static bool vfio_pci_liveupdate_can_preserve(struct liveupdate_file_handler *handler,
 					     struct file *file)
 {
@@ -245,12 +256,11 @@ static int vfio_pci_liveupdate_preserve(struct liveupdate_file_op_args *args)
 	struct vfio_pci_core_device *vdev;
 	size_t saved_size = 0, total_size;
 	struct pci_dev *pdev;
+	u64 iommufd_token = 0;
+	u32 flags = 0;
 	u16 cfg_size;
 	int err;
 
-	if (!vfio_device_is_noiommu(device))
-		return -EOPNOTSUPP;
-
 	vdev = container_of(device, struct vfio_pci_core_device, vdev);
 	pdev = vdev->pdev;
 	cfg_size = pdev->cfg_size;
@@ -263,56 +273,89 @@ static int vfio_pci_liveupdate_preserve(struct liveupdate_file_op_args *args)
 	if (!vfio_file_is_bound_cdev(args->file))
 		return -EINVAL;
 
+	/*
+	 * Devices that are not in noiommu mode are preserved together with
+	 * their iommufd state, so that their DMA keeps being translated across
+	 * kexec.
+	 */
+	if (!vfio_device_is_noiommu(device)) {
+		err = iommufd_device_preserve(args->session,
+					      device->iommufd_device,
+					      &iommufd_token);
+		if (err)
+			return err;
+
+		flags |= VFIO_PCI_SER_IOMMUFD;
+	}
+
 	if (vdev->pci_saved_state) {
 		saved_size = pci_liveupdate_saved_state_ser_size(vdev->pci_saved_state);
-		if (!saved_size)
-			return -EINVAL;
+		if (!saved_size) {
+			err = -EINVAL;
+			goto err_iommufd_unpreserve;
+		}
 	}
 
 	total_size = size_add(sizeof(*ser),
 			      size_add(size_mul(2, cfg_size), saved_size));
-	if (total_size > U32_MAX)
-		return -EOVERFLOW;
+	if (total_size > U32_MAX) {
+		err = -EOVERFLOW;
+		goto err_iommufd_unpreserve;
+	}
 
 	ser = kho_alloc_preserve(total_size);
-	if (IS_ERR(ser))
-		return PTR_ERR(ser);
+	if (IS_ERR(ser)) {
+		err = PTR_ERR(ser);
+		goto err_iommufd_unpreserve;
+	}
 
 	if (saved_size) {
 		void *dst = ser->data + 2 * (size_t)cfg_size;
 
 		err = pci_liveupdate_serialize_saved_state(vdev->pci_saved_state,
 							   dst, saved_size);
-		if (err) {
-			kho_unpreserve_free(ser);
-			return err;
-		}
+		if (err)
+			goto err_free;
 	}
 
 	err = pci_liveupdate_preserve(pdev);
-	if (err) {
-		kho_unpreserve_free(ser);
-		return err;
-	}
+	if (err)
+		goto err_free;
 
 	ser->domain = pci_domain_nr(pdev->bus);
 	ser->bdf = pci_dev_id(pdev);
 	ser->cfg_size = cfg_size;
 	ser->size = total_size;
-	ser->flags = vdev->reset_works ? VFIO_PCI_SER_RESET_WORKS : 0;
+	ser->flags = flags;
+	if (vdev->reset_works)
+		ser->flags |= VFIO_PCI_SER_RESET_WORKS;
 	ser->saved_state_size = saved_size;
 	ser->reserved = 0;
+	ser->iommufd_token = iommufd_token;
 
 	args->serialized_data = virt_to_phys(ser);
 	return 0;
+
+err_free:
+	kho_unpreserve_free(ser);
+err_iommufd_unpreserve:
+	if (flags & VFIO_PCI_SER_IOMMUFD)
+		iommufd_device_unpreserve(args->session, device->iommufd_device);
+	return err;
 }
 
 static void vfio_pci_liveupdate_unpreserve(struct liveupdate_file_op_args *args)
 {
+	struct vfio_pci_core_device_ser *ser = phys_to_virt(args->serialized_data);
 	struct vfio_device *device = vfio_device_from_file(args->file);
+	bool iommufd = ser->flags & VFIO_PCI_SER_IOMMUFD;
+
+	guard(mutex)(&device->dev_set->lock);
 
 	pci_liveupdate_unpreserve(to_pci_dev(device->dev));
-	kho_unpreserve_free(phys_to_virt(args->serialized_data));
+	kho_unpreserve_free(ser);
+	if (iommufd)
+		iommufd_device_unpreserve(args->session, device->iommufd_device);
 }
 
 static int vfio_pci_liveupdate_freeze(struct liveupdate_file_op_args *args)
@@ -486,6 +529,28 @@ vfio_pci_liveupdate_attach_incoming(struct vfio_pci_core_device *vdev,
 	return 0;
 }
 
+/*
+ * A device preserved with its iommufd state can only be retrieved together
+ * with that iommufd, so fetch it using the saved token and fail if it can't be
+ * retrieved.
+ */
+static int vfio_pci_liveupdate_retrieve_iommufd(struct liveupdate_session *s,
+						const struct vfio_pci_core_device_ser *ser)
+{
+	struct file *file;
+	int ret;
+
+	if (!IS_ENABLED(CONFIG_IOMMU_LIVEUPDATE))
+		return -EOPNOTSUPP;
+
+	ret = liveupdate_get_file_incoming(s, ser->iommufd_token, &file);
+	if (ret)
+		return ret;
+
+	fput(file);
+	return 0;
+}
+
 static int vfio_pci_liveupdate_retrieve(struct liveupdate_file_op_args *args)
 {
 	struct pci_saved_state *saved_state = NULL;
@@ -513,7 +578,11 @@ static int vfio_pci_liveupdate_retrieve(struct liveupdate_file_op_args *args)
 	vdev = container_of(device, struct vfio_pci_core_device, vdev);
 	pdev = vdev->pdev;
 
-	if (!vfio_device_is_noiommu(device)) {
+	if (ser->flags & VFIO_PCI_SER_IOMMUFD) {
+		ret = vfio_pci_liveupdate_retrieve_iommufd(args->session, ser);
+		if (ret)
+			goto out_put_dev;
+	} else if (!vfio_device_is_noiommu(device)) {
 		ret = -EOPNOTSUPP;
 		goto out_put_dev;
 	}
diff --git a/drivers/vfio/pci/vfio_pci_priv.h b/drivers/vfio/pci/vfio_pci_priv.h
index 2549a5321fce..a5f6c0d3085f 100644
--- a/drivers/vfio/pci/vfio_pci_priv.h
+++ b/drivers/vfio/pci/vfio_pci_priv.h
@@ -194,10 +194,18 @@ vfio_pci_liveupdate_check_noiommu(struct vfio_pci_core_device *vdev)
 {
 	guard(mutex)(&vdev->vdev.dev_set->lock);
 
+	/*
+	 * A device preserved with its iommufd state keeps doing DMA through
+	 * the restored IOMMU domain. The IOMMU core claims DMA ownership of its
+	 * group when it restores the domain at boot, and nothing else can own
+	 * the group before the device is opened, so an owned group means the
+	 * device was restored and can keep running.
+	 */
 	if (vdev->liveupdate_kept_running &&
-	    !vfio_device_is_noiommu(&vdev->vdev)) {
+	    !vfio_device_is_noiommu(&vdev->vdev) &&
+	    !iommu_group_dma_owner_claimed(vdev->vdev.dev->iommu_group)) {
 		pci_warn(vdev->pdev,
-			 "Stopping device kept running across Live Update: not in noiommu mode\n");
+			 "Stopping device kept running across Live Update: not in noiommu mode and IOMMU state not restored\n");
 		pci_clear_master(vdev->pdev);
 		vdev->liveupdate_stopped_early = true;
 	}
diff --git a/include/linux/kho/abi/vfio_pci.h b/include/linux/kho/abi/vfio_pci.h
index 3ace2c9d8a4f..33b410d9a026 100644
--- a/include/linux/kho/abi/vfio_pci.h
+++ b/include/linux/kho/abi/vfio_pci.h
@@ -38,7 +38,17 @@
  */
 #define VFIO_PCI_SER_RESET_WORKS	BIT(0)
 
-#define VFIO_PCI_SER_KNOWN_FLAGS	VFIO_PCI_SER_RESET_WORKS
+/**
+ * define VFIO_PCI_SER_IOMMUFD - Device was preserved with its iommufd state.
+ *
+ * Set in &vfio_pci_core_device_ser.flags when the device was not in noiommu
+ * mode and its iommufd state was preserved. The LUO token of the preserved
+ * iommufd is stored in &vfio_pci_core_device_ser.iommufd_token.
+ */
+#define VFIO_PCI_SER_IOMMUFD		BIT(1)
+
+#define VFIO_PCI_SER_KNOWN_FLAGS	(VFIO_PCI_SER_RESET_WORKS | \
+					 VFIO_PCI_SER_IOMMUFD)
 
 /**
  * struct vfio_pci_core_device_ser - Serialized state of a single VFIO PCI
@@ -54,6 +64,8 @@
  * @saved_state_size: Size in bytes of the serialized &struct pci_saved_state_ser
  *                    blob stored in @data after ``config_map``, or 0 if none.
  * @reserved: Reserved for future use; must be 0.
+ * @iommufd_token: LUO token of the preserved iommufd, valid only when
+ *                 ``VFIO_PCI_SER_IOMMUFD`` is set in @flags.
  * @data: Variable-length payload containing, in order:
  *        1. ``u8 vconfig[@cfg_size]`` — VFIO virtualized config space snapshot
  *           captured during ``freeze()``.
@@ -70,6 +82,7 @@ struct vfio_pci_core_device_ser {
 	u32 flags;
 	u32 saved_state_size;
 	u32 reserved;
+	u64 iommufd_token;
 	u8 data[];
 } __packed;
 
-- 
2.56.0.506.g506a7eabc5-goog