Plumb restore context to load*() methods.

This allows for external information to be passed to restore code.
Similar to c087777e37 ("Plumb restore context to afterLoad()").

Updates #1956.

PiperOrigin-RevId: 614125262

pkg/buffer/buffer_state.go

@@ -14,12 +14,16 @@
 
 package buffer
 
+import (
+	"context"
+)
+
 // saveData is invoked by stateify.
 func (b *Buffer) saveData() []byte {
 	return b.Flatten()
 }
 
 // loadData is invoked by stateify.
-func (b *Buffer) loadData(data []byte) {
+func (b *Buffer) loadData(_ context.Context, data []byte) {
 	*b = MakeWithData(data)
 }

pkg/cpuid/cpuid_amd64.go

@@ -18,6 +18,7 @@
 package cpuid
 
 import (
+	"context"
 	"fmt"
 	"io"
 )
@@ -56,7 +57,7 @@ func (fs *FeatureSet) saveFunction() Static {
 }
 
 // loadFunction saves the function as a static query.
-func (fs *FeatureSet) loadFunction(s Static) {
+func (fs *FeatureSet) loadFunction(_ context.Context, s Static) {
 	fs.Function = s
 }
 

pkg/segment/set_state.go

@@ -14,6 +14,10 @@
 
 package segment
 
+import (
+	"context"
+)
+
 func (s *Set) saveRoot() []FlatSegment {
 	fs := s.ExportSlice()
 	// The state package saves data in slice capacity beyond slice length; save
@@ -22,7 +26,7 @@ func (s *Set) saveRoot() []FlatSegment {
 	return fs
 }
 
-func (s *Set) loadRoot(fs []FlatSegment) {
+func (s *Set) loadRoot(_ context.Context, fs []FlatSegment) {
 	if err := s.ImportSlice(fs); err != nil {
 		panic(err)
 	}

pkg/sentry/fsimpl/erofs/save_restore.go

@@ -15,6 +15,7 @@
 package erofs
 
 import (
+	goContext "context"
 	"fmt"
 	"os"
 
@@ -62,6 +63,6 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }

pkg/sentry/fsimpl/fuse/connection.go

@@ -15,6 +15,7 @@
 package fuse
 
 import (
+	goContext "context"
 	"sync"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -193,7 +194,7 @@ func (conn *connection) saveInitializedChan() bool {
 	}
 }
 
-func (conn *connection) loadInitializedChan(closed bool) {
+func (conn *connection) loadInitializedChan(_ goContext.Context, closed bool) {
 	conn.initializedChan = make(chan struct{}, 1)
 	if closed {
 		close(conn.initializedChan)

pkg/sentry/fsimpl/fuse/dev_state.go

@@ -14,10 +14,14 @@
 
 package fuse
 
+import (
+	"context"
+)
+
 func (fd *DeviceFD) saveFullQueueCh() int {
 	return cap(fd.fullQueueCh)
 }
 
-func (fd *DeviceFD) loadFullQueueCh(capacity int) {
+func (fd *DeviceFD) loadFullQueueCh(_ context.Context, capacity int) {
 	fd.fullQueueCh = make(chan struct{}, capacity)
 }

pkg/sentry/fsimpl/gofer/save_restore.go

@@ -166,7 +166,7 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
 

pkg/sentry/fsimpl/kernfs/save_restore.go

@@ -41,6 +41,6 @@ func (d *Dentry) saveParent() *Dentry {
 }
 
 // loadParent is called by stateify.
-func (d *Dentry) loadParent(parent *Dentry) {
+func (d *Dentry) loadParent(_ context.Context, parent *Dentry) {
 	d.parent.Store(parent)
 }

pkg/sentry/fsimpl/overlay/save_restore.go

@@ -32,6 +32,6 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ context.Context, parent *dentry) {
 	d.parent.Store(parent)
 }

pkg/sentry/fsimpl/tmpfs/save_restore.go

@@ -36,7 +36,7 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
 

pkg/sentry/kernel/fd_table.go

@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	goContext "context"
 	"fmt"
 	"math"
 	"strings"
@@ -101,7 +102,7 @@ func (f *FDTable) saveDescriptorTable() map[int32]descriptor {
 	return m
 }
 
-func (f *FDTable) loadDescriptorTable(m map[int32]descriptor) {
+func (f *FDTable) loadDescriptorTable(_ goContext.Context, m map[int32]descriptor) {
 	ctx := context.Background()
 	f.initNoLeakCheck() // Initialize table.
 	f.fdBitmap = bitmap.New(uint32(math.MaxUint16))

pkg/sentry/kernel/kernel_state.go

@@ -15,6 +15,8 @@
 package kernel
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
@@ -24,7 +26,7 @@ func (k *Kernel) saveDanglingEndpoints() []tcpip.Endpoint {
 }
 
 // loadDanglingEndpoints is invoked by stateify.
-func (k *Kernel) loadDanglingEndpoints(es []tcpip.Endpoint) {
+func (k *Kernel) loadDanglingEndpoints(_ context.Context, es []tcpip.Endpoint) {
 	for _, e := range es {
 		tcpip.AddDanglingEndpoint(e)
 	}

pkg/sentry/kernel/pending_signals_state.go

@@ -14,7 +14,11 @@
 
 package kernel
 
-import "gvisor.dev/gvisor/pkg/abi/linux"
+import (
+	"context"
+
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
 
 // +stateify savable
 type savedPendingSignal struct {
@@ -37,7 +41,7 @@ func (p *pendingSignals) saveSignals() []savedPendingSignal {
 }
 
 // loadSignals is invoked by stateify.
-func (p *pendingSignals) loadSignals(pending []savedPendingSignal) {
+func (p *pendingSignals) loadSignals(_ context.Context, pending []savedPendingSignal) {
 	for _, sps := range pending {
 		p.enqueue(sps.si, sps.timer)
 	}

pkg/sentry/kernel/syscalls_state.go

@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	"context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/abi"
@@ -38,7 +39,7 @@ func (image *TaskImage) saveSt() syscallTableInfo {
 }
 
 // loadSt loads the SyscallTable.
-func (image *TaskImage) loadSt(sti syscallTableInfo) {
+func (image *TaskImage) loadSt(_ context.Context, sti syscallTableInfo) {
 	st, ok := LookupSyscallTable(sti.OS, sti.Arch)
 	if !ok {
 		panic(fmt.Sprintf("syscall table not found for OS %v, Arch %v", sti.OS, sti.Arch))

pkg/sentry/kernel/task.go

@@ -615,7 +615,7 @@ func (t *Task) savePtraceTracer() *Task {
 	return t.ptraceTracer.Load()
 }
 
-func (t *Task) loadPtraceTracer(tracer *Task) {
+func (t *Task) loadPtraceTracer(_ gocontext.Context, tracer *Task) {
 	t.ptraceTracer.Store(tracer)
 }
 
@@ -623,7 +623,7 @@ func (t *Task) saveSeccomp() *taskSeccomp {
 	return t.seccomp.Load()
 }
 
-func (t *Task) loadSeccomp(seccompData *taskSeccomp) {
+func (t *Task) loadSeccomp(_ gocontext.Context, seccompData *taskSeccomp) {
 	t.seccomp.Store(seccompData)
 }
 

pkg/sentry/kernel/thread_group.go

@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	goContext "context"
 	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -293,7 +294,7 @@ func (tg *ThreadGroup) saveOldRSeqCritical() *OldRSeqCriticalRegion {
 }
 
 // loadOldRSeqCritical is invoked by stateify.
-func (tg *ThreadGroup) loadOldRSeqCritical(r *OldRSeqCriticalRegion) {
+func (tg *ThreadGroup) loadOldRSeqCritical(_ goContext.Context, r *OldRSeqCriticalRegion) {
 	tg.oldRSeqCritical.Store(r)
 }
 

pkg/sentry/loader/vdso_state.go

@@ -15,6 +15,7 @@
 package loader
 
 import (
+	"context"
 	"debug/elf"
 )
 
@@ -40,7 +41,7 @@ func (v *VDSO) savePhdrs() []elfProgHeader {
 }
 
 // loadPhdrs is invoked by stateify.
-func (v *VDSO) loadPhdrs(s []elfProgHeader) {
+func (v *VDSO) loadPhdrs(_ context.Context, s []elfProgHeader) {
 	v.phdrs = make([]elf.ProgHeader, 0, len(s))
 	for _, h := range s {
 		v.phdrs = append(v.phdrs, elf.ProgHeader(h))

pkg/sentry/mm/save_restore.go

@@ -113,7 +113,7 @@ func (v *vma) saveRealPerms() int {
 	return b
 }
 
-func (v *vma) loadRealPerms(b int) {
+func (v *vma) loadRealPerms(_ goContext.Context, b int) {
 	if b&vmaRealPermsRead > 0 {
 		v.realPerms.Read = true
 	}

pkg/sentry/socket/netstack/netstack_state.go

@@ -15,6 +15,7 @@
 package netstack
 
 import (
+	"context"
 	"time"
 )
 
@@ -24,7 +25,7 @@ func (s *sock) saveTimestamp() int64 {
 	return s.timestamp.UnixNano()
 }
 
-func (s *sock) loadTimestamp(nsec int64) {
+func (s *sock) loadTimestamp(_ context.Context, nsec int64) {
 	s.readMu.Lock()
 	defer s.readMu.Unlock()
 	s.timestamp = time.Unix(0, nsec)

pkg/sentry/socket/socket_state.go

@@ -15,6 +15,7 @@
 package socket
 
 import (
+	"context"
 	"time"
 )
 
@@ -22,6 +23,6 @@ func (i *IPControlMessages) saveTimestamp() int64 {
 	return i.Timestamp.UnixNano()
 }
 
-func (i *IPControlMessages) loadTimestamp(nsec int64) {
+func (i *IPControlMessages) loadTimestamp(_ context.Context, nsec int64) {
 	i.Timestamp = time.Unix(0, nsec)
 }