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

5 files changed, 64 insertions, 67 deletions

diff --git a/pkg/sentry/platform/ring0/defs.go b/pkg/sentry/platform/ring0/defs.go
index 18137e55d..98d0a6de0 100644
--- a/pkg/sentry/platform/ring0/defs.go
+++ b/pkg/sentry/platform/ring0/defs.go
@@ -38,6 +38,33 @@ type Kernel struct {
 	KernelArchState
 }
 
+// Hooks are hooks for kernel functions.
+type Hooks interface {
+	// KernelSyscall is called for kernel system calls.
+	//
+	// Return from this call will restore registers and return to the kernel: the
+	// registers must be modified directly.
+	//
+	// If this function is not provided, a kernel exception results in halt.
+	//
+	// This must be go:nosplit, as this will be on the interrupt stack.
+	// Closures are permitted, as the pointer to the closure frame is not
+	// passed on the stack.
+	KernelSyscall()
+
+	// KernelException handles an exception during kernel execution.
+	//
+	// Return from this call will restore registers and return to the kernel: the
+	// registers must be modified directly.
+	//
+	// If this function is not provided, a kernel exception results in halt.
+	//
+	// This must be go:nosplit, as this will be on the interrupt stack.
+	// Closures are permitted, as the pointer to the closure frame is not
+	// passed on the stack.
+	KernelException(Vector)
+}
+
 // CPU is the per-CPU struct.
 type CPU struct {
 	// self is a self reference.
@@ -58,29 +85,8 @@ type CPU struct {
 	// calls and exceptions via the Registers function.
 	registers syscall.PtraceRegs
 
-	// KernelException handles an exception during kernel execution.
-	//
-	// Return from this call will restore registers and return to the kernel: the
-	// registers must be modified directly.
-	//
-	// If this function is not provided, a kernel exception results in halt.
-	//
-	// This must be go:nosplit, as this will be on the interrupt stack.
-	// Closures are permitted, as the pointer to the closure frame is not
-	// passed on the stack.
-	KernelException func(Vector)
-
-	// KernelSyscall is called for kernel system calls.
-	//
-	// Return from this call will restore registers and return to the kernel: the
-	// registers must be modified directly.
-	//
-	// If this function is not provided, a kernel exception results in halt.
-	//
-	// This must be go:nosplit, as this will be on the interrupt stack.
-	// Closures are permitted, as the pointer to the closure frame is not
-	// passed on the stack.
-	KernelSyscall func()
+	// hooks are kernel hooks.
+	hooks Hooks
 }
 
 // Registers returns a modifiable-copy of the kernel registers.
diff --git a/pkg/sentry/platform/ring0/entry_amd64.s b/pkg/sentry/platform/ring0/entry_amd64.s
index d48fbd2d1..afb040a6f 100644
--- a/pkg/sentry/platform/ring0/entry_amd64.s
+++ b/pkg/sentry/platform/ring0/entry_amd64.s
@@ -90,12 +90,6 @@ TEXT ·Halt(SB),NOSPLIT,$0
 	HLT
 	RET
 
-// See kernel.go.
-TEXT ·Current(SB),NOSPLIT,$0-8
-	MOVQ CPU_SELF(GS), AX
-	MOVQ AX, ret+0(FP)
-	RET
-
 // See entry_amd64.go.
 TEXT ·swapgs(SB),NOSPLIT,$0
 	SWAP_GS()
@@ -205,19 +199,12 @@ kernel:
 	MOVQ $0, CPU_ERROR_CODE(GS) // Clear error code.
 	MOVQ $0, CPU_ERROR_TYPE(GS) // Set error type to kernel.
 
-	// Load the function stored in KernelSyscall.
-	//
-	// Note that this function needs to be executed on the stack in case
-	// the runtime decides to make use of the redzone (grumble). This also
-	// protects against any functions that might not be go:nosplit, since
-	// this will cause a failure immediately.
+	// Call the syscall trampoline.
 	LOAD_KERNEL_STACK(GS)
-	MOVQ CPU_KERNEL_SYSCALL(GS), DX // Function data.
-	MOVQ 0(DX), AX                  // Function pointer.
-	PUSHQ BP                        // Push the frame pointer.
-	MOVQ SP, BP                     // Set frame pointer value.
-	CALL *AX                        // Call the function.
-	POPQ BP                         // Restore the frame pointer.
+	MOVQ CPU_SELF(GS), AX   // Load vCPU.
+	PUSHQ AX                // First argument (vCPU).
+	CALL ·kernelSyscall(SB) // Call the trampoline.
+	POPQ AX                 // Pop vCPU.
 	JMP ·resume(SB)
 
 // exception is a generic exception handler.
@@ -287,18 +274,14 @@ kernel:
 	MOVQ 0(SP), BX              // BX contains the vector.
 	ADDQ $48, SP                // Drop the exception frame.
 
-	// Load the function stored in KernelException.
-	//
-	// See note above re: the kernel stack.
+	// Call the exception trampoline.
 	LOAD_KERNEL_STACK(GS)
-	MOVQ CPU_KERNEL_EXCEPTION(GS), DX // Function data.
-	MOVQ 0(DX), AX                    // Function pointer.
-	PUSHQ BP                          // Push the frame pointer.
-	MOVQ SP, BP                       // Set frame pointer value.
-	PUSHQ BX                          // First argument (vector).
-	CALL *AX                          // Call the function.
-	POPQ BX                           // Discard the argument.
-	POPQ BP                           // Restore the frame pointer.
+	MOVQ CPU_SELF(GS), AX     // Load vCPU.
+	PUSHQ BX                  // Second argument (vector).
+	PUSHQ AX                  // First argument (vCPU).
+	CALL ·kernelException(SB) // Call the trampoline.
+	POPQ BX                   // Pop vector.
+	POPQ AX                   // Pop vCPU.
 	JMP ·resume(SB)
 
 #define EXCEPTION_WITH_ERROR(value, symbol) \
diff --git a/pkg/sentry/platform/ring0/kernel.go b/pkg/sentry/platform/ring0/kernel.go
index e70eafde2..19ac6eb7c 100644
--- a/pkg/sentry/platform/ring0/kernel.go
+++ b/pkg/sentry/platform/ring0/kernel.go
@@ -26,31 +26,41 @@ func (k *Kernel) Init(opts KernelOpts) {
 // Halt halts execution.
 func Halt()
 
-// Current returns the current CPU.
+// defaultHooks implements hooks.
+type defaultHooks struct{}
+
+// KernelSyscall implements Hooks.KernelSyscall.
 //
-// Its use is only legal in the KernelSyscall and KernelException contexts,
-// which must all be guarded go:nosplit.
-func Current() *CPU
+//go:nosplit
+func (defaultHooks) KernelSyscall() { Halt() }
+
+// KernelException implements Hooks.KernelException.
+//
+//go:nosplit
+func (defaultHooks) KernelException(Vector) { Halt() }
 
-// defaultSyscall is the default syscall hook.
+// kernelSyscall is a trampoline.
 //
 //go:nosplit
-func defaultSyscall() { Halt() }
+func kernelSyscall(c *CPU) { c.hooks.KernelSyscall() }
 
-// defaultException is the default exception hook.
+// kernelException is a trampoline.
 //
 //go:nosplit
-func defaultException(Vector) { Halt() }
+func kernelException(c *CPU, vector Vector) { c.hooks.KernelException(vector) }
 
 // Init initializes a new CPU.
 //
 // Init allows embedding in other objects.
-func (c *CPU) Init(k *Kernel) {
+func (c *CPU) Init(k *Kernel, hooks Hooks) {
 	c.self = c   // Set self reference.
 	c.kernel = k // Set kernel reference.
 	c.init()     // Perform architectural init.
 
-	// Defaults.
-	c.KernelSyscall = defaultSyscall
-	c.KernelException = defaultException
+	// Require hooks.
+	if hooks != nil {
+		c.hooks = hooks
+	} else {
+		c.hooks = defaultHooks{}
+	}
 }
diff --git a/pkg/sentry/platform/ring0/kernel_amd64.go b/pkg/sentry/platform/ring0/kernel_amd64.go
index ab562bca7..9e8c56a54 100644
--- a/pkg/sentry/platform/ring0/kernel_amd64.go
+++ b/pkg/sentry/platform/ring0/kernel_amd64.go
@@ -204,7 +204,7 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 func start(c *CPU) {
 	// Save per-cpu & FS segment.
 	WriteGS(kernelAddr(c))
-	WriteFS(uintptr(c.Registers().Fs_base))
+	WriteFS(uintptr(c.registers.Fs_base))
 
 	// Initialize floating point.
 	//
diff --git a/pkg/sentry/platform/ring0/offsets_amd64.go b/pkg/sentry/platform/ring0/offsets_amd64.go
index 753d31ef8..806e07ec0 100644
--- a/pkg/sentry/platform/ring0/offsets_amd64.go
+++ b/pkg/sentry/platform/ring0/offsets_amd64.go
@@ -34,8 +34,6 @@ func Emit(w io.Writer) {
 	fmt.Fprintf(w, "#define CPU_STACK_TOP        0x%02x\n", reflect.ValueOf(&c.stack[0]).Pointer()-reflect.ValueOf(c).Pointer()+uintptr(len(c.stack)))
 	fmt.Fprintf(w, "#define CPU_ERROR_CODE       0x%02x\n", reflect.ValueOf(&c.errorCode).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_ERROR_TYPE       0x%02x\n", reflect.ValueOf(&c.errorType).Pointer()-reflect.ValueOf(c).Pointer())
-	fmt.Fprintf(w, "#define CPU_KERNEL_EXCEPTION 0x%02x\n", reflect.ValueOf(&c.KernelException).Pointer()-reflect.ValueOf(c).Pointer())
-	fmt.Fprintf(w, "#define CPU_KERNEL_SYSCALL   0x%02x\n", reflect.ValueOf(&c.KernelSyscall).Pointer()-reflect.ValueOf(c).Pointer())
 
 	fmt.Fprintf(w, "\n// Bits.\n")
 	fmt.Fprintf(w, "#define _RFLAGS_IF           0x%02x\n", _RFLAGS_IF)