Internal change

PiperOrigin-RevId: 631677570

pkg/abi/nvgpu/frontend.go

@@ -369,6 +369,7 @@ type NVOS56Parameters struct {
 // NV_ESC_RM_ALLOC.
 //
 // +marshal
+// +stateify savable
 type NVOS64Parameters struct {
 	HRoot            Handle
 	HObjectParent    Handle

pkg/abi/nvgpu/nvgpu.go

@@ -30,6 +30,7 @@ const (
 // Handle is NvHandle, from src/common/sdk/nvidia/inc/nvtypes.h.
 //
 // +marshal
+// +stateify savable
 type Handle struct {
 	Val uint32
 }
@@ -60,6 +61,7 @@ const (
 // src/common/sdk/nvidia/inc/rs_access.h.
 //
 // +marshal
+// +stateify savable
 type RS_ACCESS_MASK struct {
 	Limbs [SDK_RS_ACCESS_MAX_LIMBS]uint32 // RsAccessLimb
 }

pkg/devutil/devutil.go

@@ -52,6 +52,7 @@ func NewGoferClient(ctx context.Context, contName string, fd int) (*GoferClient,
 	return &GoferClient{
 		clientFD: client.NewFD(devInode.ControlFD),
 		hostFD:   devHostFD,
+		contName: contName,
 	}, nil
 }
 

pkg/sentry/devices/nvproxy/BUILD

@@ -6,6 +6,13 @@ package(default_applicable_licenses = ["//:license"])
 
 licenses(["notice"])
 
+declare_mutex(
+    name = "fds_mutex",
+    out = "fds_mutex.go",
+    package = "nvproxy",
+    prefix = "fds",
+)
+
 declare_mutex(
     name = "objs_mutex",
     out = "objs_mutex.go",
@@ -28,6 +35,7 @@ go_template_instance(
 go_library(
     name = "nvproxy",
     srcs = [
+        "fds_mutex.go",
         "frontend.go",
         "frontend_mmap.go",
         "frontend_unsafe.go",

pkg/sentry/devices/nvproxy/frontend.go

@@ -81,6 +81,9 @@ func (dev *frontendDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.D
 		return nil, err
 	}
 	fd.memmapFile.fd = fd
+	fd.dev.nvp.fdsMu.Lock()
+	defer fd.dev.nvp.fdsMu.Unlock()
+	fd.dev.nvp.frontendFDs[fd] = struct{}{}
 	return &fd.vfsfd, nil
 }
 
@@ -103,8 +106,8 @@ type frontendFD struct {
 	haveMmapContext atomic.Bool `state:"nosave"`
 
 	// clients are handles of clients owned by this frontendFD. clients is
-	// protected by nvp.objsMu.
-	clients map[nvgpu.Handle]struct{} `state:"nosave"`
+	// protected by dev.nvp.objsMu.
+	clients map[nvgpu.Handle]struct{}
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
@@ -112,6 +115,10 @@ func (fd *frontendFD) Release(ctx context.Context) {
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
 
+	fd.dev.nvp.fdsMu.Lock()
+	delete(fd.dev.nvp.frontendFDs, fd)
+	fd.dev.nvp.fdsMu.Unlock()
+
 	fd.dev.nvp.objsLock()
 	defer fd.dev.nvp.objsUnlock()
 	unix.Close(int(fd.hostFD))
@@ -735,11 +742,11 @@ func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlStat
 
 // addSimpleObjDepParentLocked implements rmAllocInvoke.addObjLocked for
 // classes that require no special handling and depend only on their parents.
-func addSimpleObjDepParentLocked[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params) {
-	fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent)
+func addSimpleObjDepParentLocked[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *Params) {
+	fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRmAllocObject(fi.fd, ioctlParams, rightsRequested, allocParams), ioctlParams.HObjectParent)
 }
 
-func rmAllocSimpleParams[Params any, PParams marshalPtr[Params]](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool, objAddLocked func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params)) (uintptr, error) {
+func rmAllocSimpleParams[Params any, PParams marshalPtr[Params]](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool, objAddLocked func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *Params)) (uintptr, error) {
 	if ioctlParams.PAllocParms == 0 {
 		return rmAllocInvoke[Params](fi, ioctlParams, nil, isNVOS64, objAddLocked)
 	}
@@ -763,8 +770,8 @@ func rmAllocNoParams(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters
 }
 
 func rmAllocRootClient(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *nvgpu.Handle) {
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRootClient(fi.fd))
+	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.Handle) {
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRootClient(fi.fd, ioctlParams, rightsRequested, allocParams))
 		if fi.fd.clients == nil {
 			fi.fd.clients = make(map[nvgpu.Handle]struct{})
 		}
@@ -789,7 +796,9 @@ func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parame
 	origData := allocParams.Data
 	allocParams.Data = nvgpu.P64(uint64(eventFile.hostFD))
 
-	n, err := rmAllocInvoke(fi, ioctlParams, &allocParams, isNVOS64, addSimpleObjDepParentLocked)
+	n, err := rmAllocInvoke(fi, ioctlParams, &allocParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.NV0005_ALLOC_PARAMETERS) {
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, &osEvent{}, ioctlParams.HObjectParent)
+	})
 	if err != nil {
 		return n, err
 	}
@@ -802,22 +811,22 @@ func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parame
 }
 
 func rmAllocSMDebuggerSession(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *nvgpu.NV83DE_ALLOC_PARAMETERS) {
+	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.NV83DE_ALLOC_PARAMETERS) {
 		// Compare
 		// src/nvidia/src/kernel/gpu/gr/kernel_sm_debugger_session.c:ksmdbgssnConstruct_IMPL()
 		// => _ShareDebugger() => sessionAddDependency/sessionAddDependant();
 		// the driver indirects through a per-KernelGraphicsObject
 		// RmDebuggerSession, which we elide for dependency tracking.
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent, allocParams.HClass3DObject)
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRmAllocObject(fi.fd, ioctlParams, rightsRequested, allocParams), ioctlParams.HObjectParent, allocParams.HClass3DObject)
 	})
 }
 
 func rmAllocChannelGroup(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *nvgpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS) {
+	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS) {
 		// See
 		// src/nvidia/src/kernel/gpu/fifo/kernel_channel_group_api.c:kchangrpapiConstruct_IMPL()
 		// => refAddDependant().
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent, allocParams.HVASpace)
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRmAllocObject(fi.fd, ioctlParams, rightsRequested, allocParams), ioctlParams.HObjectParent, allocParams.HVASpace)
 		// Note: When the channel group's engine type is GR, which is always
 		// true unless MIG is enabled, kchangrpapiConstruct_IMPL() constructs a
 		// KERNEL_GRAPHICS_CONTEXT whose lifetime is the same as the channel
@@ -831,7 +840,7 @@ func rmAllocChannelGroup(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parame
 }
 
 func rmAllocChannel(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *nvgpu.NV_CHANNEL_ALLOC_PARAMS) {
+	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.NV_CHANNEL_ALLOC_PARAMS) {
 		// See
 		// src/nvidia/src/kernel/gpu/fifo/kernel_channel.c:kchannelConstruct_IMPL()
 		// => refAddDependant(). The channel's parent may be a device or
@@ -840,19 +849,19 @@ func rmAllocChannel(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters,
 		// then kchannelConstruct_IMPL() constructs one internally and frees it
 		// when the channel is destroyed, so either way no separate dependency
 		// is required.
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent, allocParams.HVASpace, allocParams.HContextShare)
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRmAllocObject(fi.fd, ioctlParams, rightsRequested, allocParams), ioctlParams.HObjectParent, allocParams.HVASpace, allocParams.HContextShare)
 	})
 }
 
 func rmAllocContextShare(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *nvgpu.NV_CTXSHARE_ALLOCATION_PARAMETERS) {
+	return rmAllocSimpleParams(fi, ioctlParams, isNVOS64, func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.NV_CTXSHARE_ALLOCATION_PARAMETERS) {
 		// See
 		// src/nvidia/src/kernel/gpu/fifo/kernel_ctxshare.c:kctxshareapiConstruct_IMPL()
 		// => refAddDependant(). The context share's parent is the channel
 		// group, so (given that we are representing graphics context
 		// dependencies as channel group dependencies) no separate dependency
 		// is required.
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent, allocParams.HVASpace)
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRmAllocObject(fi.fd, ioctlParams, rightsRequested, allocParams), ioctlParams.HObjectParent, allocParams.HVASpace)
 	})
 }
 

pkg/sentry/devices/nvproxy/frontend_unsafe.go

@@ -169,7 +169,7 @@ func ctrlSubdevGRGetInfo(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parame
 	return n, nil
 }
 
-func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params, isNVOS64 bool, addObjLocked func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params)) (uintptr, error) {
+func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params, isNVOS64 bool, addObjLocked func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *Params)) (uintptr, error) {
 	defer runtime.KeepAlive(allocParams) // since we convert to non-pointer-typed P64
 
 	// Temporarily replace application pointers with sentry pointers.
@@ -192,7 +192,7 @@ func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64
 	fi.fd.dev.nvp.objsLock()
 	n, _, errno := unix.RawSyscall(unix.SYS_IOCTL, uintptr(fi.fd.hostFD), frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS64Parameters), uintptr(unsafe.Pointer(ioctlParams)))
 	if errno == 0 && ioctlParams.Status == nvgpu.NV_OK {
-		addObjLocked(fi, ioctlParams, allocParams)
+		addObjLocked(fi, ioctlParams, rightsRequested, allocParams)
 	}
 	fi.fd.dev.nvp.objsUnlock()
 	ioctlParams.PAllocParms = origPAllocParms
@@ -230,7 +230,7 @@ func rmVidHeapControlAllocSize(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS32
 	n, err := frontendIoctlInvoke(fi, ioctlParams)
 	if err == nil && ioctlParams.Status == nvgpu.NV_OK {
 		// src/nvidia/src/kernel/mem_mgr/virtual_mem.c:virtmemConstruct_IMPL() => refAddDependant()
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, nvgpu.NV50_MEMORY_VIRTUAL, newSimpleObject(), ioctlParams.HObjectParent, ioctlParams.HVASpace)
+		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, nvgpu.NV50_MEMORY_VIRTUAL, &virtMem{}, ioctlParams.HObjectParent, ioctlParams.HVASpace)
 	}
 	fi.fd.dev.nvp.objsUnlock()
 	allocSizeParams.Address = origAddress

pkg/sentry/devices/nvproxy/nvproxy.go

@@ -44,6 +44,7 @@ func Register(vfsObj *vfs.VirtualFilesystem, versionStr string, uvmDevMajor uint
 	nvp := &nvproxy{
 		abi:         abiCons.cons(),
 		version:     version,
+		frontendFDs: make(map[*frontendFD]struct{}),
 		clients:     make(map[nvgpu.Handle]*rootClient),
 		objsFreeSet: make(map[*object]struct{}),
 	}
@@ -72,6 +73,9 @@ type nvproxy struct {
 	abi     *driverABI `state:"nosave"`
 	version DriverVersion
 
+	fdsMu       fdsMutex `state:"nosave"`
+	frontendFDs map[*frontendFD]struct{}
+
 	// See object.go.
 	// Users should call nvproxy.objsLock/Unlock() rather than locking objsMu
 	// directly.

pkg/sentry/devices/nvproxy/object.go

@@ -18,6 +18,7 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/nvgpu"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/marshal"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
 )
 
@@ -230,42 +231,76 @@ func (nvp *nvproxy) enqueueCleanup(f func()) {
 	nvp.objsCleanup = append(nvp.objsCleanup, f)
 }
 
-// simpleObject is an objectImpl tracking a driver object whose class is not
-// represented by a more specific type.
-type simpleObject struct {
-	object
+// +stateify savable
+type capturedRmAllocParams struct {
+	fd              *frontendFD
+	ioctlParams     nvgpu.NVOS64Parameters
+	rightsRequested nvgpu.RS_ACCESS_MASK
+	allocParams     []byte
 }
 
-func newSimpleObject() *simpleObject {
-	return &simpleObject{}
+func captureRmAllocParams[Params any](fd *frontendFD, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *Params) capturedRmAllocParams {
+	var allocParamsBuf []byte
+	if allocParams != nil {
+		if allocParamsMarshal, ok := any(allocParams).(marshal.Marshallable); ok {
+			allocParamsBuf = make([]byte, allocParamsMarshal.SizeBytes())
+			allocParamsMarshal.MarshalBytes(allocParamsBuf)
+		} else {
+			log.Traceback("nvproxy: allocParams %T is not marshalable")
+		}
+	}
+	return capturedRmAllocParams{
+		fd:              fd,
+		ioctlParams:     *ioctlParams,
+		rightsRequested: rightsRequested,
+		allocParams:     allocParamsBuf,
+	}
+}
+
+// rmAllocObject is an objectImpl tracking a driver object allocated by an
+// invocation of NV_ESC_RM_ALLOC whose class is not represented by a more
+// specific type.
+//
+// +stateify savable
+type rmAllocObject struct {
+	object
+
+	params capturedRmAllocParams
+}
+
+func newRmAllocObject[Params any](fd *frontendFD, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *Params) *rmAllocObject {
+	return &rmAllocObject{
+		params: captureRmAllocParams(fd, ioctlParams, rightsRequested, allocParams),
+	}
 }
 
 // Release implements objectImpl.Release.
-func (o *simpleObject) Release(ctx context.Context) {
+func (o *rmAllocObject) Release(ctx context.Context) {
 	// no-op
 }
 
 // rootClient is an objectImpl tracking a NV01_ROOT_CLIENT.
+//
+// +stateify savable
 type rootClient struct {
 	object
 
 	// These fields are protected by nvproxy.objsMu.
 	resources map[nvgpu.Handle]*object
 
-	// fd is the frontendFD that owns this client.
-	fd *frontendFD
+	params capturedRmAllocParams
 }
 
-func newRootClient(fd *frontendFD) *rootClient {
+func newRootClient(fd *frontendFD, ioctlParams *nvgpu.NVOS64Parameters, rightsRequested nvgpu.RS_ACCESS_MASK, allocParams *nvgpu.Handle) *rootClient {
 	return &rootClient{
 		resources: make(map[nvgpu.Handle]*object),
-		fd:        fd,
+		params:    captureRmAllocParams(fd, ioctlParams, rightsRequested, allocParams),
 	}
 }
 
 // Release implements objectImpl.Release.
 func (o *rootClient) Release(ctx context.Context) {
-	delete(o.fd.clients, o.handle)
+	delete(o.params.fd.clients, o.handle)
 }
 
 // osDescMem is an objectImpl tracking a NV01_MEMORY_SYSTEM_OS_DESCRIPTOR.
@@ -288,3 +323,23 @@ func (o *osDescMem) Release(ctx context.Context) {
 		}
 	})
 }
+
+// osEvent is an objectImpl tracking a NV01_EVENT_OS_EVENT.
+type osEvent struct {
+	object
+}
+
+// Release implements objectImpl.Release.
+func (o *osEvent) Release(ctx context.Context) {
+	// no-op
+}
+
+// virtMem is an objectImpl tracking a NV50_MEMORY_VIRTUAL.
+type virtMem struct {
+	object
+}
+
+// Release implements objectImpl.Release.
+func (o *virtMem) Release(ctx context.Context) {
+	// no-op
+}

runsc/boot/restore.go

@@ -225,7 +225,7 @@ func (r *restorer) restore(l *Loader) error {
 
 	log.Debugf("Restore using fdmap: %v", fdmap)
 	ctx = context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap)
-	log.Debugf("Restore using mfmap: %v", fdmap)
+	log.Debugf("Restore using mfmap: %v", mfmap)
 	ctx = context.WithValue(ctx, pgalloc.CtxMemoryFileMap, mfmap)
 	ctx = context.WithValue(ctx, devutil.CtxDevGoferClientProvider, l.k)