[PATCH v3 5/6] x86/virt/tdx: Re-initialize the extensions on runtime TDX module update

[Xu Yilun]

Runtime TDX module update introduces a mechanism to update the module
firmware while preserving and restoring TDX operations. The extensions
functionalities should also be re-initialized as part of the restoration
process.

The TDX architecture supports an update flow which allows updated
extensions to consume more memory than their original boot-time
requirement. So the arch defines the extensions re-initialization flow
the same as boot-up initialization: the host queries TDX module how much
additional memory needed, allocates it, adds it to the module via
TDH.EXT.MEM.ADD, then re-initializes the extensions via TDH.EXT.INIT.

Linux runs the updates in stop_machine() context, which prevents memory
allocation. So for Linux, a compatible update must not install updated
extensions that require additional memory.

Given that the memory for the extensions must not increase, the
re-initialization skips the memory adding steps. It is simplified as:

  - Check if the extensions are supported via TDX_FEATURES0_EXT. If not,
    skip the extensions re-initialization.

  - Re-initialize the extensions via TDH.EXT.INIT. The SEAMCALL leaf
    will fail if the updated extensions require more memory, which
    indicates the update image is not compatible.

Signed-off-by: Xu Yilun <yilun.xu@xxxxxxxxxxxxxxx>
---
v3:
 - Add this patch back cause the latest discussion decides the
   re-initialization won't be integrated in TDH.SYS.UPDATE, a dedicated
   TDH.EXT.INIT call is still needed (Dave & Rick)
 - Drop the ext_required check (Rick)
 - Rename the reinit function as reinit_tdx_module_extensions() (Tony)
 - Move the reinit function right under tdh_sys_update() (Rick)

v2:
 - Removed this patch cause the TDX module is expected to re-initialize
   the extensions on TDH.SYS.UPDATE

v1:
 - Don't update the extensions metadata any more, only check the
   metadata originated at boot time.
 - Remove memory_pool_required_pages check, let TDH.EXT.INIT fail if
   more memory required.
 - Changelog & code comments
---
 arch/x86/virt/vmx/tdx/tdx.c | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index d8df3b2b1d17..5d5f9da1ab02 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -1348,7 +1348,7 @@ static __init int tdx_ext_mem_setup(void)
 	return ret;
 }
 
-static __init int tdx_ext_init(void)
+static int tdx_ext_init(void)
 {
 	struct tdx_module_args args = {};
 	u64 ret;
@@ -1379,6 +1379,19 @@ static __init int init_tdx_module_extensions(void)
 	return tdx_ext_init();
 }
 
+/*
+ * Don't update the memory requirement metadata or try memory allocation in
+ * stop_machine(). If an incompatible update requires more memory, let the
+ * extensions re-initialization fail.
+ */
+static int reinit_tdx_module_extensions(void)
+{
+	if (!(tdx_sysinfo.features.tdx_features0 & TDX_FEATURES0_EXT))
+		return 0;
+
+	return tdx_ext_init();
+}
+
 static __init int init_tdx_module(void)
 {
 	int ret;
@@ -1562,6 +1575,10 @@ int tdx_module_run_update(void)
 	if (ret)
 		return ret;
 
+	ret = reinit_tdx_module_extensions();
+	if (ret)
+		return ret;
+
 	ret = get_tdx_sys_info_version(&tdx_sysinfo.version);
 	/*
 	 * Only fails if there is something unexpected

-- 
2.25.1