kvm: simplify floating point logic.

This reduces the number of floating point save/restore cycles required (since
we don't need to restore immediately following the switch, this always happens
in a known context) and allows the kernel hooks to capture state. This lets us
remove calls like "Current()".

PiperOrigin-RevId: 219552844
Change-Id: I7676fa2f6c18b9919718458aa888b832a7db8cab

4 files changed, 39 insertions, 28 deletions

diff --git a/pkg/sentry/platform/kvm/bluepill_amd64.go b/pkg/sentry/platform/kvm/bluepill_amd64.go
index f013d1dc9..6520682d7 100644
--- a/pkg/sentry/platform/kvm/bluepill_amd64.go
+++ b/pkg/sentry/platform/kvm/bluepill_amd64.go
@@ -47,8 +47,8 @@ func redpill() {
 // bluepillArchEnter is called during bluepillEnter.
 //
 //go:nosplit
-func bluepillArchEnter(context *arch.SignalContext64) (c *vCPU) {
-	c = vCPUPtr(uintptr(context.Rax))
+func bluepillArchEnter(context *arch.SignalContext64) *vCPU {
+	c := vCPUPtr(uintptr(context.Rax))
 	regs := c.CPU.Registers()
 	regs.R8 = context.R8
 	regs.R9 = context.R9
@@ -73,50 +73,41 @@ func bluepillArchEnter(context *arch.SignalContext64) (c *vCPU) {
 	regs.Cs = uint64(ring0.Kcode)
 	regs.Ds = uint64(ring0.Udata)
 	regs.Es = uint64(ring0.Udata)
-	regs.Fs = uint64(ring0.Udata)
 	regs.Ss = uint64(ring0.Kdata)
-
-	// ring0 uses GS exclusively, so we use GS_base to store the location
-	// of the floating point address.
-	//
-	// The address will be restored directly after running the VCPU, and
-	// will be saved again prior to halting. We rely on the fact that the
-	// SaveFloatingPointer/LoadFloatingPoint functions use the most
-	// efficient mechanism available (including compression) so the state
-	// size is guaranteed to be less than what's pointed to here.
-	regs.Gs_base = uint64(context.Fpstate)
-	return
+	return c
 }
 
-// bluepillSyscall handles kernel syscalls.
+// KernelSyscall handles kernel syscalls.
 //
 //go:nosplit
-func bluepillSyscall() {
-	regs := ring0.Current().Registers()
+func (c *vCPU) KernelSyscall() {
+	regs := c.Registers()
 	if regs.Rax != ^uint64(0) {
 		regs.Rip -= 2 // Rewind.
 	}
-	ring0.SaveFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
+	// We only trigger a bluepill entry in the bluepill function, and can
+	// therefore be guaranteed that there is no floating point state to be
+	// loaded on resuming from halt. We only worry about saving on exit.
+	ring0.SaveFloatingPoint((*byte)(c.floatingPointState))
 	ring0.Halt()
 	ring0.WriteFS(uintptr(regs.Fs_base)) // Reload host segment.
-	ring0.LoadFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 }
 
-// bluepillException handles kernel exceptions.
+// KernelException handles kernel exceptions.
 //
 //go:nosplit
-func bluepillException(vector ring0.Vector) {
-	regs := ring0.Current().Registers()
+func (c *vCPU) KernelException(vector ring0.Vector) {
+	regs := c.Registers()
 	if vector == ring0.Vector(bounce) {
 		// These should not interrupt kernel execution; point the Rip
 		// to zero to ensure that we get a reasonable panic when we
-		// attempt to return.
+		// attempt to return and a full stack trace.
 		regs.Rip = 0
 	}
-	ring0.SaveFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
+	// See above.
+	ring0.SaveFloatingPoint((*byte)(c.floatingPointState))
 	ring0.Halt()
 	ring0.WriteFS(uintptr(regs.Fs_base)) // Reload host segment.
-	ring0.LoadFloatingPoint(bytePtr(uintptr(regs.Gs_base)))
 }
 
 // bluepillArchExit is called during bluepillEnter.
@@ -142,4 +133,9 @@ func bluepillArchExit(c *vCPU, context *arch.SignalContext64) {
 	context.Rsp = regs.Rsp
 	context.Rip = regs.Rip
 	context.Eflags = regs.Eflags
+
+	// Set the context pointer to the saved floating point state. This is
+	// where the guest data has been serialized, the kernel will restore
+	// from this new pointer value.
+	context.Fpstate = uint64(uintptrValue((*byte)(c.floatingPointState)))
 }
diff --git a/pkg/sentry/platform/kvm/bluepill_unsafe.go b/pkg/sentry/platform/kvm/bluepill_unsafe.go
index 77cf7e800..2605f8c93 100644
--- a/pkg/sentry/platform/kvm/bluepill_unsafe.go
+++ b/pkg/sentry/platform/kvm/bluepill_unsafe.go
@@ -37,6 +37,13 @@ func bytePtr(addr uintptr) *byte {
 	return (*byte)(unsafe.Pointer(addr))
 }
 
+// uintptrValue returns a uintptr for the given address.
+//
+//go:nosplit
+func uintptrValue(addr *byte) uintptr {
+	return (uintptr)(unsafe.Pointer(addr))
+}
+
 // bluepillHandler is called from the signal stub.
 //
 // The world may be stopped while this is executing, and it executes on the
diff --git a/pkg/sentry/platform/kvm/machine.go b/pkg/sentry/platform/kvm/machine.go
index 4ba3a185a..deead1b5f 100644
--- a/pkg/sentry/platform/kvm/machine.go
+++ b/pkg/sentry/platform/kvm/machine.go
@@ -142,9 +142,7 @@ func (m *machine) newVCPU() *vCPU {
 		fd:      int(fd),
 		machine: m,
 	}
-	c.CPU.Init(&m.kernel)
-	c.CPU.KernelSyscall = bluepillSyscall
-	c.CPU.KernelException = bluepillException
+	c.CPU.Init(&m.kernel, c)
 	m.vCPUsByID[c.id] = c
 
 	// Ensure the signal mask is correct.
diff --git a/pkg/sentry/platform/kvm/machine_amd64.go b/pkg/sentry/platform/kvm/machine_amd64.go
index c03792a1b..5ad805b8b 100644
--- a/pkg/sentry/platform/kvm/machine_amd64.go
+++ b/pkg/sentry/platform/kvm/machine_amd64.go
@@ -63,6 +63,10 @@ type vCPUArchState struct {
 	//
 	// This starts above fixedKernelPCID.
 	PCIDs *pagetables.PCIDs
+
+	// floatingPointState is the floating point state buffer used in guest
+	// to host transitions. See usage in bluepill_amd64.go.
+	floatingPointState *arch.FloatingPointData
 }
 
 const (
@@ -149,6 +153,12 @@ func (c *vCPU) initArchState() error {
 		return err
 	}
 
+	// Allocate some floating point state save area for the local vCPU.
+	// This will be saved prior to leaving the guest, and we restore from
+	// this always. We cannot use the pointer in the context alone because
+	// we don't know how large the area there is in reality.
+	c.floatingPointState = arch.NewFloatingPointData()
+
 	// Set the time offset to the host native time.
 	return c.setSystemTime()
 }