[PATCH v2 0/2] x86/fpu: Fix dynamic fpstate leak on exec()

[PATCH v2 0/2] x86/fpu: Fix dynamic fpstate leak on exec()

[Guixiong Wei]

An exec() after using an XFD-controlled xfeature resets fpu->fpstate to
its embedded storage without freeing the dynamically allocated state.
Fix the leak in fpstate_reset() and add an AMX regression selftest.

The fix preserves the old pointer, installs and initializes the embedded
fpstate, and then frees the detached allocation. Since fpstate_reset()
now owns cleanup, fpu_clone() initializes dst_fpu->fpstate to NULL before
invoking it.

The selftest performs ten XTILEDATA request, XRSTOR and self-exec cycles
in the same task and checks the associated /proc/vmallocinfo entries.
It fails with 10 leaked allocations on the unfixed kernel and passes
with zero on the fixed kernel.

Changes since v1:
- Rename the fix and describe the memory leak explicitly.
- Move cleanup into fpstate_reset().
- Detach the old fpstate before freeing it.
- Initialize dst_fpu->fpstate to NULL before resetting it.
- Add the requested AMX regression selftest as a separate patch.

v1: https://lore.kernel.org/r/20260929151013.81562-2-weiguixiong@bytedance.com

Guixiong Wei (2):
  x86/fpu: Fix memory leak with dynamic fpstate and exec()
  selftests/x86/amx: Test dynamic fpstate cleanup across exec()

 arch/x86/include/asm/fpu/api.h    |   6 +-
 arch/x86/kernel/fpu/core.c        |   5 ++
 arch/x86/kernel/fpu/xstate.c      |  10 +--
 arch/x86/kernel/process.c         |   2 +-
 tools/testing/selftests/x86/amx.c | 144 +++++++++++++++++++++++++++++-
 5 files changed, 156 insertions(+), 11 deletions(-)


base-commit: 72d3fcf802c45d00b300f25b848a93c3a2bd7c7e
-- 
2.50.1 (Apple Git-155)

[PATCH v2 1/2] x86/fpu: Fix memory leak with dynamic fpstate and exec()

[Guixiong Wei]

The fpstate embedded in struct fpu only accommodates the default
xfeatures. After a task obtains permission for an XFD-controlled
xfeature, its first use raises #NM. The handler calls
__xfd_enable_feature(), which uses fpstate_realloc() to install a larger
fpstate allocated with vzalloc().

On exec(), fpu_flush_thread() calls fpstate_reset(), which overwrites
fpu->fpstate with the embedded fpstate. This leaks the dynamic
allocation because arch_release_task_struct() cannot free it when the
task exits.

Make fpstate_reset() own the complete lifetime transition. Preserve the
old pointer, install and initialize the embedded fpstate, and then free
the old allocation. This detaches the allocation from fpu->fpstate
before freeing it and makes future reset callers release dynamic state
automatically.

Since fpstate_reset() now releases the old state, initialize
dst_fpu->fpstate to NULL before invoking it from fpu_clone().

Make fpstate_free() operate on the fpstate pointer itself so it can free
the detached allocation. Use it for the existing reallocation and task
release paths as well.

On an AMX-capable KVM guest, 1000 iterations left 1000 dynamic fpstate
allocations on an unfixed kernel and none with this change.

Fixes: 500afbf645a0 ("x86/fpu/xstate: Add fpstate_realloc()/free()")
Cc: stable@vger.kernel.org
Signed-off-by: Guixiong Wei <weiguixiong@bytedance.com>
---
 arch/x86/include/asm/fpu/api.h |  6 +++---
 arch/x86/kernel/fpu/core.c     |  5 +++++
 arch/x86/kernel/fpu/xstate.c   | 10 ++++------
 arch/x86/kernel/process.c      |  2 +-
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/arch/x86/include/asm/fpu/api.h b/arch/x86/include/asm/fpu/api.h
index 90c63fe19c0fb..3e0c3b7dce73b 100644
--- a/arch/x86/include/asm/fpu/api.h
+++ b/arch/x86/include/asm/fpu/api.h
@@ -123,11 +123,11 @@ extern void fpu__resume_cpu(void);
 DECLARE_PER_CPU(bool, kernel_fpu_allowed);
 DECLARE_PER_CPU(struct fpu *, fpu_fpregs_owner_ctx);
 
-/* Process cleanup */
+/* Dynamic fpstate cleanup */
 #ifdef CONFIG_X86_64
-extern void fpstate_free(struct fpu *fpu);
+extern void fpstate_free(struct fpstate *fpstate);
 #else
-static inline void fpstate_free(struct fpu *fpu) { }
+static inline void fpstate_free(struct fpstate *fpstate) { }
 #endif
 
 /* fpstate-related functions which are exported to KVM */
diff --git a/arch/x86/kernel/fpu/core.c b/arch/x86/kernel/fpu/core.c
index d1aeecd57f5ed..bcebe1f0d82b7 100644
--- a/arch/x86/kernel/fpu/core.c
+++ b/arch/x86/kernel/fpu/core.c
@@ -594,6 +594,8 @@ static void __fpstate_reset(struct fpstate *fpstate)
 
 void fpstate_reset(struct fpu *fpu)
 {
+	struct fpstate *oldfpstate = fpu->fpstate;
+
 	/* Set the fpstate pointer to the default fpstate */
 	fpu->fpstate = &fpu->__fpstate;
 	__fpstate_reset(fpu->fpstate);
@@ -610,6 +612,8 @@ void fpstate_reset(struct fpu *fpu)
 	 * guest FPUs, which allows for common guest and userspace ABI.
 	 */
 	fpu->guest_perm.__user_state_size = fpu_user_cfg.default_size;
+
+	fpstate_free(oldfpstate);
 }
 
 static inline void fpu_inherit_perms(struct fpu *dst_fpu)
@@ -670,6 +674,7 @@ int fpu_clone(struct task_struct *dst, u64 clone_flags, bool minimal,
 	/* The new task's FPU state cannot be valid in the hardware. */
 	dst_fpu->last_cpu = -1;
 
+	dst_fpu->fpstate = NULL;
 	fpstate_reset(dst_fpu);
 
 	if (!cpu_feature_enabled(X86_FEATURE_FPU))
diff --git a/arch/x86/kernel/fpu/xstate.c b/arch/x86/kernel/fpu/xstate.c
index a7b6524a9dea2..ea3736a74fd43 100644
--- a/arch/x86/kernel/fpu/xstate.c
+++ b/arch/x86/kernel/fpu/xstate.c
@@ -1556,10 +1556,10 @@ static int __init xfd_update_static_branch(void)
 }
 arch_initcall(xfd_update_static_branch)
 
-void fpstate_free(struct fpu *fpu)
+void fpstate_free(struct fpstate *fpstate)
 {
-	if (fpu->fpstate && fpu->fpstate != &fpu->__fpstate)
-		vfree(fpu->fpstate);
+	if (fpstate && fpstate->is_valloc)
+		vfree(fpstate);
 }
 
 /**
@@ -1640,9 +1640,7 @@ static int fpstate_realloc(u64 xfeatures, unsigned int ksize,
 		xfd_update_state(fpu->fpstate);
 	fpregs_unlock();
 
-	/* Only free valloc'ed state */
-	if (curfps && curfps->is_valloc)
-		vfree(curfps);
+	fpstate_free(curfps);
 
 	return 0;
 }
diff --git a/arch/x86/kernel/process.c b/arch/x86/kernel/process.c
index 346c438ac8801..b7306f257f4b0 100644
--- a/arch/x86/kernel/process.c
+++ b/arch/x86/kernel/process.c
@@ -118,7 +118,7 @@ int arch_dup_task_struct(struct task_struct *dst, struct task_struct *src)
 void arch_release_task_struct(struct task_struct *tsk)
 {
 	if (fpu_state_size_dynamic() && !(tsk->flags & (PF_KTHREAD | PF_USER_WORKER)))
-		fpstate_free(x86_task_fpu(tsk));
+		fpstate_free(x86_task_fpu(tsk)->fpstate);
 }
 #endif
 
-- 
2.50.1 (Apple Git-155)

[PATCH v2 2/2] selftests/x86/amx: Test dynamic fpstate cleanup across exec()

[Guixiong Wei]

An AMX task gets a dynamically allocated fpstate after requesting
XTILEDATA permission and first loading tile data. Test that exec()
releases this allocation instead of making it unreachable.

Run ten request, XRSTOR and self-exec cycles in the same task. Count the
matching allocations in /proc/vmallocinfo before the test, while the
first allocation is live and after the task exits. Skip the test when
/proc/vmallocinfo is unavailable or when concurrent global vmalloc
changes prevent an isolated measurement.

The test observed 10 leaked allocations on an unfixed kernel and zero
with the fix applied.

Signed-off-by: Guixiong Wei <weiguixiong@bytedance.com>
---
 tools/testing/selftests/x86/amx.c | 144 +++++++++++++++++++++++++++++-
 1 file changed, 143 insertions(+), 1 deletion(-)

diff --git a/tools/testing/selftests/x86/amx.c b/tools/testing/selftests/x86/amx.c
index 40769c16de1bb..cf098dfbe50c0 100644
--- a/tools/testing/selftests/x86/amx.c
+++ b/tools/testing/selftests/x86/amx.c
@@ -3,8 +3,10 @@
 #define _GNU_SOURCE
 #include <err.h>
 #include <errno.h>
+#include <limits.h>
 #include <setjmp.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <string.h>
 #include <stdbool.h>
 #include <unistd.h>
@@ -30,6 +32,9 @@
 #define XFEATURE_MASK_XTILEDATA	(1 << XFEATURE_XTILEDATA)
 #define XFEATURE_MASK_XTILE	(XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA)
 
+#define EXEC_TEST_ARG		"--exec-test"
+#define EXEC_TEST_ITERS		10
+
 struct xstate_info xtiledata;
 
 /* The helpers for managing XSAVE buffer and tile states: */
@@ -478,11 +483,144 @@ static void test_fork(void)
 	_exit(0);
 }
 
-int main(void)
+static int count_dynamic_fpstates(void)
+{
+	char *line = NULL;
+	size_t line_size = 0;
+	FILE *fp;
+	int count = 0;
+
+	fp = fopen("/proc/vmallocinfo", "r");
+	if (!fp)
+		return -errno;
+
+	while (getline(&line, &line_size, fp) >= 0) {
+		if (strstr(line, "__xfd_enable_feature") ||
+		    strstr(line, "fpstate_realloc"))
+			count++;
+	}
+
+	if (ferror(fp))
+		count = -EIO;
+
+	free(line);
+	fclose(fp);
+	return count;
+}
+
+static int parse_exec_arg(const char *arg)
+{
+	char *end;
+	long value;
+
+	errno = 0;
+	value = strtol(arg, &end, 10);
+	if (errno || *end || value < 0 || value > INT_MAX)
+		fatal_error("invalid exec test argument: %s", arg);
+
+	return value;
+}
+
+static int run_exec_test(int iterations, int baseline)
+{
+	char iterations_arg[16];
+	char baseline_arg[16];
+	int allocated;
+
+	if (!iterations)
+		return 0;
+
+	req_xtiledata_perm();
+	if (!load_rand_tiledata(stashed_xsave))
+		fatal_error("failed to load tiledata before exec()");
+
+	if (iterations == EXEC_TEST_ITERS) {
+		allocated = count_dynamic_fpstates();
+		if (allocated != baseline + 1)
+			return KSFT_SKIP;
+	}
+
+	snprintf(iterations_arg, sizeof(iterations_arg), "%d", iterations - 1);
+	snprintf(baseline_arg, sizeof(baseline_arg), "%d", baseline);
+	execl("/proc/self/exe", "amx", EXEC_TEST_ARG, iterations_arg,
+	      baseline_arg, NULL);
+	fatal_error("exec");
+}
+
+static void test_exec(void)
+{
+	char iterations_arg[16];
+	char baseline_arg[16];
+	int before, after, delta;
+	pid_t child;
+	int status;
+
+	printf("[RUN]\tCheck dynamic fpstate cleanup across exec().\n");
+
+	before = count_dynamic_fpstates();
+	if (before < 0) {
+		printf("[SKIP]\tCannot read /proc/vmallocinfo: %s\n",
+		       strerror(-before));
+		return;
+	}
+
+	child = fork();
+	if (child < 0)
+		fatal_error("fork");
+	if (!child) {
+		snprintf(iterations_arg, sizeof(iterations_arg), "%d",
+			 EXEC_TEST_ITERS);
+		snprintf(baseline_arg, sizeof(baseline_arg), "%d", before);
+		execl("/proc/self/exe", "amx", EXEC_TEST_ARG, iterations_arg,
+		      baseline_arg, NULL);
+		fatal_error("exec");
+	}
+
+	if (waitpid(child, &status, 0) != child)
+		fatal_error("waitpid");
+	if (WIFEXITED(status) && WEXITSTATUS(status) == KSFT_SKIP) {
+		printf("[SKIP]\tDynamic fpstate allocation was not isolated.\n");
+		return;
+	}
+	if (!WIFEXITED(status) || WEXITSTATUS(status))
+		fatal_error("exec test child");
+
+	after = count_dynamic_fpstates();
+	if (after < 0) {
+		printf("[SKIP]\tCannot read /proc/vmallocinfo after exec(): %s\n",
+		       strerror(-after));
+		return;
+	}
+
+	delta = after - before;
+	if (!delta) {
+		printf("[OK]\tDynamic fpstate allocations were freed across exec().\n");
+		return;
+	}
+
+	if (delta == EXEC_TEST_ITERS)
+		errx(1, "[FAIL]\texec() leaked %d dynamic fpstate allocations",
+		     delta);
+
+	printf("[SKIP]\tDynamic fpstate vmalloc usage changed concurrently.\n");
+}
+
+int main(int argc, char **argv)
 {
 	unsigned long features;
+	int iterations = 0;
+	int baseline = 0;
+	bool exec_test;
 	long rc;
 
+	exec_test = argc == 4 && !strcmp(argv[1], EXEC_TEST_ARG);
+	if (exec_test) {
+		iterations = parse_exec_arg(argv[2]);
+		baseline = parse_exec_arg(argv[3]);
+		if (!iterations)
+			return 0;
+	}
+
 	rc = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_SUPP, &features);
 	if (rc || (features & XFEATURE_MASK_XTILE) != XFEATURE_MASK_XTILE) {
 		ksft_print_msg("no AMX support\n");
@@ -498,6 +636,10 @@ int main(void)
 	init_stashed_xsave();
 	sethandler(SIGILL, handle_noperm, 0);
 
+	if (exec_test)
+		return run_exec_test(iterations, baseline);
+
+	test_exec();
 	test_dynamic_state();
 
 	/* Request permission for the following tests */
-- 
2.50.1 (Apple Git-155)