[PATCH 6/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. As part of the
restoration process, the host must re-initialize the extensions to
restore their functionality.

Linux runs the updates in stop_machine() context, which prevents memory
allocation. This introduces a hard restriction that the updated TDX
environment must not consume more memory for the extensions.

The impact of the memory allocation restriction can be mitigated by
another requirement. Runtime updates should keep the add-on features
unchanged across updates, so that existing TDX users can continue to
operate without disruption. This requirement minimizes the chance of
increased memory demand. As a result, the restriction only affects the
compatibility rule for choosing the update image.

To adhere to these requirements, the post-update initialization for the
extensions can be simplified as:

  - Check if the extensions were originally initialized during boot up.
    If not, skip the re-initialization.
  - Assume no more memory needed, skip the memory adding step.
  - Re-initialize the extensions via TDH.EXT.INIT. The SEAMCALL leaf
    will fail if the updated module requires more memory, or if it drops
    the extensions initialization entirely, which indicates the update
    image is not compatible.

Signed-off-by: Xu Yilun <yilun.xu@xxxxxxxxxxxxxxx>
---
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 873b8393f32f..1ca3996f32dc 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -1289,7 +1289,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;
@@ -1332,6 +1332,19 @@ static __init int init_tdx_module_extensions(void)
 	return tdx_ext_init();
 }
 
+/*
+ * Don't update the extensions metadata, just follow the requirement originated
+ * during TDX module initialization. Let the extensions re-initialization fail
+ * if more memory is needed, or if ext_required is dropped after updates.
+ */
+static int update_tdx_module_extensions(void)
+{
+	if (!tdx_sysinfo.ext.ext_required)
+		return 0;
+
+	return tdx_ext_init();
+}
+
 static __init int init_tdx_module(void)
 {
 	int ret;
@@ -1532,6 +1545,10 @@ int tdx_module_run_update(void)
 	 */
 	WARN_ON_ONCE(ret);
 
+	ret = update_tdx_module_extensions();
+	if (ret)
+		return ret;
+
 	tdx_module_state.initialized = true;
 	return 0;
 }
-- 
2.25.1