Plumb restore context to afterLoad()

This allows for external information to be passed to restore code, like
host FDs to be remapped.

Updates #1956

PiperOrigin-RevId: 612540749

pkg/cpuid/cpuid.go

@@ -38,7 +38,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
-// contextID is the package for context.Context.Value keys.
+// contextID is the package for anyContext.Context.Value keys.
 type contextID int
 
 const (
@@ -51,13 +51,13 @@ const (
 	_AT_HWCAP2 = 26
 )
 
-// context represents context.Context.
-type context interface {
+// anyContext represents context.Context.
+type anyContext interface {
 	Value(key any) any
 }
 
 // FromContext returns the FeatureSet from the context, if available.
-func FromContext(ctx context) FeatureSet {
+func FromContext(ctx anyContext) FeatureSet {
 	v := ctx.Value(CtxFeatureSet)
 	if v == nil {
 		return FeatureSet{} // Panics if used.

pkg/cpuid/static_amd64.go

@@ -17,6 +17,8 @@
 
 package cpuid
 
+import "context"
+
 // Static is a static CPUID function.
 //
 // +stateify savable
@@ -90,7 +92,7 @@ func (s Static) ToFeatureSet() FeatureSet {
 }
 
 // afterLoad calls normalize.
-func (s Static) afterLoad() {
+func (s Static) afterLoad(context.Context) {
 	s.normalize()
 }
 

pkg/refs/refs_template.go

@@ -17,6 +17,7 @@
 package refs_template
 
 import (
+	"context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/atomicbitops"
@@ -156,7 +157,7 @@ func (r *Refs) DecRef(destroy func()) {
 	}
 }
 
-func (r *Refs) afterLoad() {
+func (r *Refs) afterLoad(context.Context) {
 	if r.ReadRefs() > 0 {
 		refs.Register(r)
 	}

pkg/sentry/arch/arch_x86_impl.go

@@ -19,6 +19,8 @@
 package arch
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/sentry/arch/fpu"
 )
 
@@ -35,6 +37,6 @@ type State struct {
 }
 
 // afterLoad is invoked by stateify.
-func (s *State) afterLoad() {
+func (s *State) afterLoad(context.Context) {
 	s.afterLoadFPState()
 }

pkg/sentry/devices/nvproxy/save_restore.go

@@ -15,6 +15,7 @@
 package nvproxy
 
 import (
+	goContext "context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/abi/nvgpu"
@@ -30,7 +31,7 @@ func (n *nvproxy) beforeSave() {
 	n.objsLive = nil
 }
 
-func (n *nvproxy) afterLoad() {
+func (n *nvproxy) afterLoad(goContext.Context) {
 	Init()
 	abiCons, ok := abis[n.version]
 	if !ok {

pkg/sentry/fsimpl/devtmpfs/save_restore.go

@@ -14,8 +14,10 @@
 
 package devtmpfs
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (fst *FilesystemType) afterLoad() {
+func (fst *FilesystemType) afterLoad(context.Context) {
 	if fst.fs != nil {
 		// Ensure that we don't create another filesystem.
 		fst.initOnce.Do(func() {})

pkg/sentry/fsimpl/fuse/save_restore.go

@@ -14,6 +14,8 @@
 
 package fuse
 
-func (fRes *futureResponse) afterLoad() {
+import "context"
+
+func (fRes *futureResponse) afterLoad(context.Context) {
 	fRes.ch = make(chan struct{})
 }

pkg/sentry/fsimpl/gofer/save_restore.go

@@ -15,6 +15,7 @@
 package gofer
 
 import (
+	goContext "context"
 	"fmt"
 	"io"
 
@@ -128,7 +129,7 @@ func (d *dentry) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (d *dentry) afterLoad() {
+func (d *dentry) afterLoad(goContext.Context) {
 	d.readFD = atomicbitops.FromInt32(-1)
 	d.writeFD = atomicbitops.FromInt32(-1)
 	d.mmapFD = atomicbitops.FromInt32(-1)
@@ -138,12 +139,12 @@ func (d *dentry) afterLoad() {
 }
 
 // afterLoad is invoked by stateify.
-func (d *directfsDentry) afterLoad() {
+func (d *directfsDentry) afterLoad(goContext.Context) {
 	d.controlFD = -1
 }
 
 // afterLoad is invoked by stateify.
-func (d *dentryPlatformFile) afterLoad() {
+func (d *dentryPlatformFile) afterLoad(goContext.Context) {
 	if d.hostFileMapper.IsInited() {
 		// Ensure that we don't call d.hostFileMapper.Init() again.
 		d.hostFileMapperInitOnce.Do(func() {})
@@ -151,7 +152,7 @@ func (d *dentryPlatformFile) afterLoad() {
 }
 
 // afterLoad is invoked by stateify.
-func (fd *specialFileFD) afterLoad() {
+func (fd *specialFileFD) afterLoad(goContext.Context) {
 	fd.handle.fd = -1
 	if fd.hostFileMapper.IsInited() {
 		// Ensure that we don't call fd.hostFileMapper.Init() again.

pkg/sentry/fsimpl/host/save_restore.go

@@ -15,6 +15,7 @@
 package host
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -57,7 +58,7 @@ func (i *inode) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (i *inode) afterLoad() {
+func (i *inode) afterLoad(context.Context) {
 	if i.epollable {
 		if err := unix.SetNonblock(i.hostFD, true); err != nil {
 			panic(fmt.Sprintf("host.inode.afterLoad: failed to set host FD %d non-blocking: %v", i.hostFD, err))

pkg/sentry/fsimpl/iouringfs/iouringfs_state.go

@@ -14,6 +14,8 @@
 
 package iouringfs
 
+import "context"
+
 // beforeSave is invoked by stateify.
 func (fd *FileDescription) beforeSave() {
 	if fd.running.Load() != 0 {
@@ -22,7 +24,7 @@ func (fd *FileDescription) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (fd *FileDescription) afterLoad() {
+func (fd *FileDescription) afterLoad(context.Context) {
 	// Remap shared buffers.
 	fd.remap = true
 	fd.runC = make(chan struct{}, 1)

pkg/sentry/fsimpl/kernfs/save_restore.go

@@ -15,18 +15,20 @@
 package kernfs
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/refs"
 )
 
 // afterLoad is invoked by stateify.
-func (d *Dentry) afterLoad() {
+func (d *Dentry) afterLoad(context.Context) {
 	if d.refs.Load() >= 0 {
 		refs.Register(d)
 	}
 }
 
 // afterLoad is invoked by stateify.
-func (i *inodePlatformFile) afterLoad() {
+func (i *inodePlatformFile) afterLoad(context.Context) {
 	if i.fileMapper.IsInited() {
 		// Ensure that we don't call i.fileMapper.Init() again.
 		i.fileMapperInitOnce.Do(func() {})

pkg/sentry/fsimpl/overlay/save_restore.go

@@ -15,10 +15,12 @@
 package overlay
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/refs"
 )
 
-func (d *dentry) afterLoad() {
+func (d *dentry) afterLoad(context.Context) {
 	if d.refs.Load() != -1 {
 		refs.Register(d)
 	}

pkg/sentry/fsimpl/tmpfs/save_restore.go

@@ -15,6 +15,7 @@
 package tmpfs
 
 import (
+	goContext "context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/context"
@@ -23,7 +24,7 @@ import (
 )
 
 // afterLoad is called by stateify.
-func (fs *filesystem) afterLoad() {
+func (fs *filesystem) afterLoad(goContext.Context) {
 	if !fs.privateMF {
 		fs.mf = fs.mfp.MemoryFile()
 	}

pkg/sentry/fsutil/host_file_mapper_state.go

@@ -14,7 +14,9 @@
 
 package fsutil
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (f *HostFileMapper) afterLoad() {
+func (f *HostFileMapper) afterLoad(context.Context) {
 	f.mappings = make(map[uint64]mapping)
 }

pkg/sentry/inet/namespace.go

@@ -15,6 +15,8 @@
 package inet
 
 import (
+	goContext "context"
+
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/nsfs"
 	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
@@ -151,7 +153,7 @@ func (n *Namespace) init() {
 }
 
 // afterLoad is invoked by stateify.
-func (n *Namespace) afterLoad() {
+func (n *Namespace) afterLoad(goContext.Context) {
 	n.init()
 }
 

pkg/sentry/kernel/pipe/save_restore.go

@@ -15,11 +15,13 @@
 package pipe
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/safemem"
 )
 
 // afterLoad is called by stateify.
-func (p *Pipe) afterLoad() {
+func (p *Pipe) afterLoad(context.Context) {
 	p.bufBlocks[0] = safemem.BlockFromSafeSlice(p.buf)
 	p.bufBlocks[1] = p.bufBlocks[0]
 	p.bufBlockSeq = safemem.BlockSeqFromSlice(p.bufBlocks[:])

pkg/sentry/kernel/task.go

@@ -630,7 +630,7 @@ func (t *Task) loadSeccomp(seccompData *taskSeccomp) {
 }
 
 // afterLoad is invoked by stateify.
-func (t *Task) afterLoad() {
+func (t *Task) afterLoad(gocontext.Context) {
 	t.updateInfoLocked()
 	if ts := t.seccomp.Load().(*taskSeccomp); ts != nil {
 		ts.populateCache(t)

pkg/sentry/kernel/timekeeper_state.go

@@ -15,6 +15,8 @@
 package kernel
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/sentry/time"
 )
 
@@ -36,6 +38,6 @@ func (t *Timekeeper) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (t *Timekeeper) afterLoad() {
+func (t *Timekeeper) afterLoad(context.Context) {
 	t.restored = make(chan struct{})
 }

pkg/sentry/mm/aio_context.go

@@ -126,52 +126,52 @@ type AIOContext struct {
 }
 
 // destroy marks the context dead.
-func (ctx *AIOContext) destroy() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	ctx.dead = true
-	ctx.checkForDone()
+func (aio *AIOContext) destroy() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	aio.dead = true
+	aio.checkForDone()
 }
 
 // Preconditions: ctx.mu must be held by caller.
-func (ctx *AIOContext) checkForDone() {
-	if ctx.dead && ctx.outstanding == 0 {
-		close(ctx.requestReady)
-		ctx.requestReady = nil
+func (aio *AIOContext) checkForDone() {
+	if aio.dead && aio.outstanding == 0 {
+		close(aio.requestReady)
+		aio.requestReady = nil
 	}
 }
 
 // Prepare reserves space for a new request, returning nil if available.
 // Returns EAGAIN if the context is busy and EINVAL if the context is dead.
-func (ctx *AIOContext) Prepare() error {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	if ctx.dead {
+func (aio *AIOContext) Prepare() error {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	if aio.dead {
 		// Context died after the caller looked it up.
 		return linuxerr.EINVAL
 	}
-	if ctx.outstanding >= ctx.maxOutstanding {
+	if aio.outstanding >= aio.maxOutstanding {
 		// Context is busy.
 		return linuxerr.EAGAIN
 	}
-	ctx.outstanding++
+	aio.outstanding++
 	return nil
 }
 
 // PopRequest pops a completed request if available, this function does not do
 // any blocking. Returns false if no request is available.
-func (ctx *AIOContext) PopRequest() (any, bool) {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) PopRequest() (any, bool) {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
 	// Is there anything ready?
-	if e := ctx.results.Front(); e != nil {
-		if ctx.outstanding == 0 {
+	if e := aio.results.Front(); e != nil {
+		if aio.outstanding == 0 {
 			panic("AIOContext outstanding is going negative")
 		}
-		ctx.outstanding--
-		ctx.results.Remove(e)
-		ctx.checkForDone()
+		aio.outstanding--
+		aio.results.Remove(e)
+		aio.checkForDone()
 		return e.data, true
 	}
 	return nil, false
@@ -179,17 +179,17 @@ func (ctx *AIOContext) PopRequest() (any, bool) {
 
 // FinishRequest finishes a pending request. It queues up the data
 // and notifies listeners.
-func (ctx *AIOContext) FinishRequest(data any) {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) FinishRequest(data any) {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
 	// Push to the list and notify opportunistically. The channel notify
 	// here is guaranteed to be safe because outstanding must be non-zero.
 	// The requestReady channel is only closed when outstanding reaches zero.
-	ctx.results.PushBack(&ioResult{data: data})
+	aio.results.PushBack(&ioResult{data: data})
 
 	select {
-	case ctx.requestReady <- struct{}{}:
+	case aio.requestReady <- struct{}{}:
 	default:
 	}
 }
@@ -197,46 +197,46 @@ func (ctx *AIOContext) FinishRequest(data any) {
 // WaitChannel returns a channel that is notified when an AIO request is
 // completed. Returns nil if the context is destroyed and there are no more
 // outstanding requests.
-func (ctx *AIOContext) WaitChannel() chan struct{} {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	return ctx.requestReady
+func (aio *AIOContext) WaitChannel() chan struct{} {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	return aio.requestReady
 }
 
 // Dead returns true if the context has been destroyed.
-func (ctx *AIOContext) Dead() bool {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	return ctx.dead
+func (aio *AIOContext) Dead() bool {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	return aio.dead
 }
 
 // CancelPendingRequest forgets about a request that hasn't yet completed.
-func (ctx *AIOContext) CancelPendingRequest() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) CancelPendingRequest() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
-	if ctx.outstanding == 0 {
+	if aio.outstanding == 0 {
 		panic("AIOContext outstanding is going negative")
 	}
-	ctx.outstanding--
-	ctx.checkForDone()
+	aio.outstanding--
+	aio.checkForDone()
 }
 
 // Drain drops all completed requests. Pending requests remain untouched.
-func (ctx *AIOContext) Drain() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) Drain() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
-	if ctx.outstanding == 0 {
+	if aio.outstanding == 0 {
 		return
 	}
-	size := uint32(ctx.results.Len())
-	if ctx.outstanding < size {
+	size := uint32(aio.results.Len())
+	if aio.outstanding < size {
 		panic("AIOContext outstanding is going negative")
 	}
-	ctx.outstanding -= size
-	ctx.results.Reset()
-	ctx.checkForDone()
+	aio.outstanding -= size
+	aio.results.Reset()
+	aio.checkForDone()
 }
 
 // aioMappable implements memmap.MappingIdentity and memmap.Mappable for AIO

pkg/sentry/mm/aio_context_state.go

@@ -14,7 +14,9 @@
 
 package mm
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (ctx *AIOContext) afterLoad() {
-	ctx.requestReady = make(chan struct{}, 1)
+func (aio *AIOContext) afterLoad(context.Context) {
+	aio.requestReady = make(chan struct{}, 1)
 }