Restore FS on resume.

Previously, the vCPU FS was always correct because it relied on the
reset coming out of the switch. When that doesn't occur, for example,
using bluepill directly, the FS value can be incorrect leading to
strange corruption.

This change is necessary for a subsequent change that enforces guest
mode for page table modifications, and it may reduce test flakiness.
(The problematic path may occur in tests, but does not occur in the
actual platform.)

PiperOrigin-RevId: 198648137
Change-Id: I513910a973dd8666c9a1d18cf78990964d6a644d

3 files changed, 15 insertions, 13 deletions

diff --git a/pkg/sentry/platform/kvm/bluepill_amd64.go b/pkg/sentry/platform/kvm/bluepill_amd64.go
index a2baefb7d..b364e3ef7 100644
--- a/pkg/sentry/platform/kvm/bluepill_amd64.go
+++ b/pkg/sentry/platform/kvm/bluepill_amd64.go
@@ -98,6 +98,7 @@ func bluepillSyscall() {
 	}
 	ring0.SaveFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 	ring0.Halt()
+	ring0.WriteFS(uintptr(regs.Fs_base)) // Reload host segment.
 	ring0.LoadFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 }
 
@@ -114,6 +115,7 @@ func bluepillException(vector ring0.Vector) {
 	}
 	ring0.SaveFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 	ring0.Halt()
+	ring0.WriteFS(uintptr(regs.Fs_base)) // Reload host segment.
 	ring0.LoadFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 }
 
diff --git a/pkg/sentry/platform/ring0/kernel_amd64.go b/pkg/sentry/platform/ring0/kernel_amd64.go
index 02d6d0de4..58ac4b4b2 100644
--- a/pkg/sentry/platform/ring0/kernel_amd64.go
+++ b/pkg/sentry/platform/ring0/kernel_amd64.go
@@ -200,8 +200,8 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 
 	// Perform the switch.
 	swapgs()                                         // GS will be swapped on return.
-	wrfs(uintptr(regs.Fs_base))                      // Set application FS.
-	wrgs(uintptr(regs.Gs_base))                      // Set application GS.
+	WriteFS(uintptr(regs.Fs_base))                   // Set application FS.
+	WriteGS(uintptr(regs.Gs_base))                   // Set application GS.
 	LoadFloatingPoint(switchOpts.FloatingPointState) // Copy in floating point.
 	jumpToKernel()                                   // Switch to upper half.
 	writeCR3(uintptr(userCR3))                       // Change to user address space.
@@ -213,7 +213,7 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	writeCR3(uintptr(kernelCR3))                     // Return to kernel address space.
 	jumpToUser()                                     // Return to lower half.
 	SaveFloatingPoint(switchOpts.FloatingPointState) // Copy out floating point.
-	wrfs(uintptr(c.registers.Fs_base))               // Restore kernel FS.
+	WriteFS(uintptr(c.registers.Fs_base))            // Restore kernel FS.
 	return
 }
 
@@ -225,8 +225,8 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 //go:nosplit
 func start(c *CPU) {
 	// Save per-cpu & FS segment.
-	wrgs(kernelAddr(c))
-	wrfs(uintptr(c.Registers().Fs_base))
+	WriteGS(kernelAddr(c))
+	WriteFS(uintptr(c.Registers().Fs_base))
 
 	// Initialize floating point.
 	//
diff --git a/pkg/sentry/platform/ring0/lib_amd64.go b/pkg/sentry/platform/ring0/lib_amd64.go
index f1ed5bfb4..de2842b5a 100644
--- a/pkg/sentry/platform/ring0/lib_amd64.go
+++ b/pkg/sentry/platform/ring0/lib_amd64.go
@@ -43,8 +43,8 @@ func xsave(*byte)
 // xsaveopt uses xsaveopt to save floating point state.
 func xsaveopt(*byte)
 
-// wrfs sets the GS address (set by init).
-var wrfs func(addr uintptr)
+// WriteFS sets the GS address (set by init).
+var WriteFS func(addr uintptr)
 
 // wrfsbase writes to the GS base address.
 func wrfsbase(addr uintptr)
@@ -52,8 +52,8 @@ func wrfsbase(addr uintptr)
 // wrfsmsr writes to the GS_BASE MSR.
 func wrfsmsr(addr uintptr)
 
-// wrgs sets the GS address (set by init).
-var wrgs func(addr uintptr)
+// WriteGS sets the GS address (set by init).
+var WriteGS func(addr uintptr)
 
 // wrgsbase writes to the GS base address.
 func wrgsbase(addr uintptr)
@@ -119,10 +119,10 @@ func Init(featureSet *cpuid.FeatureSet) {
 		LoadFloatingPoint = fxrstor
 	}
 	if hasFSGSBASE {
-		wrfs = wrfsbase
-		wrgs = wrgsbase
+		WriteFS = wrfsbase
+		WriteGS = wrgsbase
 	} else {
-		wrfs = wrfsmsr
-		wrgs = wrgsmsr
+		WriteFS = wrfsmsr
+		WriteGS = wrgsmsr
 	}
 }