From 4cc749881e544cc8edd5dc00b5033437cbf8b5fa Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Fri, 2 Jun 2023 09:35:49 -0700
Subject: [PATCH] cpuid: ExtendedStateSize returns the state size just for the
 enabled features

It is impossible to enable any features from user-space, so
it is no reason to return the maximum size with all features enabled.

The kvm and systrap platforms uses ExtendedStateSize to generate fpu states
that are restored via rt_sigreturn, but the kernel rejects frames with
incorrect sizes.

PiperOrigin-RevId: 537333319
---
 pkg/cpuid/cpuid_amd64.go  | 19 ++++++++++---------
 pkg/cpuid/static_amd64.go |  1 +
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/pkg/cpuid/cpuid_amd64.go b/pkg/cpuid/cpuid_amd64.go
index 187e96668..044eed079 100644
--- a/pkg/cpuid/cpuid_amd64.go
+++ b/pkg/cpuid/cpuid_amd64.go
@@ -360,22 +360,23 @@ func (fs FeatureSet) Intel() bool {
 //
 // If xSaveInfo isn't supported, cpuid will not fault but will
 // return bogus values.
-var maxXsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ecx
+var (
+	xsaveSize    = native(In{Eax: uint32(xSaveInfo)}).Ebx
+	maxXsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ecx
+)
 
 // ExtendedStateSize returns the number of bytes needed to save the "extended
-// state" for this processor and the boundary it must be aligned to. Extended
-// state includes floating point registers, and other cpu state that's not
-// associated with the normal task context.
+// state" for the enabled features and the boundary it must be aligned to.
+// Extended state includes floating point registers, and other cpu state that's
+// not associated with the normal task context.
 //
-// Note: We can save some space here with an optimization where we use a
-// smaller chunk of memory depending on features that are actually enabled.
-// Currently we just use the largest possible size for simplicity (which is
-// about 2.5K worst case, with avx512).
+// Note: the return value matches the size of signal FP state frames.
+// Look at check_xstate_in_sigframe() in the kernel sources for more details.
 //
 //go:nosplit
 func (fs FeatureSet) ExtendedStateSize() (size, align uint) {
 	if fs.UseXsave() {
-		return uint(maxXsaveSize), 64
+		return uint(xsaveSize), 64
 	}
 
 	// If we don't support xsave, we fall back to fxsave, which requires
diff --git a/pkg/cpuid/static_amd64.go b/pkg/cpuid/static_amd64.go
index 7f3ace846..09bcf16b2 100644
--- a/pkg/cpuid/static_amd64.go
+++ b/pkg/cpuid/static_amd64.go
@@ -102,6 +102,7 @@ func (s Static) normalize() {
 		in := In{Eax: uint32(xSaveInfo)}
 		out := s[in]
 		out.Ecx = maxXsaveSize
+		out.Ebx = xsaveSize
 		s[in] = out
 	}
 }