[PATCH v4 3/6] x86/virt/tdx: Add extra memory to TDX module for the extensions

[Xu Yilun <yilun.xu@linux.intel.com>]

TDX module extensions need memory for their execution environment to
serve SEAMCALL leafs. The TDX architecture implements the extensions
in such a way that they use the memory outside of SEAM range, so the
kernel should add the memory upfront at initialization time.

Introduce a new memory adding process backed by a new SEAMCALL leaf
TDH.EXT.MEM.ADD. The kernel queries TDX module how much memory needed,
allocates it, adds it to the module, and never gets it back.

The TDX module accepts the memory in the form of an HPA array. This
array is passed via a single 64-bit SEAMCALL leaf parameter, which
encodes two values: the PFN of the container page holding the array, and
the number of entries in the array. Create a helper to encode this
format and name it after the TDX module term: HPA_LIST_INFO.

TDX module extensions consume tens of megabytes memory that will never
be returned to the host. Use contiguous page allocation to isolate these
large blocks entirely. This is a simple way to avoid permanent memory
fragmentation: requiring the full size to be contiguous is more
expensive than needed but should be good during the boot time. Print the
allocation amount on TDX module extensions initialization for
visibility.

The TDX module provides a metadata field "required_pages" to indicate
the required memory size. This metadata is present only when the module
supports the extensions and is meaningful only after the add-on feature
configuration. Don't read and cache the metadata in tdx_sysinfo at the
very beginning of TDX module initialization. Instead read it right
before allocating memory for the extensions, when it is guaranteed to be
valid.

AI was used under supervision to review code and workshop logs.

Signed-off-by: Xu Yilun <yilun.xu@linux.intel.com>
Reviewed-by: Tony Lindgren <tony.lindgren@linux.intel.com>
Reviewed-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
An alternative solution is to use a loop that gives memory on a memory
error code - TDX_EXT_MEMORY_POOL_REQUIRED, add one page per iteration
until TDH.EXT.INIT succeeds. Something like:

	do {
		ret = tdh_sys_init();
		if (ret == TDX_EXT_MEMORY_POOL_REQUIRED)
			tdh_ext_mem_add(); //single page
	} while (ret == TDX_EXT_MEMORY_POOL_REQUIRED);

This approach is slightly simpler as we don't have to query the module
for the total memory, no memory pre-allocation or segmentation math.
But allocating a single 4K page per iteration may cause permanent memory
fragmentation.

v4:
 - Rename memory_pool_required_pages to required_pages (Rick)
 - Rename local var sysinfo_ext to ext
 - Collect Reviewed-by tags

v3:
 - Merge remaining pre-check operations from the previous patch (Rick)
 - Don't save the metadata in tdx_sysinfo (Rick)
 - Tweak the code comment for memory_pool_required_pages == 0 (Chao)
 - Re-phrase the code comment for struct tdx_hpa_list (Tony)
 - Put the error handling comments on top of memory add loop, to avoid
   moving them around when DPAMT support is added in the following
   patch.
 - Print memory size before adding memory; it won't be freed afterward,
   whether the operation succeeds or fails.
 - Collect Reviewed-by tags

v2:
 - Add code comments for container page allocation (Kiryl)
 - Adjust comments/changelog for contiguous allocation of extensions
   memory (Kiryl)
 - s/PFN array/HPA array in commit log (Kiryl)
 - State the alternative solution (Rick)
 - Add code comment for struct tdx_hpa_list

v1:
 - Fix return value for SEAMCALL wrappers (Chao)
 - Print SEAMCALL error code for SEAMCALL wrappers (Xiaoyao)
 - Rename local vars to make the ext memory adding loop clear (Rick)
 - Remove input parameters for tdx_ext_mem_setup() (Kevin)
 - Add a Macro for tdh_hpa_list size.
 - Change the SEAMALL wrapper parameter type,
   struct page *hpa_list => struct tdx_hpa_list *hpa_list
 - changelog & code comments
---
 arch/x86/include/asm/tdx.h                  |   1 +
 arch/x86/include/asm/tdx_global_metadata.h  |   4 +
 arch/x86/virt/vmx/tdx/tdx.h                 |   1 +
 arch/x86/virt/vmx/tdx/tdx.c                 | 138 ++++++++++++++++++++++++++++
 arch/x86/virt/vmx/tdx/tdx_global_metadata.c |  14 +++
 5 files changed, 158 insertions(+)

diff --git a/arch/x86/include/asm/tdx.h b/arch/x86/include/asm/tdx.h
index 1fcfaed515fe..89e7628796b0 100644
--- a/arch/x86/include/asm/tdx.h
+++ b/arch/x86/include/asm/tdx.h
@@ -37,6 +37,7 @@
 #define TDX_FEATURES0_TD_PRESERVING	BIT_ULL(1)
 #define TDX_FEATURES0_NO_RBP_MOD	BIT_ULL(18)
 #define TDX_FEATURES0_DYNAMIC_PAMT	BIT_ULL(36)
+#define TDX_FEATURES0_EXT		BIT_ULL(39)
 
 #ifndef __ASSEMBLER__
 
diff --git a/arch/x86/include/asm/tdx_global_metadata.h b/arch/x86/include/asm/tdx_global_metadata.h
index 8a3cc1a2a41e..f18e63b671c2 100644
--- a/arch/x86/include/asm/tdx_global_metadata.h
+++ b/arch/x86/include/asm/tdx_global_metadata.h
@@ -47,6 +47,10 @@ struct tdx_sys_info_handoff {
 	u16 module_hv;
 };
 
+struct tdx_sys_info_ext {
+	u32 required_pages;
+};
+
 struct tdx_sys_info {
 	struct tdx_sys_info_version version;
 	struct tdx_sys_info_features features;
diff --git a/arch/x86/virt/vmx/tdx/tdx.h b/arch/x86/virt/vmx/tdx/tdx.h
index db209541d3cd..2adc5d31d258 100644
--- a/arch/x86/virt/vmx/tdx/tdx.h
+++ b/arch/x86/virt/vmx/tdx/tdx.h
@@ -50,6 +50,7 @@
 #define TDH_SYS_UPDATE			53
 #define TDH_PHYMEM_PAMT_ADD		58
 #define TDH_PHYMEM_PAMT_REMOVE		59
+#define TDH_EXT_MEM_ADD			61
 #define TDH_SYS_DISABLE			69
 
 /* TDX page types */
diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index f08278a47aff..744ef95e448f 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -1220,6 +1220,140 @@ static __init int init_tdmrs(struct tdmr_info_list *tdmr_list)
 	return 0;
 }
 
+#define TDX_HPA_LIST_MAX_NR_PAGES	(PAGE_SIZE / sizeof(u64))
+
+/*
+ * TDX module in-memory ABI to specify the memory being added to the TDX
+ * module. An array of HPAs, each element points to a 4K page.
+ */
+struct tdx_hpa_list {
+	u64 phys[TDX_HPA_LIST_MAX_NR_PAGES];
+};
+
+static_assert(sizeof(struct tdx_hpa_list) == PAGE_SIZE);
+
+#define HPA_LIST_INFO_FIRST_ENTRY	GENMASK_U64(11, 3)
+#define HPA_LIST_INFO_PFN		GENMASK_U64(51, 12)
+#define HPA_LIST_INFO_LAST_ENTRY	GENMASK_U64(63, 55)
+
+static __init u64 to_hpa_list_info(struct tdx_hpa_list *hpa_list,
+				   unsigned int nr_pages)
+{
+	return FIELD_PREP(HPA_LIST_INFO_FIRST_ENTRY, 0) |
+	       FIELD_PREP(HPA_LIST_INFO_PFN, PFN_DOWN(__pa(hpa_list))) |
+	       FIELD_PREP(HPA_LIST_INFO_LAST_ENTRY, nr_pages - 1);
+}
+
+static __init int tdx_ext_mem_add(struct tdx_hpa_list *hpa_list,
+				  unsigned int nr_pages)
+{
+	struct tdx_module_args args = {
+		.rcx = to_hpa_list_info(hpa_list, nr_pages),
+	};
+	u64 ret;
+
+	do {
+		/*
+		 * The TDX module overwrites RCX to track progress when this
+		 * SEAMCALL leaf is interrupted. Use seamcall_ret() to save and
+		 * pass the updated value back on retry.
+		 */
+		ret = seamcall_ret(TDH_EXT_MEM_ADD, &args);
+	} while (ret == TDX_INTERRUPTED_RESUMABLE);
+
+	if (ret != TDX_SUCCESS) {
+		pr_err("TDH.EXT.MEM.ADD failed: 0x%016llx\n", ret);
+		return -EIO;
+	}
+
+	return 0;
+}
+
+static __init int tdx_ext_mem_setup(void)
+{
+	struct tdx_hpa_list *hpa_list;
+	struct tdx_sys_info_ext ext;
+	unsigned int added_pages;
+	struct page *page;
+	int ret;
+
+	ret = get_tdx_sys_info_ext(&ext);
+	if (ret)
+		return ret;
+
+	/*
+	 * Skip the memory setup if no memory is required. This may happen when
+	 * no add-on features requiring TDX module extensions are configured
+	 * via TDH.SYS.CONFIG.
+	 */
+	if (!ext.required_pages)
+		return 0;
+
+	/*
+	 * Allocate the container page for the HPA_LIST. tdx_hpa_list is
+	 * guaranteed to be page-sized by static_assert(), so kzalloc()
+	 * guarantees the page alignment.
+	 */
+	hpa_list = kzalloc_obj(*hpa_list);
+	if (!hpa_list)
+		return -ENOMEM;
+
+	/*
+	 * Memory for TDX module extensions is never reclaimed and can be tens
+	 * of megabytes. Allocating a physically contiguous chunk is a simple
+	 * way to avoid permanent memory fragmentation: requiring the full size
+	 * to be contiguous is more expensive than needed but should be good
+	 * during the boot time.
+	 */
+	page = alloc_contig_pages(ext.required_pages, GFP_KERNEL,
+				  numa_mem_id(), &node_online_map);
+	if (!page) {
+		ret = -ENOMEM;
+		goto out_free_hpa_list;
+	}
+
+	/* Print the amount so users know the cost. */
+	pr_info("%lu KB allocated for TDX module extensions\n",
+		ext.required_pages * PAGE_SIZE / 1024);
+
+	/*
+	 * Following TDX operations shouldn't fail, and if they do, things are
+	 * broken enough that complex error handling isn't worth it.
+	 * Intentionally leak all pages on failure, including un-added pages.
+	 */
+	added_pages = 0;
+	while (added_pages < ext.required_pages) {
+		unsigned int chunk_pages = min(ext.required_pages - added_pages,
+					       TDX_HPA_LIST_MAX_NR_PAGES);
+		struct page *chunk = page + added_pages;
+		unsigned int i;
+
+		for (i = 0; i < chunk_pages; i++)
+			hpa_list->phys[i] = page_to_phys(chunk + i);
+
+		ret = tdx_ext_mem_add(hpa_list, chunk_pages);
+		if (ret) {
+			WARN(1, "TDX module rejected memory for extensions, stranded all pages\n");
+			break;
+		}
+
+		added_pages += chunk_pages;
+	}
+
+out_free_hpa_list:
+	kfree(hpa_list);
+
+	return ret;
+}
+
+static __init int init_tdx_module_extensions(void)
+{
+	if (!(tdx_sysinfo.features.tdx_features0 & TDX_FEATURES0_EXT))
+		return 0;
+
+	return tdx_ext_mem_setup();
+}
+
 static __init int init_tdx_module(void)
 {
 	int ret;
@@ -1278,6 +1412,10 @@ static __init int init_tdx_module(void)
 	if (ret)
 		goto err_reset_pamts;
 
+	ret = init_tdx_module_extensions();
+	if (ret)
+		goto err_reset_pamts;
+
 	pr_info("%lu KB allocated for PAMT\n", tdmrs_count_pamt_kb(&tdx_tdmr_list));
 
 out_put_tdxmem:
diff --git a/arch/x86/virt/vmx/tdx/tdx_global_metadata.c b/arch/x86/virt/vmx/tdx/tdx_global_metadata.c
index 98ebf17aab1c..9ab4d021017f 100644
--- a/arch/x86/virt/vmx/tdx/tdx_global_metadata.c
+++ b/arch/x86/virt/vmx/tdx/tdx_global_metadata.c
@@ -125,6 +125,20 @@ static int get_tdx_sys_info_handoff(struct tdx_sys_info_handoff *sysinfo_handoff
 	return 0;
 }
 
+static __init int get_tdx_sys_info_ext(struct tdx_sys_info_ext *sysinfo_ext)
+{
+	int ret;
+	u64 val;
+
+	ret = read_sys_metadata_field(0x3100000200000000, &val);
+	if (ret)
+		return ret;
+
+	sysinfo_ext->required_pages = val;
+
+	return 0;
+}
+
 static __init int get_tdx_sys_info(struct tdx_sys_info *sysinfo)
 {
 	int ret = 0;

-- 
2.25.1