systrap: save context pointer on sysmsg

We don't need to calculate an address from context_id each time.

PiperOrigin-RevId: 518640998

pkg/sentry/platform/systrap/stub_unsafe.go

@@ -222,6 +222,8 @@ func stubInit() {
 	*p = deepSleepTimeout
 	p = (*uint64)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_handshake_timeout)))
 	*p = handshakeTimeout
+	p = (*uint64)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_context_region)))
+	*p = uint64(stubContextRegion)
 	archState := (*sysmsg.ArchState)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_arch_state)))
 	archState.Init()
 	exp := (*uint64)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_context_decoupling_exp)))

pkg/sentry/platform/systrap/subprocess.go

@@ -1039,14 +1039,13 @@ func (s *subprocess) createSysmsgThread(tregs *arch.Registers, c *context, ac *a
 	sysThread.setMsg(sysmsg.StackAddrToMsg(sentryStackAddr))
 	sysThread.msg.Init(threadID)
 	if contextDecouplingExp {
-		sysThread.msg.ContextID = invalidContextID
+		sysThread.msg.Context = 0
 	} else {
 		c.sharedContext.setThreadID(threadID)
-		sysThread.msg.ContextID = c.sharedContext.contextID
+		sysThread.msg.Context = uint64(stubContextRegion + uintptr(c.sharedContext.contextID)*sysmsg.AllocatedSizeofThreadContextStruct)
 	}
 	sysThread.msg.Self = uint64(sysmsgStackAddr + sysmsg.MsgOffsetFromSharedStack)
 	sysThread.msg.SyshandlerStack = uint64(sysmsg.StackAddrToSyshandlerStack(sysThread.sysmsgPerThreadMemAddr()))
-	sysThread.msg.ContextRegion = uint64(stubContextRegion)
 	sysThread.msg.Syshandler = uint64(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_syshandler))
 
 	sysThread.msg.State.Set(sysmsg.ThreadStateInitializing)

pkg/sentry/platform/systrap/sysmsg/build.bzl

@@ -9,6 +9,7 @@ def cc_pie_obj(name, srcs, outs):
         srcs = srcs,
         outs = outs,
         cmd = "$(CC)  $(CC_FLAGS)  " +
+              "-Wall -Werror -Wno-unused-command-line-argument " +
               "-fpie " +
               # -01 is required for clang to avoid making use of memcpy when
               # building for ARM64. For some reason when no optimization is turned

pkg/sentry/platform/systrap/sysmsg/sighandler_amd64.c

@@ -205,7 +205,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     return;
   }
 
-  struct thread_context *ctx = thread_context_addr(sysmsg);
+  struct thread_context *ctx = sysmsg->context;
 
   if (signo == SIGCHLD) {
     // If the current thread is in syshandler, an interrupt has to be postponed,
@@ -254,8 +254,6 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
 
   switch (signo) {
     case SIGSYS: {
-      int si_sysno = siginfo->si_syscall;
-      int i;
       ctx_state = CONTEXT_STATE_SYSCALL;
 
       // Check whether this syscall can be replaced on a function call or not.
@@ -354,7 +352,7 @@ void __syshandler() {
   int state = __atomic_load_n(&sysmsg->state, __ATOMIC_ACQUIRE);
   if (state != THREAD_STATE_PREP) panic(state);
 
-  struct thread_context *ctx = thread_context_addr(sysmsg);
+  struct thread_context *ctx = sysmsg->context;
 
   enum context_state ctx_state = CONTEXT_STATE_SYSCALL_TRAP;
   ctx->signo = SIGSYS;

pkg/sentry/platform/systrap/sysmsg/sighandler_arm64.c

@@ -104,7 +104,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     return;
   }
 
-  struct thread_context *ctx = thread_context_addr(sysmsg);
+  struct thread_context *ctx = sysmsg->context;
 
   uint32_t ctx_state = CONTEXT_STATE_INVALID;
   ctx->signo = signo;
@@ -181,7 +181,8 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
 void restore_state(struct sysmsg *sysmsg, struct thread_context *ctx,
                    void *_ucontext) {
   ucontext_t *ucontext = _ucontext;
-  struct fpsimd_context *fpctx = &ucontext->uc_mcontext.__reserved;
+  struct fpsimd_context *fpctx =
+      (struct fpsimd_context *)&ucontext->uc_mcontext.__reserved;
   uint8_t *fpStatePointer = (uint8_t *)&fpctx->fpsr;
 
   if (__export_context_decoupling_exp &&

pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S

@@ -18,15 +18,6 @@
 // Helper macros:
 ////////////////////////////////////////
 
-// load_thread_context loads the address of the thread context slot for the current
-// context.
-// Clobbers %rflags; loads address into %rcx.
-.macro load_thread_context_addr
-  movl %gs:offsetof_sysmsg_context_id, %ecx
-  shl $THREAD_CONTEXT_STRUCT_BITSHIFT, %rcx
-  add %gs:offsetof_sysmsg_context_region, %rcx
-.endm
-
 // prepare_enter_syshandler does the following:
 //   - saves all registers that are restorable onto the thread_context struct.
 //   - loads the address of the thread_context struct into %rcx.
@@ -37,7 +28,7 @@
   // load_thread_context_addr overwrites %rcx.
   push %rcx
 
-  load_thread_context_addr
+  movq %gs:offsetof_sysmsg_context, %rcx
 
   // Registers listed in order as written in ptregs:
   movq %r15, offsetof_thread_context_ptregs_r15(%rcx)
@@ -200,7 +191,7 @@ __export_syshandler:
 .type asm_restore_state, @function;
 asm_restore_state:
   // thread_context may have changed, therefore we reload it into %rcx anew.
-  load_thread_context_addr
+  movq %gs:offsetof_sysmsg_context, %rcx
   restore_fpstate
 
   prepare_exit_syshandler

pkg/sentry/platform/systrap/sysmsg/sysmsg.go

@@ -152,13 +152,10 @@ type Msg struct {
 	// State indicates to the sentry what the sysmsg thread is doing at a given
 	// moment.
 	State ThreadState
-	// ContextRegion defines the ThreadContext memory region start within
-	// the sysmsg thread address space.
-	ContextRegion uint64
 	// ContextID is the ID of the ThreadContext struct that the current
 	// sysmsg thread is is processing. This ID is used in the {sig|sys}handler
 	// to find the offset to the correct ThreadContext struct location.
-	ContextID uint32
+	Context uint64
 
 	// FaultJump is the size of a faulted instruction.
 	FaultJump int32
@@ -352,7 +349,7 @@ func (m *Msg) String() string {
 	fmt.Fprintf(&b, "sysmsg.Msg{msg: %x state %d", m.Self, m.State)
 	fmt.Fprintf(&b, " err %x line %d debug %x", m.Err, m.Line, m.Debug)
 	fmt.Fprintf(&b, " app stack %x", m.AppStack)
-	fmt.Fprintf(&b, " contextID %d", m.ContextID)
+	fmt.Fprintf(&b, " context %x", m.Context)
 	b.WriteString("}")
 
 	return b.String()

pkg/sentry/platform/systrap/sysmsg/sysmsg.h

@@ -16,6 +16,7 @@
 #define THIRD_PARTY_GVISOR_PKG_SENTRY_PLATFORM_SYSTRAP_SYSMSG_SYSMSG_H_
 
 #include <stdint.h>
+#include <sys/types.h>
 #include <sys/user.h>
 
 #include "sysmsg_offsets.h"  // NOLINT
@@ -48,6 +49,8 @@ enum thread_state {
   THREAD_STATE_INITIALIZING,
 };
 
+struct thread_context;
+
 // sysmsg contains the current state of the sysmsg thread. See: sysmsg.go:Msg
 struct sysmsg {
   struct sysmsg *self;
@@ -57,8 +60,7 @@ struct sysmsg {
   uint64_t app_stack;
   uint32_t interrupt;
   uint32_t state;
-  uint64_t context_region;
-  uint32_t context_id;
+  struct thread_context *context;
 
   // The fields above have offsets defined in sysmsg_offsets*.h
 
@@ -142,13 +144,6 @@ static struct sysmsg *sysmsg_addr(void *sp) {
   return (struct sysmsg *)(sp + MSG_OFFSET_FROM_START);
 }
 
-static struct thread_context *thread_context_addr(struct sysmsg *sysmsg) {
-  uint64_t tcid = __atomic_load_n(&sysmsg->context_id, __ATOMIC_ACQUIRE);
-  return (struct thread_context *)(sysmsg->context_region +
-                                   tcid *
-                                       ALLOCATED_SIZEOF_THREAD_CONTEXT_STRUCT);
-}
-
 long __syscall(long n, long a1, long a2, long a3, long a4, long a5, long a6);
 
 struct __kernel_timespec;
@@ -158,7 +153,7 @@ long sys_futex(uint32_t *addr, int op, int val, struct __kernel_timespec *tv,
 static void __panic(int err, long line) {
   void *sp = sysmsg_sp();
   struct sysmsg *sysmsg = sysmsg_addr(sp);
-  struct thread_context *ctx = thread_context_addr(sysmsg);
+  struct thread_context *ctx = sysmsg->context;
   sysmsg->err = err;
   sysmsg->err_line = line;
   // Normally sentry waits on sysmsg->state.

pkg/sentry/platform/systrap/sysmsg/sysmsg_lib.c

@@ -79,6 +79,14 @@ static __inline__ unsigned long rdtsc(void) {
 static __inline__ void spinloop(void) { asm volatile("yield" : : : "memory"); }
 #endif
 
+void *__export_context_region;
+
+static struct thread_context *thread_context_addr(uint32_t tcid) {
+  return (struct thread_context *)(__export_context_region +
+                                   tcid *
+                                       ALLOCATED_SIZEOF_THREAD_CONTEXT_STRUCT);
+}
+
 void memcpy(uint8_t *dest, uint8_t *src, size_t n) {
   for (size_t i = 0; i < n; i += 1) {
     dest[i] = src[i];
@@ -180,8 +188,8 @@ struct thread_context *queue_get_context(struct sysmsg *sysmsg) {
     if (context_id > MAX_STUB_THREADS) {
       panic(context_id);
     }
-    sysmsg->context_id = context_id;
-    struct thread_context *ctx = thread_context_addr(sysmsg);
+    struct thread_context *ctx = thread_context_addr(context_id);
+    sysmsg->context = ctx;
     __atomic_store_n(&ctx->acked, 1, __ATOMIC_RELEASE);
     __atomic_store_n(&ctx->thread_id, sysmsg->thread_id, __ATOMIC_RELEASE);
     return ctx;
@@ -235,12 +243,11 @@ struct thread_context *switch_context(struct sysmsg *sysmsg,
     }
   }
 
-  uint32_t old_ctx_id = sysmsg->context_id;
+  struct thread_context *old_ctx = sysmsg->context;
 
   ctx = get_context(sysmsg);
 
-  if (old_ctx_id != sysmsg->context_id ||
-      ctx->last_thread_id != sysmsg->thread_id) {
+  if (old_ctx != ctx || ctx->last_thread_id != sysmsg->thread_id) {
     ctx->fpstate_changed = 1;
   }
 
@@ -332,10 +339,7 @@ void verify_offsets() {
   BUILD_BUG_ON(offsetof_sysmsg_app_stack != offsetof(struct sysmsg, app_stack));
   BUILD_BUG_ON(offsetof_sysmsg_interrupt != offsetof(struct sysmsg, interrupt));
   BUILD_BUG_ON(offsetof_sysmsg_state != offsetof(struct sysmsg, state));
-  BUILD_BUG_ON(offsetof_sysmsg_context_id !=
-               offsetof(struct sysmsg, context_id));
-  BUILD_BUG_ON(offsetof_sysmsg_context_region !=
-               offsetof(struct sysmsg, context_region));
+  BUILD_BUG_ON(offsetof_sysmsg_context != offsetof(struct sysmsg, context));
 
   BUILD_BUG_ON(offsetof_thread_context_fpstate !=
                offsetof(struct thread_context, fpstate));

pkg/sentry/platform/systrap/sysmsg/sysmsg_offsets.h

@@ -39,8 +39,7 @@
 #define offsetof_sysmsg_app_stack 0x20
 #define offsetof_sysmsg_interrupt 0x28
 #define offsetof_sysmsg_state 0x2c
-#define offsetof_sysmsg_context_region 0x30
-#define offsetof_sysmsg_context_id 0x38
+#define offsetof_sysmsg_context 0x30
 
 #define offsetof_thread_context_fpstate 0x0
 #define offsetof_thread_context_fpstate_changed MAX_FPSTATE_LEN