[PATCH 4/5] iommu/vt-d: Add helpers to set up and tear down TDX extensions

[Lu Baolu]

Implement per-IOMMU TDX Connect bring-up/tear-down helpers.

For each active IOMMU that advertises TDXCS and has DMA translation
enabled, the driver allocates and populates the IOMMU_MT parameter
layout required by TDH.IOMMU.SETUP:

- two contiguous invalidation-queue buffers, and
- IOMMU_MT_PAGES_COUNT metadata pages.

The driver invokes tdh_iommu_setup() to transition the IOMMU into Secure
TDX Mode, then stores the returned tdx_iommu_id and allocation state for
later cleanup.

On teardown, or on partial-init rollback, the driver calls
tdh_iommu_clear() for each initialized IOMMU and frees all associated
IOMMU_MT pages.

Completion handling:

- TDX_SUCCESS: setup completed successfully.
- TDX_OPERAND_INVALID: treated as "not supported on this IOMMU" and
  skipped.
- all other return codes: treated as failures and abort initialization
  with rollback.

No explicit clflush is needed for pages shared with the TDX module in
this flow, as cache coherency is handled architecturally.

Signed-off-by: Lu Baolu <baolu.lu@xxxxxxxxxxxxxxx>
---
 drivers/iommu/intel/iommu.h |  17 ++++
 drivers/iommu/intel/dmar.c  |   3 +
 drivers/iommu/intel/tdxc.c  | 150 ++++++++++++++++++++++++++++++++++++
 3 files changed, 170 insertions(+)

diff --git a/drivers/iommu/intel/iommu.h b/drivers/iommu/intel/iommu.h
index 452a381e6a40..7bf70113c4c8 100644
--- a/drivers/iommu/intel/iommu.h
+++ b/drivers/iommu/intel/iommu.h
@@ -680,6 +680,17 @@ struct iommu_pmu {
 #define IOMMU_IRQ_ID_OFFSET_PRQ		(DMAR_UNITS_SUPPORTED)
 #define IOMMU_IRQ_ID_OFFSET_PERF	(2 * DMAR_UNITS_SUPPORTED)
 
+/*
+ * Represents a list of pages for TDX Module defined IOMMU_MT object.
+ * Typically it uses a "root page" as the medium to exchange a list of
+ * data pages between host and TDX Module.
+ */
+struct tdxc_pages {
+	u64 *root;
+	void **pages;
+	unsigned int nr_entries;
+};
+
 struct intel_iommu {
 	void __iomem	*reg; /* Pointer to hardware regs, virtual addr */
 	u64 		reg_phys; /* physical address of hw register set */
@@ -736,6 +747,12 @@ struct intel_iommu {
 	void *perf_statistic;
 
 	struct iommu_pmu *pmu;
+#ifdef CONFIG_INTEL_IOMMU_TDX_CONNECT
+	/* mutex to protect below tdx state data */
+	struct mutex tdx_lock;
+	u64 tdx_iommu_id;
+	struct tdxc_pages *mt_pages;
+#endif
 };
 
 /* PCI domain-device relationship */
diff --git a/drivers/iommu/intel/dmar.c b/drivers/iommu/intel/dmar.c
index ba675b08cd20..310b8d0c4132 100644
--- a/drivers/iommu/intel/dmar.c
+++ b/drivers/iommu/intel/dmar.c
@@ -1175,6 +1175,9 @@ static int alloc_iommu(struct dmar_drhd_unit *drhd)
 	ida_init(&iommu->domain_ida);
 	mutex_init(&iommu->did_lock);
 	iommu->max_domain_id = cap_ndoms(iommu->cap);
+#ifdef CONFIG_INTEL_IOMMU_TDX_CONNECT
+	mutex_init(&iommu->tdx_lock);
+#endif
 
 	ver = readl(iommu->reg + DMAR_VER_REG);
 	pr_info("%s: reg_base_addr %llx ver %d:%d cap %llx ecap %llx\n",
diff --git a/drivers/iommu/intel/tdxc.c b/drivers/iommu/intel/tdxc.c
index 559c752a1527..b5dfdeeb23db 100644
--- a/drivers/iommu/intel/tdxc.c
+++ b/drivers/iommu/intel/tdxc.c
@@ -10,12 +10,128 @@
 #include <linux/pci.h>
 #include <asm/vmx.h>
 #include <asm/tdx.h>
+
+#include "../iommu-pages.h"
 #include "iommu.h"
 
+#define IQ_BUFFERS_NUM		2
+#define IQ_BUFFER_PAGES		2
+#define IQ_BUFFER_SIZE		SZ_8K
+
 bool intel_tdxc_initialized;
 
+static void free_mt_pages(struct tdxc_pages *array)
+{
+	if (!array)
+		return;
+
+	for (int i = 0; i < array->nr_entries; i++)
+		iommu_free_pages(array->pages[i]);
+
+	iommu_free_pages(array->root);
+	kfree(array->pages);
+	kfree(array);
+}
+
+DEFINE_FREE(free_mt_pages, struct tdxc_pages *, free_mt_pages(_T))
+
+static void **alloc_mt_pages(unsigned int nr_entries, int node)
+{
+	void **pages;
+	void *vaddr;
+	int i;
+
+	pages = kzalloc_objs(*pages, nr_entries);
+	if (!pages)
+		return NULL;
+
+	/* Allocate two contiguous buffers for the invalidation queue. */
+	pages[0] = iommu_alloc_pages_node_sz(node, GFP_KERNEL, IQ_BUFFER_SIZE);
+	if (!pages[0])
+		goto free_pages;
+
+	pages[1] = iommu_alloc_pages_node_sz(node, GFP_KERNEL, IQ_BUFFER_SIZE);
+	if (!pages[1])
+		goto free_pages;
+
+	/* Allocate the required number of pages for the IOMMU metadata. */
+	for (i = IQ_BUFFERS_NUM; i < nr_entries; i++) {
+		vaddr = iommu_alloc_pages_node_sz(node, GFP_KERNEL, SZ_4K);
+		if (!vaddr)
+			goto free_pages;
+		pages[i] = vaddr;
+	}
+
+	return pages;
+free_pages:
+	for (i = 0; i < nr_entries; i++) {
+		if (!pages[i])
+			break;
+
+		iommu_free_pages(pages[i]);
+	}
+	kfree(pages);
+
+	return NULL;
+}
+
+static void populate_mt_pages(struct tdxc_pages *array)
+{
+	unsigned int nr_entries = array->nr_entries;
+	void **pages = array->pages;
+	u64 *entries = array->root;
+	int i;
+
+	/*
+	 * Populate the parameter for the TDH_IOMMU_SETUP SEAMCALL according to
+	 * the format defined in "Table 3.35: Structure of IOMMU_MT Parameter"
+	 * of the ABI reference specification.
+	 */
+	for (i = 0; i < nr_entries; i++) {
+		entries[i] = __pa(pages[i]);
+		if (i < IQ_BUFFERS_NUM)
+			entries[i] |= IQ_BUFFER_PAGES;
+	}
+}
+
+static struct tdxc_pages *tdxc_alloc_mt_pages(struct intel_iommu *iommu,
+					      unsigned int nr_mt_pages)
+{
+	unsigned int nr_entries = nr_mt_pages + IQ_BUFFERS_NUM;
+	struct tdxc_pages *array;
+
+	if (!nr_mt_pages || nr_mt_pages > (PAGE_SIZE / sizeof(u64) - IQ_BUFFERS_NUM))
+		return NULL;
+
+	array = kzalloc_obj(*array);
+	if (!array)
+		return NULL;
+
+	array->root = iommu_alloc_pages_node_sz(iommu->node, GFP_KERNEL, SZ_4K);
+	if (!array->root)
+		goto free_array;
+
+	array->nr_entries = nr_entries;
+	array->pages = alloc_mt_pages(nr_entries, iommu->node);
+	if (!array->pages)
+		goto free_root;
+
+	populate_mt_pages(array);
+
+	return array;
+
+free_root:
+	iommu_free_pages(array->root);
+free_array:
+	kfree(array);
+	return NULL;
+}
+
 static int intel_iommu_bringup_tdxc(struct intel_iommu *iommu, unsigned int nr_pages)
 {
+	struct dmar_drhd_unit *drhd = iommu->drhd;
+	u64 r, tdx_iommu_id;
+
 	/*
 	 * Nothing to do if the iommu doesn't support TDX extension or the
 	 * DMA translation has not been enabled.
@@ -23,12 +139,46 @@ static int intel_iommu_bringup_tdxc(struct intel_iommu *iommu, unsigned int nr_p
 	if (!ecap_tdxcs(iommu->ecap) || !(iommu->gcmd & DMA_GCMD_TE))
 		return 0;
 
+	struct tdxc_pages *iommu_mt __free(free_mt_pages) =
+			tdxc_alloc_mt_pages(iommu, nr_pages);
+	if (!iommu_mt)
+		return -ENOMEM;
+
+	guard(mutex)(&iommu->tdx_lock);
+	r = tdh_iommu_setup(drhd->reg_base_addr, iommu_mt->root, &tdx_iommu_id);
+	/* TDX Extension is not supported on this iommu. Nothing to do. */
+	if ((r & TDX_SEAMCALL_STATUS_MASK) == TDX_OPERAND_INVALID)
+		return 0;
+	if (r) {
+		pr_err("%s: TDH.IOMMU.SETUP failed, status 0x%llx\n", iommu->name, r);
+		return -EFAULT;
+	}
+
+	iommu->tdx_iommu_id = tdx_iommu_id;
+	iommu->mt_pages = no_free_ptr(iommu_mt);
+
 	/* Bring-up is not complete yet; report as unsupported for now. */
 	return -EOPNOTSUPP;
 }
 
 static void intel_iommu_teardown_tdxc(struct intel_iommu *iommu)
 {
+	u64 r;
+
+	guard(mutex)(&iommu->tdx_lock);
+
+	if (!iommu->mt_pages)
+		return;
+
+	r = tdh_iommu_clear(iommu->tdx_iommu_id);
+	if (r) {
+		pr_err("%s: TDH.IOMMU.CLEAR failed, status 0x%llx\n", iommu->name, r);
+		return;
+	}
+
+	free_mt_pages(iommu->mt_pages);
+	iommu->mt_pages = NULL;
+	iommu->tdx_iommu_id = 0;
 }
 
 void intel_tdxc_exit(void)
-- 
2.43.0