From 9aca9a74291fc7d884df9bd106c0f064751f3665 Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Thu, 2 Nov 2023 11:18:20 -0700
Subject: [PATCH] Elide WriteGS if the GS is already correct.

This caches the GS value in the ring0.CPU struct, in order to avoid a
full pipeline flush in the switch path.

Fixes #2010

Co-authored-by: Adin Scannell <adin@scannell.ca>
Co-authored-by: Lai Jiangshan <jiangshan.ljs@antfin.com>
PiperOrigin-RevId: 578919437
---
 pkg/ring0/defs_amd64.go   |  3 +++
 pkg/ring0/entry_amd64.s   | 36 +++++++++++++++++++++++++++++++-----
 pkg/ring0/kernel_amd64.go |  1 +
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/pkg/ring0/defs_amd64.go b/pkg/ring0/defs_amd64.go
index a9cfeee46..cb928327b 100644
--- a/pkg/ring0/defs_amd64.go
+++ b/pkg/ring0/defs_amd64.go
@@ -137,10 +137,13 @@ type CPUArchState struct {
 
 	*kernelEntry
 
+	appGsBase uint64
+
 	// Copies of global variables, stored in CPU so that they can be used by
 	// syscall and exception handlers (in the upper address space).
 	hasXSAVE    bool
 	hasXSAVEOPT bool
+	hasFSGSBASE bool
 }
 
 // ErrorCode returns the last error code.
diff --git a/pkg/ring0/entry_amd64.s b/pkg/ring0/entry_amd64.s
index 020b9b941..d2f870a52 100644
--- a/pkg/ring0/entry_amd64.s
+++ b/pkg/ring0/entry_amd64.s
@@ -16,16 +16,18 @@
 #include "textflag.h"
 
 // CPU offsets.
-#define CPU_REGISTERS    64  // +checkoffset . CPU.registers
-#define CPU_FPU_STATE    280 // +checkoffset . CPU.floatingPointState
+#define CPU_REGISTERS    72  // +checkoffset . CPU.registers
+#define CPU_FPU_STATE    288 // +checkoffset . CPU.floatingPointState
 #define CPU_ARCH_STATE   16  // +checkoffset . CPU.CPUArchState
 #define CPU_ERROR_CODE   CPU_ARCH_STATE+0  // +checkoffset . CPUArchState.errorCode
 #define CPU_ERROR_TYPE   CPU_ARCH_STATE+8  // +checkoffset . CPUArchState.errorType
 #define CPU_VECTOR       CPU_ARCH_STATE+16 // +checkoffset . CPUArchState.vector
 #define CPU_FAULT_ADDR   CPU_ARCH_STATE+24 // +checkoffset . CPUArchState.faultAddr
 #define CPU_ENTRY        CPU_ARCH_STATE+32 // +checkoffset . CPUArchState.kernelEntry
-#define CPU_HAS_XSAVE    CPU_ARCH_STATE+40 // +checkoffset . CPUArchState.hasXSAVE
-#define CPU_HAS_XSAVEOPT CPU_ARCH_STATE+41 // +checkoffset . CPUArchState.hasXSAVEOPT
+#define CPU_APP_GS_BASE        CPU_ARCH_STATE+40 // +checkoffset . CPUArchState.appGsBase
+#define CPU_HAS_XSAVE    CPU_ARCH_STATE+48 // +checkoffset . CPUArchState.hasXSAVE
+#define CPU_HAS_XSAVEOPT CPU_ARCH_STATE+49 // +checkoffset . CPUArchState.hasXSAVEOPT
+#define CPU_HAS_FSGSBASE CPU_ARCH_STATE+50 // +checkoffset . CPUArchState.hasFSGSBASE
 
 #define ENTRY_SCRATCH0   256 // +checkoffset . kernelEntry.scratch0
 #define ENTRY_STACK_TOP  264 // +checkoffset . kernelEntry.stackTop
@@ -240,10 +242,13 @@ fprestore_done:
 	MOVQ regs+8(FP), R8
 	SWAP_GS()
 	MOVQ PTRACE_GS_BASE(R8), AX
+	CMPQ AX, CPU_APP_GS_BASE(SI)
+	JE skip_gs
+	MOVQ AX, CPU_APP_GS_BASE(SI)
 	PUSHQ AX
 	CALL ·writeGS(SB)
 	POPQ AX
-
+skip_gs:
 	// Call sysret() or iret().
 	MOVQ userCR3+24(FP), CX
 	MOVQ needIRET+32(FP), R9
@@ -411,6 +416,12 @@ TEXT ·start(SB),NOSPLIT|NOFRAME,$0
 	CALL ·writeFS(SB)
 	POPQ BX
 
+	MOVQ CPU_APP_GS_BASE(AX),BX
+	PUSHQ BX
+	CALL ·writeGS(SB)
+	POPQ BX
+	SWAP_GS()
+
 	// First argument (CPU) already at bottom of stack.
 	CALL ·startGo(SB) // Call Go hook.
 	JMP ·resume(SB)   // Restore to registers.
@@ -439,6 +450,14 @@ user:
 	MOVQ CX,  PTRACE_RAX(AX)               // Save everything else.
 	MOVQ CX,  PTRACE_ORIGRAX(AX)
 
+	CMPB CPU_HAS_FSGSBASE(GS), $1
+	JNE sysenter_skip_gs
+	SWAP_GS()
+	BYTE $0xf3; BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0xcb; // rdgsbase rbx
+	MOVQ BX, PTRACE_GS_BASE(AX)
+	SWAP_GS()
+
+sysenter_skip_gs:
 	MOVQ ENTRY_CPU_SELF(GS), AX            // Load vCPU.
 	MOVQ CPU_REGISTERS+PTRACE_RSP(AX), SP  // Get stacks.
 	MOVQ $0, CPU_ERROR_CODE(AX)            // Clear error code.
@@ -561,6 +580,13 @@ user:
 	POPQ BX                                 // Restore original AX.
 	MOVQ BX, PTRACE_RAX(AX)                 // Save it.
 	MOVQ BX, PTRACE_ORIGRAX(AX)
+	CMPB CPU_HAS_FSGSBASE(GS), $1
+	JNE exception_skip_gs
+	SWAP_GS()
+	BYTE $0xf3; BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0xcb; // rdgsbase rbx
+	MOVQ BX, PTRACE_GS_BASE(AX)
+	SWAP_GS()
+exception_skip_gs:
 	MOVQ 16(SP), BX; MOVQ BX, PTRACE_RIP(AX)
 	MOVQ 24(SP), CX; MOVQ CX, PTRACE_CS(AX)
 	MOVQ 32(SP), DX; MOVQ DX, PTRACE_FLAGS(AX)
diff --git a/pkg/ring0/kernel_amd64.go b/pkg/ring0/kernel_amd64.go
index 0d7f34106..a0428460a 100644
--- a/pkg/ring0/kernel_amd64.go
+++ b/pkg/ring0/kernel_amd64.go
@@ -144,6 +144,7 @@ func (c *CPU) init(cpuID int) {
 
 	c.hasXSAVE = hasXSAVE
 	c.hasXSAVEOPT = hasXSAVEOPT
+	c.hasFSGSBASE = hasFSGSBASE
 }
 
 // StackTop returns the kernel's stack address.