Move the code that manages floating-point state to a separate package

This change is inspired by Adin's cl/355256448.

PiperOrigin-RevId: 364695931

5 files changed, 12 insertions, 9 deletions

diff --git a/pkg/ring0/BUILD b/pkg/ring0/BUILD
index d1b14efdb..885958456 100644
--- a/pkg/ring0/BUILD
+++ b/pkg/ring0/BUILD
@@ -80,6 +80,7 @@ go_library(
         "//pkg/ring0/pagetables",
         "//pkg/safecopy",
         "//pkg/sentry/arch",
+        "//pkg/sentry/arch/fpu",
         "//pkg/usermem",
     ],
 )
diff --git a/pkg/ring0/defs.go b/pkg/ring0/defs.go
index e2561e4c2..b6e2012e8 100644
--- a/pkg/ring0/defs.go
+++ b/pkg/ring0/defs.go
@@ -17,6 +17,7 @@ package ring0
 import (
 	"gvisor.dev/gvisor/pkg/ring0/pagetables"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/arch/fpu"
 )
 
 // Kernel is a global kernel object.
@@ -96,7 +97,7 @@ type SwitchOpts struct {
 
 	// FloatingPointState is a byte pointer where floating point state is
 	// saved and restored.
-	FloatingPointState *byte
+	FloatingPointState *fpu.State
 
 	// PageTables are the application page tables.
 	PageTables *pagetables.PageTables
diff --git a/pkg/ring0/gen_offsets/BUILD b/pkg/ring0/gen_offsets/BUILD
index 15b93d61c..f421e1687 100644
--- a/pkg/ring0/gen_offsets/BUILD
+++ b/pkg/ring0/gen_offsets/BUILD
@@ -35,6 +35,7 @@ go_binary(
         "//pkg/cpuid",
         "//pkg/ring0/pagetables",
         "//pkg/sentry/arch",
+        "//pkg/sentry/arch/fpu",
         "//pkg/usermem",
     ],
 )
diff --git a/pkg/ring0/kernel_amd64.go b/pkg/ring0/kernel_amd64.go
index 36a60700e..33c259757 100644
--- a/pkg/ring0/kernel_amd64.go
+++ b/pkg/ring0/kernel_amd64.go
@@ -239,17 +239,17 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	regs.Ss = uint64(Udata)   // Ditto.
 
 	// Perform the switch.
-	swapgs()                                         // GS will be swapped on return.
-	WriteFS(uintptr(regs.Fs_base))                   // escapes: no. Set application FS.
-	WriteGS(uintptr(regs.Gs_base))                   // escapes: no. Set application GS.
-	LoadFloatingPoint(switchOpts.FloatingPointState) // escapes: no. Copy in floating point.
+	swapgs()                                                       // GS will be swapped on return.
+	WriteFS(uintptr(regs.Fs_base))                                 // escapes: no. Set application FS.
+	WriteGS(uintptr(regs.Gs_base))                                 // escapes: no. Set application GS.
+	LoadFloatingPoint(switchOpts.FloatingPointState.BytePointer()) // escapes: no. Copy in floating point.
 	if switchOpts.FullRestore {
 		vector = iret(c, regs, uintptr(userCR3))
 	} else {
 		vector = sysret(c, regs, uintptr(userCR3))
 	}
-	SaveFloatingPoint(switchOpts.FloatingPointState) // escapes: no. Copy out floating point.
-	WriteFS(uintptr(c.registers.Fs_base))            // escapes: no. Restore kernel FS.
+	SaveFloatingPoint(switchOpts.FloatingPointState.BytePointer()) // escapes: no. Copy out floating point.
+	WriteFS(uintptr(c.registers.Fs_base))                          // escapes: no. Restore kernel FS.
 	return
 }
 
diff --git a/pkg/ring0/kernel_arm64.go b/pkg/ring0/kernel_arm64.go
index 41909b3a0..7975e5f92 100644
--- a/pkg/ring0/kernel_arm64.go
+++ b/pkg/ring0/kernel_arm64.go
@@ -62,7 +62,7 @@ func IsCanonical(addr uint64) bool {
 //go:nosplit
 func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	storeAppASID(uintptr(switchOpts.UserASID))
-	storeEl0Fpstate(switchOpts.FloatingPointState)
+	storeEl0Fpstate(switchOpts.FloatingPointState.BytePointer())
 
 	if switchOpts.Flush {
 		FlushTlbByASID(uintptr(switchOpts.UserASID))
@@ -82,7 +82,7 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 
 	fpDisableTrap = CPACREL1()
 	if fpDisableTrap != 0 {
-		SaveFloatingPoint(switchOpts.FloatingPointState)
+		SaveFloatingPoint(switchOpts.FloatingPointState.BytePointer())
 	}
 
 	vector = c.vecCode