Refactor context-related shared memory usage.

This change introduces an abstraction for most accesses to shared thread-context
memory. In general, there are very few cases where accessing this memory is not
supposed to be atomic, so it makes sense to abstract these accesses into
getters/setters that perform the actions atomically. After this change, we
should treat most direct accesses through sharedContext.shared as suspect.

Additionally this cleanup allows the new sharedContext instance to become the
context interruptor. When doing this it is no longer required to use locks, as
was done in context.NotifyInterrupt.

PiperOrigin-RevId: 517166307

pkg/sentry/platform/systrap/BUILD

@@ -28,15 +28,14 @@ go_library(
         "filters_arm64.go",
         "lib_amd64.s",
         "lib_arm64.s",
+        "shared_context.go",
         "stub_amd64.s",
         "stub_arm64.s",
         "stub_defs.go",
         "stub_unsafe.go",
         "subprocess.go",
         "subprocess_amd64.go",
-        "subprocess_amd64_unsafe.go",
         "subprocess_arm64.go",
-        "subprocess_arm64_unsafe.go",
         "subprocess_linux.go",
         "subprocess_linux_unsafe.go",
         "subprocess_pool.go",

pkg/sentry/platform/systrap/shared_context.go

@@ -0,0 +1,168 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package systrap
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/sentry/platform"
+	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
+)
+
+const (
+	ackReset uint32 = 0
+)
+
+// sharedContext is an abstraction for interactions that the sentry has to
+// perform with memory shared between it and the stub threads used for contexts.
+//
+// Any access to shared memory should most likely have a getter/setter through
+// this struct. This is due to the following reasons:
+//   - The memory needs to be read or modified atomically because there is no
+//     (trusted) synchronization between the sentry and the stub processes.
+//   - Data read from shared memory may require validation before it can be used.
+type sharedContext struct {
+	// subprocess is the subprocess that this sharedContext instance belongs to.
+	subprocess *subprocess
+	// contextID is the ID corresponding to the sysmsg.ThreadContext memory slot
+	// that is used for this sharedContext.
+	contextID uint64
+	// shared is the handle to the shared memory that the sentry task goroutine
+	// reads from and writes to.
+	// NOTE: Using this handle directly without a getter from this function should
+	//       most likely be avoided due to concerns listed above.
+	shared *sysmsg.ThreadContext
+}
+
+func (s *subprocess) getSharedContext() (*sharedContext, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	id, ok := s.threadContextPool.Get()
+	if !ok {
+		return nil, fmt.Errorf("subprocess has too many active tasks (%d); failed to create a new one", maxGuestContexts)
+	}
+	s.IncRef()
+	sc := sharedContext{
+		subprocess: s,
+		contextID:  id,
+		shared:     s.getThreadContextFromID(id),
+	}
+	sc.shared.Init(invalidThreadID)
+
+	return &sc, nil
+}
+
+func (sc *sharedContext) release() {
+	if sc == nil {
+		return
+	}
+	sc.subprocess.threadContextPool.Put(sc.contextID)
+	sc.subprocess.DecRef(sc.subprocess.release)
+}
+
+func (sc *sharedContext) isActiveInSubprocess(s *subprocess) bool {
+	if sc == nil {
+		return false
+	}
+	return sc.subprocess == s
+}
+
+// NotifyInterrupt implements interrupt.Receiver.NotifyInterrupt.
+func (sc *sharedContext) NotifyInterrupt() {
+	// If this context is not being worked on right now we need to mark it as
+	// interrupted so the next executor does not start working on it.
+	atomic.StoreUint32(&sc.shared.Interrupt, 1)
+	if sc.threadID() == invalidThreadID {
+		return
+	}
+	sc.subprocess.sysmsgThreadsMu.Lock()
+	defer sc.subprocess.sysmsgThreadsMu.Unlock()
+
+	threadID := atomic.LoadUint32(&sc.shared.ThreadID)
+	sysmsgThread, ok := sc.subprocess.sysmsgThreads[threadID]
+	if !ok {
+		// This is either an invalidThreadID or another garbage value; either way we
+		// don't know which thread to interrupt; best we can do is mark the context.
+		return
+	}
+
+	t := sysmsgThread.thread
+	atomic.StoreUint64(&sysmsgThread.msg.InterruptedContextID, sc.contextID)
+	if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(t.tgid), uintptr(t.tid), uintptr(platform.SignalInterrupt)); e != 0 {
+		panic(fmt.Sprintf("failed to interrupt the child process %d: %v", t.tid, e))
+	}
+}
+
+func (sc *sharedContext) state() sysmsg.ContextState {
+	return sc.shared.State.Get()
+}
+
+func (sc *sharedContext) setState(state sysmsg.ContextState) {
+	sc.shared.State.Set(state)
+}
+
+func (sc *sharedContext) setInterrupt() {
+	atomic.StoreUint32(&sc.shared.Interrupt, 1)
+}
+
+func (sc *sharedContext) clearInterrupt() {
+	atomic.StoreUint32(&sc.shared.Interrupt, 0)
+}
+
+func (sc *sharedContext) setFPStateChanged() {
+	atomic.StoreUint64(&sc.shared.FPStateChanged, 1)
+}
+
+func (sc *sharedContext) threadID() uint32 {
+	return atomic.LoadUint32(&sc.shared.ThreadID)
+}
+
+func (sc *sharedContext) setThreadID(threadID uint32) {
+	if contextDecouplingExp {
+		panic("context decoupled systrap should never explicitly set ThreadID")
+	}
+	atomic.StoreUint32(&sc.shared.ThreadID, threadID)
+}
+
+// EnableSentryFastPath indicates that the polling mode is enabled for the
+// Sentry. It has to be called before putting the context into the context queue.
+// This function is used if contextDecouplingExp=true because the fastpath
+// is negotiated in ThreadContext.
+func (sc *sharedContext) enableSentryFastPath() {
+	atomic.StoreUint32(&sc.shared.SentryFastPath, 1)
+}
+
+// DisableSentryFastPath indicates that the polling mode for the sentry is
+// disabled for the Sentry.
+// This function is used if contextDecouplingExp=true because the fastpath
+// is negotiated in ThreadContext.
+func (sc *sharedContext) disableSentryFastPath() {
+	atomic.StoreUint32(&sc.shared.SentryFastPath, 0)
+}
+
+func (sc *sharedContext) isAcked() bool {
+	return atomic.LoadUint32(&sc.shared.Acked) != ackReset
+}
+
+func (sc *sharedContext) resetAcked() {
+	atomic.StoreUint32(&sc.shared.Acked, ackReset)
+}
+
+func (sc *sharedContext) sleepOnState(state sysmsg.ContextState) {
+	sc.shared.SleepOnState(state)
+}

pkg/sentry/platform/systrap/subprocess.go

@@ -680,9 +680,9 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 	// Reset necessary registers.
 	regs := &ac.StateData().Regs
 	s.resetSysemuRegs(regs)
-	ctx := s.getThreadContextFromID(c.cid)
-	ctx.Regs = regs.PtraceRegs
-	restoreArchSpecificState(ctx, ac)
+	ctx := c.sharedContext
+	ctx.shared.Regs = regs.PtraceRegs
+	restoreArchSpecificState(ctx.shared, ac)
 
 	// Get sysmsg thread bound to the context; no-op if contextDecoupling is on.
 	sysThread, err := s.getSysmsgThread(regs, c, ac)
@@ -691,36 +691,36 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 	}
 
 	// Check for interrupts, and ensure that future interrupts signal the context.
-	if !c.interrupt.Enable(c) {
+	if !c.interrupt.Enable(c.sharedContext) {
 		// Pending interrupt; simulate.
-		ctx.Interrupt = 0
+		ctx.clearInterrupt()
 		c.signalInfo = linux.SignalInfo{Signo: int32(platform.SignalInterrupt)}
 		return false, false, nil
 	}
 	defer c.interrupt.Disable()
 
 	if contextDecouplingExp {
-		s.restoreFPState(nil, ctx, 0, c, ac)
+		restoreFPState(nil, ctx, 0, c, ac)
 
 		// Place the context onto the context queue.
-		ctx.State.Set(sysmsg.ContextStateNone)
-		s.contextQueue.add(uint32(c.cid))
+		ctx.setState(sysmsg.ContextStateNone)
+		s.contextQueue.add(uint32(ctx.contextID))
 		s.waitOnState(ctx)
 
 		// Check if there's been an error.
-		tid := atomic.LoadUint32(&ctx.ThreadID)
-		if tid != invalidThreadID {
-			if sysThread, ok := s.sysmsgThreads[tid]; ok && sysThread.msg.Err != 0 {
+		threadID := ctx.threadID()
+		if threadID != invalidThreadID {
+			if sysThread, ok := s.sysmsgThreads[threadID]; ok && sysThread.msg.Err != 0 {
 				msg := sysThread.msg
 				panic(fmt.Sprintf("stub thread %d failed: err 0x%x line %d: %s", sysThread.thread.tid, msg.Err, msg.Line, msg))
 			}
-			log.Warningf("systrap: found unexpected ThreadContext.ThreadID field, expected %d found %d", invalidThreadID, tid)
+			log.Warningf("systrap: found unexpected ThreadContext.ThreadID field, expected %d found %d", invalidThreadID, threadID)
 		}
 	} else {
 		msg := sysThread.msg
 		t := sysThread.thread
 
-		s.restoreFPState(msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
+		restoreFPState(msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
 
 		msg.EnableSentryFastPath()
 		sysThread.waitEvent(sysmsg.ThreadStateDone)
@@ -730,7 +730,7 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 			panic(fmt.Sprintf("stub thread %d failed: err %d line %d: %s", t.tid, msg.Err, msg.Line, msg))
 		}
 
-		if ctx.State != sysmsg.ContextStateSyscallTrap {
+		if ctx.state() != sysmsg.ContextStateSyscallTrap {
 			var err error
 			sysThread.fpuStateToMsgOffset, err = msg.FPUStateOffset()
 			if err != nil {
@@ -738,27 +738,28 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 			}
 		}
 
-		retrieveArchSpecificState(ctx, ac)
+		retrieveArchSpecificState(ctx.shared, ac)
 	}
 
-	regs.PtraceRegs = ctx.Regs
+	regs.PtraceRegs = ctx.shared.Regs
 	// We have a signal. We verify however, that the signal was
 	// either delivered from the kernel or from this process. We
 	// don't respect other signals.
-	c.signalInfo = ctx.SignalInfo
-	if ctx.State == sysmsg.ContextStateSyscallCanBePatched {
-		ctx.State = sysmsg.ContextStateSyscall
+	c.signalInfo = ctx.shared.SignalInfo
+	ctxState := ctx.state()
+	if ctxState == sysmsg.ContextStateSyscallCanBePatched {
+		ctxState = sysmsg.ContextStateSyscall
 		shouldPatchSyscall = true
 	}
 
-	if ctx.State == sysmsg.ContextStateSyscall || ctx.State == sysmsg.ContextStateSyscallTrap {
+	if ctxState == sysmsg.ContextStateSyscall || ctxState == sysmsg.ContextStateSyscallTrap {
 		if maybePatchSignalInfo(regs, &c.signalInfo) {
 			return false, false, nil
 		}
 		updateSyscallRegs(regs)
 		return true, shouldPatchSyscall, nil
-	} else if ctx.State != sysmsg.ContextStateFault {
-		panic(fmt.Sprintf("unknown context state: %v", ctx.State))
+	} else if ctxState != sysmsg.ContextStateFault {
+		panic(fmt.Sprintf("unknown context state: %v", ctxState))
 	}
 
 	return false, false, nil
@@ -773,23 +774,21 @@ const (
 	threadKickTimeout = uint64(20000)
 )
 
-func (s *subprocess) waitOnState(ctx *sysmsg.ThreadContext) {
-	// ackedEvents is always reset to 0 at the end of this function.
-	ackedEvents := uint32(0)
+func (s *subprocess) waitOnState(ctx *sharedContext) {
 	kicked := false
 	slowPath := false
 	start := cputicks()
 	handshake := false
-	for curState := ctx.State.Get(); curState == sysmsg.ContextStateNone; curState = ctx.State.Get() {
+	for curState := ctx.state(); curState == sysmsg.ContextStateNone; curState = ctx.state() {
 		if !slowPath {
 			delta := uint64(cputicks() - start)
 			if delta > decoupledDeepSleepTimeout {
-				ctx.DisableSentryFastPath()
+				ctx.disableSentryFastPath()
 				slowPath = true
 				continue
 			}
 
-			if !handshake && ackedEvents != atomic.LoadUint32(&ctx.Acked) {
+			if !handshake && ctx.isAcked() {
 				handshake = true
 				continue
 			}
@@ -802,12 +801,12 @@ func (s *subprocess) waitOnState(ctx *sysmsg.ThreadContext) {
 				s.kickSysmsgThread()
 			}
 
-			ctx.SleepOnState(curState)
+			ctx.sleepOnState(curState)
 		}
 	}
 
-	atomic.StoreUint32(&ctx.Acked, 0)
-	ctx.EnableSentryFastPath()
+	ctx.resetAcked()
+	ctx.enableSentryFastPath()
 }
 
 func (s *subprocess) kickSysmsgThread() {
@@ -894,18 +893,17 @@ func (s *subprocess) Unmap(addr hostarch.Addr, length uint64) {
 }
 
 func (s *subprocess) PullFullState(c *context, ac *arch.Context64) error {
-	if s != c.subprocess {
+	if !c.sharedContext.isActiveInSubprocess(s) {
 		panic("Attempted to PullFullState for context that is not used in subprocess")
 	}
-	ctx := s.getThreadContextFromID(c.cid)
 	if contextDecouplingExp {
-		s.saveFPState(nil, ctx, 0, c, ac)
+		saveFPState(nil, c.sharedContext, 0, c, ac)
 	} else {
 		sysThread, err := s.getSysmsgThread(&ac.StateData().Regs, c, ac)
 		if err != nil {
 			return err
 		}
-		s.saveFPState(sysThread.msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
+		saveFPState(sysThread.msg, c.sharedContext, sysThread.fpuStateToMsgOffset, c, ac)
 	}
 	return nil
 }
@@ -1023,8 +1021,8 @@ func (s *subprocess) createSysmsgThread(tregs *arch.Registers, c *context, ac *a
 	if contextDecouplingExp {
 		sysThread.msg.ContextID = uint64(invalidContextID)
 	} else {
-		s.getThreadContextFromID(c.cid).ThreadID = threadID
-		sysThread.msg.ContextID = c.cid
+		c.sharedContext.setThreadID(threadID)
+		sysThread.msg.ContextID = c.sharedContext.contextID
 	}
 	sysThread.msg.Self = uint64(sysmsgStackAddr + sysmsg.MsgOffsetFromSharedStack)
 	sysThread.msg.SyshandlerStack = uint64(sysmsg.StackAddrToSyshandlerStack(sysThread.sysmsgPerThreadMemAddr()))
@@ -1100,60 +1098,19 @@ func (s *subprocess) PostFork() {
 	s.usertrap.PostFork() // +checklocksforce: PreFork acquires, above.
 }
 
-// registerContext registers the context to an ID specific to this subprocess.
-// It will return an error if too many contexts are already active in this
-// subprocess.
-func (s *subprocess) registerContext(c *context) error {
-	s.mu.Lock()
-	c.mu.Lock()
-	// Unlock manually for the sake of not holding the lock while initializing
-	// context memory.
-	locked := true
-	unlock := func() {
-		if locked {
-			c.mu.Unlock()
-			s.mu.Unlock()
-			locked = false
+// activateContext activates the context in this subprocess.
+// No-op if the context is already active within the subprocess; if not,
+// deactivates it from its last subprocess.
+func (s *subprocess) activateContext(c *context) error {
+	if !c.sharedContext.isActiveInSubprocess(s) {
+		c.sharedContext.release()
+		c.sharedContext = nil
+
+		shared, err := s.getSharedContext()
+		if err != nil {
+			return err
 		}
+		c.sharedContext = shared
 	}
-	defer unlock()
-
-	if s == c.subprocess && c.cid != invalidContextID {
-		return nil
-	}
-
-	id, ok := s.threadContextPool.Get()
-	if !ok {
-		return fmt.Errorf("subprocess has too many active threads (%d); failed to create a new one", maxGuestContexts)
-	}
-	s.IncRef()
-	c.cid = id
-	c.subprocess = s
-	c.FullStateChanged()
-	unlock()
-
-	threadContext := s.getThreadContextFromID(id)
-	threadContext.Init(invalidThreadID)
 	return nil
 }
-
-// unregisterContext releases all references held for this context.
-//
-// Precondition: context c must have been active within subprocess s.
-func (s *subprocess) unregisterContext(c *context) {
-	if s == nil {
-		return
-	}
-	c.mu.Lock()
-	cid := c.cid
-	c.cid = invalidContextID
-	c.subprocess = nil
-	c.mu.Unlock()
-
-	s.mu.Lock()
-	delete(s.faultedContexts, c)
-	s.threadContextPool.Put(cid)
-	s.mu.Unlock()
-
-	s.DecRef(s.release)
-}

pkg/sentry/platform/systrap/subprocess_amd64_unsafe.go

@@ -1,65 +0,0 @@
-// Copyright 2018 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//go:build amd64
-// +build amd64
-
-package systrap
-
-import (
-	"unsafe"
-
-	"gvisor.dev/gvisor/pkg/sentry/arch"
-	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
-)
-
-//go:nosplit
-func isFPStateInContextRegion(ctx *sysmsg.ThreadContext) bool {
-	// If context decoupling experiment is ON then both the sighandler and
-	// syshandler save FPState to the context region since contexts will move
-	// threads. Otherwise only syshandler will save FPState to the region.
-	return contextDecouplingExp || ctx.State == sysmsg.ContextStateSyscallTrap
-}
-
-func (s *subprocess) saveFPState(msg *sysmsg.Msg, ctx *sysmsg.ThreadContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
-	fpState := ac.FloatingPointData().BytePointer()
-	dst := unsafeSlice(uintptr(unsafe.Pointer(fpState)), c.fpLen)
-	var src []byte
-	if isFPStateInContextRegion(ctx) {
-		src = ctx.FPState[:]
-	} else {
-		src = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), c.fpLen)
-	}
-	copy(dst, src)
-}
-
-// restoreFPStateDecoupledContext writes FPState from c to the thread context
-// shared memory region if there is any need to do so.
-func (s *subprocess) restoreFPState(msg *sysmsg.Msg, ctx *sysmsg.ThreadContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
-	if !c.needRestoreFPState {
-		return
-	}
-	c.needRestoreFPState = false
-	ctx.FPStateChanged = 1
-
-	fpState := ac.FloatingPointData().BytePointer()
-	src := unsafeSlice(uintptr(unsafe.Pointer(fpState)), c.fpLen)
-	var dst []byte
-	if isFPStateInContextRegion(ctx) {
-		dst = ctx.FPState[:]
-	} else {
-		dst = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), c.fpLen)
-	}
-	copy(dst, src)
-}

pkg/sentry/platform/systrap/subprocess_arm64_unsafe.go

@@ -1,65 +0,0 @@
-// Copyright 2019 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//go:build arm64
-// +build arm64
-
-package systrap
-
-import (
-	"unsafe"
-
-	"gvisor.dev/gvisor/pkg/sentry/arch"
-	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
-)
-
-//go:nosplit
-func isFPStateInContextRegion(ctx *sysmsg.ThreadContext) bool {
-	// If context decoupling experiment is ON then both the sighandler and
-	// syshandler save FPState to the context region since contexts will move
-	// threads. Otherwise only syshandler will save FPState to the region.
-	return contextDecouplingExp || ctx.State == sysmsg.ContextStateSyscallTrap
-}
-
-func (s *subprocess) restoreFPState(msg *sysmsg.Msg, ctx *sysmsg.ThreadContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
-	// c.needRestoreFPState is changed only from the task goroutine, so it can
-	// be accessed without locks.
-	if !c.needRestoreFPState {
-		return
-	}
-	c.needRestoreFPState = false
-	ctx.FPStateChanged = 1
-
-	fpState := ac.FloatingPointData().BytePointer()
-	src := unsafeSlice(uintptr(unsafe.Pointer(fpState)), c.fpLen)
-	var dst []byte
-	if isFPStateInContextRegion(ctx) {
-		dst = ctx.FPState[:]
-	} else {
-		dst = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), c.fpLen)
-	}
-	copy(dst, src)
-}
-
-func (s *subprocess) saveFPState(msg *sysmsg.Msg, ctx *sysmsg.ThreadContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
-	fpState := ac.FloatingPointData().BytePointer()
-	dst := unsafeSlice(uintptr(unsafe.Pointer(fpState)), c.fpLen)
-	var src []byte
-	if isFPStateInContextRegion(ctx) {
-		src = ctx.FPState[:]
-	} else {
-		src = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), c.fpLen)
-	}
-	copy(dst, src)
-}

pkg/sentry/platform/systrap/subprocess_unsafe.go

@@ -26,6 +26,7 @@ import (
 	"unsafe"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
@@ -68,3 +69,43 @@ func mmapContextQueueForSentry(memoryFile *pgalloc.MemoryFile, opts pgalloc.Allo
 
 	return fr, (*contextQueue)(unsafe.Pointer(addr))
 }
+
+//go:nosplit
+func isFPStateInContextRegion(ctx *sharedContext) bool {
+	// If context decoupling experiment is ON then both the sighandler and
+	// syshandler save FPState to the context region since contexts will move
+	// threads. Otherwise only syshandler will save FPState to the region.
+	return contextDecouplingExp || ctx.state() == sysmsg.ContextStateSyscallTrap
+}
+
+func saveFPState(msg *sysmsg.Msg, ctx *sharedContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
+	fpState := ac.FloatingPointData().BytePointer()
+	dst := unsafeSlice(uintptr(unsafe.Pointer(fpState)), archState.FpLen())
+	var src []byte
+	if isFPStateInContextRegion(ctx) {
+		src = ctx.shared.FPState[:]
+	} else {
+		src = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), archState.FpLen())
+	}
+	copy(dst, src)
+}
+
+// restoreFPStateDecoupledContext writes FPState from c to the thread context
+// shared memory region if there is any need to do so.
+func restoreFPState(msg *sysmsg.Msg, ctx *sharedContext, fpuToMsgOffset uint64, c *context, ac *arch.Context64) {
+	if !c.needRestoreFPState {
+		return
+	}
+	c.needRestoreFPState = false
+	ctx.setFPStateChanged()
+
+	fpState := ac.FloatingPointData().BytePointer()
+	src := unsafeSlice(uintptr(unsafe.Pointer(fpState)), archState.FpLen())
+	var dst []byte
+	if isFPStateInContextRegion(ctx) {
+		dst = ctx.shared.FPState[:]
+	} else {
+		dst = unsafeSlice(uintptr(unsafe.Pointer(msg))+uintptr(fpuToMsgOffset), archState.FpLen())
+	}
+	copy(dst, src)
+}

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

@@ -337,22 +337,6 @@ func (m *Msg) DisableSentryFastPath() {
 	atomic.StoreUint32(&m.sentryFastPath, 0)
 }
 
-// EnableSentryFastPath indicates that the polling mode is enabled for the
-// Sentry. It has to be called before putting the context into the context queue.
-// This function is used if contextDecouplingExp=true because the fastpath
-// is negotiated in ThreadContext
-func (c *ThreadContext) EnableSentryFastPath() {
-	c.SentryFastPath = 1
-}
-
-// DisableSentryFastPath indicates that the polling mode for the sentry is
-// disabled for the Sentry.
-// This function is used if contextDecouplingExp=true because the fastpath
-// is negotiated in ThreadContext.
-func (c *ThreadContext) DisableSentryFastPath() {
-	atomic.StoreUint32(&c.SentryFastPath, 0)
-}
-
 // FPUStateOffset returns the offset of a saved FPU state to the msg.
 func (m *Msg) FPUStateOffset() (uint64, error) {
 	offset := m.fpState

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

@@ -63,6 +63,11 @@ func (s *ArchState) Init() {
 	}
 }
 
+// FpLen returns the FP state length for AMD64.
+func (s *ArchState) FpLen() int {
+	return int(s.fpLen)
+}
+
 func (s *ArchState) String() string {
 	var b strings.Builder
 	fmt.Fprintf(&b, "sysmsg.ArchState{")

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

@@ -43,6 +43,11 @@ func (s *ArchState) Init() {
 	s.fpLen = uint32(fpLenUint)
 }
 
+// FpLen returns the FP state length for ARM.
+func (s *ArchState) FpLen() int {
+	return int(s.fpLen)
+}
+
 func (s *ArchState) String() string {
 	var b strings.Builder
 	fmt.Fprintf(&b, "sysmsg.ArchState{")

pkg/sentry/platform/systrap/systrap.go

@@ -52,18 +52,16 @@ import (
 	"fmt"
 	"os"
 	"sync"
-	"sync/atomic"
 
-	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	pkgcontext "gvisor.dev/gvisor/pkg/context"
-	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/memutil"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/pkg/sentry/platform/interrupt"
+	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/usertrap"
 )
 
@@ -100,6 +98,9 @@ var (
 
 	// stubInitialized controls one-time stub initialization.
 	stubInitialized sync.Once
+
+	// archState stores architecture-specific details used in the platform.
+	archState sysmsg.ArchState
 )
 
 // context is an implementation of the platform context.
@@ -110,16 +111,15 @@ type context struct {
 	// interrupt is the interrupt context.
 	interrupt interrupt.Forwarder
 
+	// sharedContext is everything related to this context that is resident in
+	// shared memory with the stub thread.
+	// sharedContext is only accessed on the Task goroutine, therefore it is not
+	// mutex protected.
+	sharedContext *sharedContext
+
 	// mu protects the following fields.
 	mu sync.Mutex
 
-	// subprocess is the current subprocess used to execute the context.
-	subprocess *subprocess
-
-	// cid is the ID of the context in the address space of the current
-	// subprocess used to run it.
-	cid uint64
-
 	// If lastFaultSP is non-nil, the last context switch was due to a fault
 	// received while executing lastFaultSP. Only context.Switch may set
 	// lastFaultSP to a non-nil value.
@@ -137,9 +137,6 @@ type context struct {
 	// application code. (Note: Unused if contextDecouplingExp=true).
 	sysmsgThread *sysmsgThread
 
-	// fpLen is the size of the floating point context.
-	fpLen int
-
 	// needRestoreFPState indicates that the FPU state has been changed by
 	// the Sentry and has to be updated on the stub thread.
 	needRestoreFPState bool
@@ -174,13 +171,10 @@ func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *
 
 	as := mm.AddressSpace()
 	s := as.(*subprocess)
-
-	if s != c.subprocess {
-		c.subprocess.unregisterContext(c)
-		if err := s.registerContext(c); err != nil {
-			return nil, hostarch.NoAccess, err
-		}
+	if err := s.activateContext(c); err != nil {
+		return nil, hostarch.NoAccess, err
 	}
+
 restart:
 	isSyscall, needPatch, err := s.switchToApp(c, ac)
 	if err != nil {
@@ -281,52 +275,15 @@ func (c *context) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
-// NotifyInterrupt implements interrupt.Receiver.NotifyInterrupt.
-//
-// Another reasonable existing object to implement NotifyInterrupt would be
-// sysmsg.ThreadContext, because we can write the correct host TID into it
-// to know which thread to send the signal to. However, because it is in shared
-// memory, one subprocess can overwrite it to have the sentry send an interrupt
-// to a completely different subprocess.
-// For this reason we use systrap.context and check that the target thread
-// is actually valid within the subprocess.
-func (c *context) NotifyInterrupt() {
-	c.mu.Lock()
-	s := c.subprocess
-	cid := c.cid
-	c.mu.Unlock()
-
-	if s == nil || cid == invalidContextID {
-		return
-	}
-
-	threadContext := s.getThreadContextFromID(cid)
-	atomic.StoreUint32(&threadContext.Interrupt, 1)
-	threadID := atomic.LoadUint32(&threadContext.ThreadID)
-
-	s.sysmsgThreadsMu.Lock()
-	defer s.sysmsgThreadsMu.Unlock()
-
-	sysmsgThread, ok := s.sysmsgThreads[threadID]
-	if !ok {
-		// This is either an invalidThreadID or another garbage value; either way we
-		// don't know which thread to interrupt; best we can do is mark the context.
-		return
-	}
-
-	t := sysmsgThread.thread
-	atomic.StoreUint64(&sysmsgThread.msg.InterruptedContextID, cid)
-	if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(t.tgid), uintptr(t.tid), uintptr(platform.SignalInterrupt)); e != 0 {
-		panic(fmt.Sprintf("failed to interrupt the child process %d: %v", t.tid, e))
-	}
-}
-
 // Release releases all platform resources used by the context.
 func (c *context) Release() {
 	if c.sysmsgThread != nil {
 		c.sysmsgThread.destroy()
 	}
-	c.subprocess.unregisterContext(c)
+	if c.sharedContext != nil {
+		c.sharedContext.release()
+		c.sharedContext = nil
+	}
 }
 
 // PrepareSleep implements platform.Context.platform.PrepareSleep.
@@ -356,6 +313,9 @@ func (*Systrap) MinUserAddress() hostarch.Addr {
 
 // New returns a new seccomp-based implementation of the platform interface.
 func New() (*Systrap, error) {
+	// CPUID information has been initialized at this point.
+	archState.Init()
+
 	mf, err := createMemoryFile()
 	if err != nil {
 		return nil, err
@@ -412,11 +372,7 @@ func (p *Systrap) NewAddressSpace(any) (platform.AddressSpace, <-chan struct{},
 
 // NewContext returns an interruptible context.
 func (*Systrap) NewContext(ctx pkgcontext.Context) platform.Context {
-	fs := cpuid.FromContext(ctx)
-	fpLen, _ := fs.ExtendedStateSize()
 	return &context{
-		cid:                 invalidContextID,
-		fpLen:               int(fpLen),
 		needRestoreFPState:  true,
 		needToPullFullState: false,
 	}