Liveupdate prepare callback preserves the iommus of the devices that are
attached to a preserved iommu domain. It does this for only pcie devices
by iterating through all the pcie devices and checking whether the
attached iommu domain is preserved.

Only stub implementation of the iommu and device preservation are added
with this commit.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 drivers/iommu/intel/dmar.c       |   4 +
 drivers/iommu/intel/iommu.h      |   3 +
 drivers/iommu/intel/liveupdate.c | 150 ++++++++++++++++++++++++++++++-
 3 files changed, 156 insertions(+), 1 deletion(-)

diff --git a/drivers/iommu/intel/dmar.c b/drivers/iommu/intel/dmar.c
index 248bc7e9b035..cd6ce519c1da 100644
--- a/drivers/iommu/intel/dmar.c
+++ b/drivers/iommu/intel/dmar.c
@@ -1101,6 +1101,10 @@ static int alloc_iommu(struct dmar_drhd_unit *drhd)
 	ida_init(&iommu->domain_ida);
 	mutex_init(&iommu->did_lock);
 
+#ifdef CONFIG_LIVEUPDATE
+	atomic_set(&iommu->preserved, 0);
+#endif
+
 	ver = readl(iommu->reg + DMAR_VER_REG);
 	pr_info("%s: reg_base_addr %llx ver %d:%d cap %llx ecap %llx\n",
 		iommu->name,
diff --git a/drivers/iommu/intel/iommu.h b/drivers/iommu/intel/iommu.h
index 6b69232efffd..93ac55eb49f0 100644
--- a/drivers/iommu/intel/iommu.h
+++ b/drivers/iommu/intel/iommu.h
@@ -758,6 +758,9 @@ struct intel_iommu {
 	void *perf_statistic;
 
 	struct iommu_pmu *pmu;
+#ifdef CONFIG_LIVEUPDATE
+	atomic_t preserved;
+#endif
 };
 
 /* PCI domain-device relationship */
diff --git a/drivers/iommu/intel/liveupdate.c b/drivers/iommu/intel/liveupdate.c
index a15feef4d9ca..94aabf025a60 100644
--- a/drivers/iommu/intel/liveupdate.c
+++ b/drivers/iommu/intel/liveupdate.c
@@ -6,24 +6,172 @@
 
 #define pr_fmt(fmt)    "iommu: liveupdate: " fmt
 
+#include <linux/kexec_handover.h>
 #include <linux/liveupdate.h>
 #include <linux/module.h>
+#include <linux/pci.h>
 
 #include "iommu.h"
 
+struct iommu_unit_ser {
+	u64 phys_addr;
+	u64 root_table;
+};
+
+struct device_ser {
+	u64 bdf;
+	u64 pasid_table;
+	u64 pasid_order;
+	u64 iommu_phys;
+};
+
+struct iommu_ser {
+	u64 nr_iommus;
+	u64 nr_devices;
+
+	union {
+		u64 iommu_units_phys;
+		struct iommu_unit_ser *iommu_units;
+	};
+
+	union {
+		u64 devices_phys;
+		struct device_ser *devices;
+	};
+};
+
 int intel_iommu_domain_liveupdate_preserve(struct iommu_domain *domain)
 {
 	pr_warn("Not implemented\n");
 	return 0;
+}
+
+static bool is_device_domain_preserved(struct device *dev)
+{
+	struct device_domain_info *info = dev_iommu_priv_get(dev);
 
+	return atomic_read(&info->domain->domain.preserved) == 1;
 }
 
-static int intel_liveupdate_prepare(struct liveupdate_subsystem *handle, u64 *data)
+static int preserve_device_state(struct pci_dev *dev, struct device_ser *ser)
 {
 	pr_warn("Not implemented\n");
 	return 0;
 }
 
+static int preserve_iommu_state(struct intel_iommu *iommu,
+				struct iommu_unit_ser *ser)
+{
+	pr_warn("Not implemented\n");
+	return 0;
+}
+
+static void unpreserve_state(struct iommu_ser *ser)
+{
+	pr_warn("Not implemented\n");
+}
+
+static int preserve_state(struct iommu_ser *ser)
+{
+	struct device_domain_info *info;
+	struct pci_dev *pdev = NULL;
+	struct dmar_drhd_unit *drhd;
+	struct intel_iommu *iommu;
+	int ret = 0;
+
+	for_each_pci_dev(pdev) {
+		if (!is_device_domain_preserved(&pdev->dev))
+			continue;
+
+		info = dev_iommu_priv_get(&pdev->dev);
+		if (!info)
+			return -EINVAL;
+
+		if (ser->devices)
+			ret = preserve_device_state(pdev, &ser->devices[ser->nr_devices]);
+
+		if (ret)
+			return ret;
+
+		atomic_set(&info->iommu->preserved, 1);
+		ser->nr_devices++;
+	}
+
+	for_each_iommu(iommu, drhd) {
+		if (!atomic_read(&iommu->preserved))
+			continue;
+
+		atomic_set(&iommu->preserved, 0);
+		if (ser->iommu_units)
+			ret = preserve_iommu_state(iommu, &ser->iommu_units[ser->nr_iommus]);
+
+		if (ret)
+			return ret;
+
+		ser->nr_iommus++;
+	}
+
+	return 0;
+}
+
+static struct iommu_ser *alloc_preserve_state_mem(void)
+{
+	struct iommu_ser *ser_ptr;
+	struct iommu_ser ser;
+	struct folio *folio;
+	size_t sz;
+	int ret;
+
+	memset(&ser, 0, sizeof(ser));
+	ret = preserve_state(&ser);
+	if (ret)
+		goto error;
+
+	sz = sizeof(struct iommu_ser) +
+			(ser.nr_iommus * sizeof(struct iommu_unit_ser)) +
+			(ser.nr_devices * sizeof(struct device_ser));
+
+	folio = folio_alloc(GFP_KERNEL, get_order(sz));
+	if (!folio)
+		return ERR_PTR(-ENOMEM);
+
+	ret = kho_preserve_folio(folio);
+	if (ret)
+		goto error_preserve;
+
+	ser_ptr = folio_address(folio);
+	memset(ser_ptr, 0, sz);
+	ser_ptr->iommu_units = (void *)(ser_ptr + 1);
+	ser_ptr->devices = (void *)(ser_ptr->iommu_units + ser.nr_iommus);
+
+	return ser_ptr;
+
+error_preserve:
+	folio_put(folio);
+error:
+	return ERR_PTR(ret);
+}
+
+static int intel_liveupdate_prepare(struct liveupdate_subsystem *handle, u64 *data)
+{
+	struct iommu_ser *ser;
+	int ret;
+
+	guard_liveupdate_state_write();
+	ser = alloc_preserve_state_mem();
+	if (IS_ERR(ser))
+		return PTR_ERR(ser);
+
+	ret = preserve_state(ser);
+	if (ret)
+		unpreserve_state(ser);
+
+	if (!ret)
+		*data = __pa(ser);
+
+	return ret;
+}
+
 static void intel_liveupdate_cancel(struct liveupdate_subsystem *handle, u64 data)
 {
 	pr_warn("Not implemented\n");
-- 
2.51.0.536.g15c5d4f767-goog