nvproxy: track driver object dependencies

This is necessary to prevent osDescMem pinned page leaks when the osDescMem is
freed indirectly, e.g. by closing the `/dev/nvidiactl` FD that owns the
containing client.

- Switch from passing separate sentryIoctlParams to temporarily modifying
  ioctlParams in-place, in order to ensure that callers of *Invoke functions
  always observe ioctlParams updated by driver copy-out (including e.g. updated
  handles).

- Only pass NVOS64Parameters (rather than NVOS21Parameters) in NV_ESC_RM_ALLOC
  calls to the driver; this should behave equivalently (new comment in
  rmAllocInvoke()) and simplifies passing them to object registration callbacks
  (objAddLocked).

PiperOrigin-RevId: 627630863

pkg/abi/nvgpu/ctrl.go

@@ -259,6 +259,13 @@ const (
 	NV503C_CTRL_CMD_UNREGISTER_VIDMEM = 0x503c0105
 )
 
+// +marshal
+type NV503C_CTRL_REGISTER_VA_SPACE_PARAMS struct {
+	HVASpace     Handle
+	Pad          [4]byte
+	VASpaceToken uint64
+}
+
 // From src/common/sdk/nvidia/inc/ctrl/ctrl83de/ctrl83dedebug.h:
 const (
 	NV83DE_CTRL_CMD_DEBUG_SET_EXCEPTION_MASK        = 0x83de0309

pkg/abi/nvgpu/status.go

@@ -16,6 +16,7 @@ package nvgpu
 
 // Status codes, from src/common/sdk/nvidia/inc/nvstatuscodes.h.
 const (
+	NV_OK                   = 0x00000000
 	NV_ERR_INVALID_ADDRESS  = 0x0000001e
 	NV_ERR_INVALID_ARGUMENT = 0x0000001f
 	NV_ERR_INVALID_CLASS    = 0x00000022

pkg/sentry/devices/nvproxy/BUILD

@@ -1,5 +1,6 @@
 load("//pkg/sync/locking:locking.bzl", "declare_mutex")
 load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools/go_generics:defs.bzl", "go_template_instance")
 
 package(default_applicable_licenses = ["//:license"])
 
@@ -12,6 +13,18 @@ declare_mutex(
     prefix = "objs",
 )
 
+go_template_instance(
+    name = "object_free_list",
+    out = "object_free_list.go",
+    package = "nvproxy",
+    prefix = "objectFree",
+    template = "//pkg/ilist:generic_list",
+    types = {
+        "Element": "*object",
+        "Linker": "*object",
+    },
+)
+
 go_library(
     name = "nvproxy",
     srcs = [
@@ -20,6 +33,8 @@ go_library(
         "frontend_unsafe.go",
         "nvproxy.go",
         "nvproxy_unsafe.go",
+        "object.go",
+        "object_free_list.go",
         "objs_mutex.go",
         "save_restore.go",
         "seccomp_filters.go",

pkg/sentry/devices/nvproxy/frontend.go

@@ -101,13 +101,25 @@ type frontendFD struct {
 	queue waiter.Queue
 
 	haveMmapContext atomic.Bool
+
+	// clients are handles of clients owned by this frontendFD. clients is
+	// protected by nvp.objsMu.
+	clients map[nvgpu.Handle]struct{} `state:"nosave"`
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
-func (fd *frontendFD) Release(context.Context) {
+func (fd *frontendFD) Release(ctx context.Context) {
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
+
+	fd.nvp.objsLock()
+	defer fd.nvp.objsUnlock()
 	unix.Close(int(fd.hostFD))
+	// src/nvidia/arch/nvalloc/unix/src/osapi.c:rm_cleanup_file_private() =>
+	// RmFreeUnusedClients()
+	for h := range fd.clients {
+		fd.nvp.objFree(ctx, h, h)
+	}
 }
 
 // EventRegister implements waiter.Waitable.EventRegister.
@@ -150,7 +162,7 @@ func (fd *frontendFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 		panic("Ioctl should be called from a task context")
 	}
 
-	if log.IsLogging(log.Debug) {
+	if ctx.IsLogging(log.Debug) {
 		ctx.Debugf("nvproxy: frontend ioctl: nr = %d = %#x, argSize = %d", nr, nr, argSize)
 	}
 
@@ -244,11 +256,9 @@ func frontendRegisterFD(fi *frontendIoctlState) (uintptr, error) {
 	if !ok {
 		return 0, linuxerr.EINVAL
 	}
-	sentryIoctlParams := nvgpu.IoctlRegisterFD{
-		CtlFD: ctlFile.hostFD,
-	}
+	ioctlParams.CtlFD = ctlFile.hostFD
 	// The returned ctl_fd can't change, so skip copying out.
-	return frontendIoctlInvoke(fi, &sentryIoctlParams)
+	return frontendIoctlInvoke(fi, &ioctlParams)
 }
 
 func rmAllocOSEvent(fi *frontendIoctlState) (uintptr, error) {
@@ -268,17 +278,16 @@ func rmAllocOSEvent(fi *frontendIoctlState) (uintptr, error) {
 	if !ok {
 		return 0, linuxerr.EINVAL
 	}
-	sentryIoctlParams := ioctlParams
-	sentryIoctlParams.FD = uint32(eventFile.hostFD)
 
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	origFD := ioctlParams.FD
+	ioctlParams.FD = uint32(eventFile.hostFD)
+	n, err := frontendIoctlInvoke(fi, &ioctlParams)
+	ioctlParams.FD = origFD
 	if err != nil {
 		return n, err
 	}
 
-	outIoctlParams := sentryIoctlParams
-	outIoctlParams.FD = ioctlParams.FD
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 
@@ -302,17 +311,16 @@ func rmFreeOSEvent(fi *frontendIoctlState) (uintptr, error) {
 	if !ok {
 		return 0, linuxerr.EINVAL
 	}
-	sentryIoctlParams := ioctlParams
-	sentryIoctlParams.FD = uint32(eventFile.hostFD)
 
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	origFD := ioctlParams.FD
+	ioctlParams.FD = uint32(eventFile.hostFD)
+	n, err := frontendIoctlInvoke(fi, &ioctlParams)
+	ioctlParams.FD = origFD
 	if err != nil {
 		return n, err
 	}
 
-	outIoctlParams := sentryIoctlParams
-	outIoctlParams.FD = ioctlParams.FD
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 
@@ -347,9 +355,8 @@ func rmAllocOSDescriptor(fi *frontendIoctlState, ioctlParams *nvgpu.IoctlNVOS02P
 	// Compare src/nvidia/arch/nvalloc/unix/src/escape.c:RmAllocOsDescriptor()
 	// => RmCreateOsDescriptor().
 	failWithStatus := func(status uint32) error {
-		outIoctlParams := *ioctlParams
-		outIoctlParams.Params.Status = status
-		_, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr)
+		ioctlParams.Params.Status = status
+		_, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr)
 		return err
 	}
 	appAddr := addrFromP64(ioctlParams.Params.PMemory)
@@ -382,15 +389,13 @@ func rmAllocOSDescriptor(fi *frontendIoctlState, ioctlParams *nvgpu.IoctlNVOS02P
 	if errno != 0 {
 		return 0, errno
 	}
-	cu := cleanup.Make(func() {
-		unix.RawSyscall(unix.SYS_MUNMAP, m, uintptr(arLen), 0)
-	})
-	defer cu.Clean()
+	defer unix.RawSyscall(unix.SYS_MUNMAP, m, uintptr(arLen), 0)
 	// Mirror application mappings into the reserved range.
 	prs, err := fi.t.MemoryManager().Pin(fi.ctx, appAR, at, false /* ignorePermissions */)
-	cu.Add(func() {
+	unpinCleanup := cleanup.Make(func() {
 		mm.Unpin(prs)
 	})
+	defer unpinCleanup.Clean()
 	if err != nil {
 		return 0, err
 	}
@@ -409,35 +414,34 @@ func rmAllocOSDescriptor(fi *frontendIoctlState, ioctlParams *nvgpu.IoctlNVOS02P
 			ims = ims.Tail()
 		}
 	}
-	sentryIoctlParams := *ioctlParams
-	sentryIoctlParams.Params.PMemory = nvgpu.P64(uint64(m))
+	origPMemory := ioctlParams.Params.PMemory
+	ioctlParams.Params.PMemory = nvgpu.P64(uint64(m))
 	// NV01_MEMORY_SYSTEM_OS_DESCRIPTOR shouldn't use ioctlParams.FD; clobber
 	// it to be sure.
-	sentryIoctlParams.FD = -1
+	origFD := ioctlParams.FD
+	ioctlParams.FD = -1
 
-	fi.fd.nvp.objsMu.Lock()
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	fi.fd.nvp.objsLock()
+	n, err := frontendIoctlInvoke(fi, ioctlParams)
+	if err == nil && ioctlParams.Params.Status == nvgpu.NV_OK {
+		// Transfer ownership of pinned pages to an osDescMem object, to be
+		// unpinned when the driver OsDescMem is freed.
+		fi.fd.nvp.objAdd(fi.ctx, ioctlParams.Params.HRoot, ioctlParams.Params.HObjectNew, nvgpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR, &osDescMem{
+			pinnedRanges: prs,
+		}, ioctlParams.Params.HObjectParent)
+		unpinCleanup.Release()
+		if fi.ctx.IsLogging(log.Debug) {
+			fi.ctx.Debugf("nvproxy: pinned %d bytes for OS descriptor with handle %v", arLen, ioctlParams.Params.HObjectNew)
+		}
+	}
+	fi.fd.nvp.objsUnlock()
+	ioctlParams.Params.PMemory = origPMemory
+	ioctlParams.FD = origFD
 	if err != nil {
-		fi.fd.nvp.objsMu.Unlock()
 		return n, err
 	}
-	// Transfer ownership of pinned pages to an osDescMem object, to be
-	// unpinned when the driver OsDescMem is freed.
-	o := &osDescMem{
-		pinnedRanges: prs,
-	}
-	o.object.init(o)
-	fi.fd.nvp.objsLive[sentryIoctlParams.Params.HObjectNew] = &o.object
-	fi.fd.nvp.objsMu.Unlock()
-	cu.Release()
-	fi.ctx.Infof("nvproxy: pinned pages for OS descriptor with handle %v", sentryIoctlParams.Params.HObjectNew)
-	// Unmap the reserved range, which is no longer required.
-	unix.RawSyscall(unix.SYS_MUNMAP, m, uintptr(arLen), 0)
 
-	outIoctlParams := sentryIoctlParams
-	outIoctlParams.Params.PMemory = ioctlParams.Params.PMemory
-	outIoctlParams.FD = ioctlParams.FD
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 
@@ -453,20 +457,15 @@ func rmFree(fi *frontendIoctlState) (uintptr, error) {
 		return 0, err
 	}
 
-	fi.fd.nvp.objsMu.Lock()
+	fi.fd.nvp.objsLock()
 	n, err := frontendIoctlInvoke(fi, &ioctlParams)
+	if err == nil && ioctlParams.Status == nvgpu.NV_OK {
+		fi.fd.nvp.objFree(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectOld)
+	}
+	fi.fd.nvp.objsUnlock()
 	if err != nil {
-		fi.fd.nvp.objsMu.Unlock()
 		return n, err
 	}
-	o, ok := fi.fd.nvp.objsLive[ioctlParams.HObjectOld]
-	if ok {
-		delete(fi.fd.nvp.objsLive, ioctlParams.HObjectOld)
-	}
-	fi.fd.nvp.objsMu.Unlock()
-	if ok {
-		o.Release(fi.ctx)
-	}
 
 	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
@@ -485,7 +484,7 @@ func rmControl(fi *frontendIoctlState) (uintptr, error) {
 
 	// Cmd determines the type of Params.
 	if log.IsLogging(log.Debug) {
-		fi.ctx.Debugf("nvproxy: control command %#x", ioctlParams.Cmd)
+		fi.ctx.Debugf("nvproxy: control command %#x, object %#x", ioctlParams.Cmd, ioctlParams.HObject.Val)
 	}
 	if ioctlParams.Cmd&nvgpu.RM_GSS_LEGACY_MASK != 0 {
 		// This is a "legacy GSS control" that is implemented by the GPU System
@@ -550,9 +549,8 @@ func rmControlSimple(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters
 }
 
 func ctrlCmdFailWithStatus(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters, status uint32) error {
-	outIoctlParams := *ioctlParams
-	outIoctlParams.Status = status
-	_, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr)
+	ioctlParams.Status = status
+	_, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr)
 	return err
 }
 
@@ -625,6 +623,31 @@ func ctrlDevGpuGetClasslist(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Par
 	return n, nil
 }
 
+func ctrlRegisterVASpace(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters) (uintptr, error) {
+	var ctrlParams nvgpu.NV503C_CTRL_REGISTER_VA_SPACE_PARAMS
+	if ctrlParams.SizeBytes() != int(ioctlParams.ParamsSize) {
+		return 0, linuxerr.EINVAL
+	}
+	if _, err := ctrlParams.CopyIn(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+		return 0, err
+	}
+	fi.fd.nvp.objsLock()
+	n, err := rmControlInvoke(fi, ioctlParams, &ctrlParams)
+	if err == nil && ioctlParams.Status == nvgpu.NV_OK {
+		// src/nvidia/src/kernel/gpu/bus/third_party_p2p.c:CliAddThirdPartyP2PVASpace()
+		// => refAddDependant()
+		fi.fd.nvp.objAddDep(ioctlParams.HClient, ioctlParams.HObject, ctrlParams.HVASpace)
+	}
+	fi.fd.nvp.objsUnlock()
+	if err != nil {
+		return n, err
+	}
+	if _, err := ctrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+		return n, err
+	}
+	return n, nil
+}
+
 func ctrlSubdevFIFODisableChannels(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters) (uintptr, error) {
 	var ctrlParams nvgpu.NV2080_CTRL_FIFO_DISABLE_CHANNELS_PARAMS
 	if ctrlParams.SizeBytes() != int(ioctlParams.ParamsSize) {
@@ -659,7 +682,7 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 	default:
 		return 0, linuxerr.EINVAL
 	}
-	// Copy in parameters and convert to NVOS64ParametersV535, which is a super
+	// Copy in parameters and convert to NVOS64Parameters, which is a super
 	// set of all parameter types we support.
 	buf := nvgpu.GetRmAllocParamObj(isNVOS64)
 	if _, err := buf.CopyIn(fi.t, fi.ioctlParamsAddr); err != nil {
@@ -678,6 +701,9 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 	// ("External Class") to the type of pAllocParms ("Alloc Param Info") and
 	// the class whose constructor interprets it ("Internal Class").
 	// - Add symbol and parameter type definitions to //pkg/abi/nvgpu.
+	// - Check constructor for calls to refAddDependant(),
+	// sessionAddDependant(), or sessionAddDependency(), which need to be
+	// mirrored by dependencies in the call to nvproxy.objAddLocked().
 	// - Add handling below.
 	handler := fi.fd.nvp.abi.allocationClass[ioctlParams.HClass]
 	if handler == nil {
@@ -697,18 +723,32 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 	return handler(fi, &ioctlParams, isNVOS64)
 }
 
+// rmAllocSimple implements NV_ESC_RM_ALLOC for classes whose parameters don't
+// contain any pointers or file descriptors requiring translation, and whose
+// objects require no special handling and depend only on their parents.
+//
 // Unlike frontendIoctlSimple and rmControlSimple, rmAllocSimple requires the
 // parameter type since the parameter's size is otherwise unknown.
 func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
+	return rmAllocSimpleParams[Params, PParams](fi, ioctlParams, isNVOS64, addSimpleObjDepParentLocked)
+}
+
+// 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.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), 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) {
 	if ioctlParams.PAllocParms == 0 {
-		return rmAllocInvoke[byte](fi, ioctlParams, nil, isNVOS64)
+		return rmAllocInvoke[Params](fi, ioctlParams, nil, isNVOS64, objAddLocked)
 	}
 
 	var allocParams Params
 	if _, err := (PParams)(&allocParams).CopyIn(fi.t, addrFromP64(ioctlParams.PAllocParms)); err != nil {
 		return 0, err
 	}
-	n, err := rmAllocInvoke(fi, ioctlParams, &allocParams, isNVOS64)
+	n, err := rmAllocInvoke(fi, ioctlParams, &allocParams, isNVOS64, objAddLocked)
 	if err != nil {
 		return n, err
 	}
@@ -719,7 +759,17 @@ func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlStat
 }
 
 func rmAllocNoParams(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
-	return rmAllocInvoke[byte](fi, ioctlParams, nil, isNVOS64)
+	return rmAllocInvoke[byte](fi, ioctlParams, nil, isNVOS64, addSimpleObjDepParentLocked)
+}
+
+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.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newRootClient(fi.fd))
+		if fi.fd.clients == nil {
+			fi.fd.clients = make(map[nvgpu.Handle]struct{})
+		}
+		fi.fd.clients[ioctlParams.HObjectNew] = struct{}{}
+	})
 }
 
 func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
@@ -736,22 +786,76 @@ func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parame
 	if !ok {
 		return 0, linuxerr.EINVAL
 	}
-	sentryAllocParams := allocParams
-	sentryAllocParams.Data = nvgpu.P64(uint64(eventFile.hostFD))
+	origData := allocParams.Data
+	allocParams.Data = nvgpu.P64(uint64(eventFile.hostFD))
 
-	n, err := rmAllocInvoke(fi, ioctlParams, &sentryAllocParams, isNVOS64)
+	n, err := rmAllocInvoke(fi, ioctlParams, &allocParams, isNVOS64, addSimpleObjDepParentLocked)
 	if err != nil {
 		return n, err
 	}
 
-	outAllocParams := sentryAllocParams
-	outAllocParams.Data = allocParams.Data
-	if _, err := outAllocParams.CopyOut(fi.t, addrFromP64(ioctlParams.PAllocParms)); err != nil {
+	allocParams.Data = origData
+	if _, err := allocParams.CopyOut(fi.t, addrFromP64(ioctlParams.PAllocParms)); err != nil {
 		return n, err
 	}
 	return n, nil
 }
 
+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) {
+		// 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.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), 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) {
+		// See
+		// src/nvidia/src/kernel/gpu/fifo/kernel_channel_group_api.c:kchangrpapiConstruct_IMPL()
+		// => refAddDependant().
+		fi.fd.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), 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
+		// group's (the graphics context is freed when the channel group is).
+		// Channels, context shares, and graphics objects depend on this
+		// graphics context rather than the channel group. Consequently, if MIG
+		// is enabled, these might not depend on the channel group at all.
+		// Since nvproxy currently does not support MIG, we represent these
+		// dependencies as unconditionally on the channel group instead.
+	})
+}
+
+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) {
+		// See
+		// src/nvidia/src/kernel/gpu/fifo/kernel_channel.c:kchannelConstruct_IMPL()
+		// => refAddDependant(). The channel's parent may be a device or
+		// channel group; if it is a channel group then the channel depends on
+		// it via the parent relationship, and if it is not a channel group
+		// then kchannelConstruct_IMPL() constructs one internally and frees it
+		// when the channel is destroyed, so either way no separate dependency
+		// is required.
+		fi.fd.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), 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) {
+		// 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.nvp.objAdd(fi.ctx, ioctlParams.HRoot, ioctlParams.HObjectNew, ioctlParams.HClass, newSimpleObject(), ioctlParams.HObjectParent, allocParams.HVASpace)
+	})
+}
+
 func rmVidHeapControl(fi *frontendIoctlState) (uintptr, error) {
 	var ioctlParams nvgpu.NVOS32Parameters
 	if fi.ioctlParamsSize != nvgpu.SizeofNVOS32Parameters {
@@ -762,7 +866,7 @@ func rmVidHeapControl(fi *frontendIoctlState) (uintptr, error) {
 	}
 
 	// Function determines the type of Data.
-	if log.IsLogging(log.Debug) {
+	if fi.ctx.IsLogging(log.Debug) {
 		fi.ctx.Debugf("nvproxy: VID_HEAP_CONTROL function %d", ioctlParams.Function)
 	}
 	// See
@@ -798,17 +902,16 @@ func rmMapMemory(fi *frontendIoctlState) (uintptr, error) {
 		fi.ctx.Warningf("nvproxy: attempted to reuse FD %d for NV_ESC_RM_MAP_MEMORY", ioctlParams.FD)
 		return 0, linuxerr.EINVAL
 	}
-	sentryIoctlParams := ioctlParams
-	sentryIoctlParams.FD = mapFile.hostFD
+	origFD := ioctlParams.FD
+	ioctlParams.FD = mapFile.hostFD
 
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	n, err := frontendIoctlInvoke(fi, &ioctlParams)
 	if err != nil {
 		return n, err
 	}
 
-	outIoctlParams := sentryIoctlParams
-	outIoctlParams.FD = ioctlParams.FD
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	ioctlParams.FD = origFD
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 

pkg/sentry/devices/nvproxy/frontend_unsafe.go

@@ -32,62 +32,56 @@ func frontendIoctlInvoke[Params any](fi *frontendIoctlState, sentryParams *Param
 	return n, nil
 }
 
-func frontendIoctlInvokePtr(fi *frontendIoctlState, sentryParams uintptr) (uintptr, error) {
-	n, _, errno := unix.RawSyscall(unix.SYS_IOCTL, uintptr(fi.fd.hostFD), frontendIoctlCmd(fi.nr, fi.ioctlParamsSize), sentryParams)
-	if errno != 0 {
-		return n, errno
-	}
-	return n, nil
-}
-
 func rmControlInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters, ctrlParams *Params) (uintptr, error) {
 	defer runtime.KeepAlive(ctrlParams) // since we convert to non-pointer-typed P64
-	sentryIoctlParams := *ioctlParams
-	sentryIoctlParams.Params = p64FromPtr(unsafe.Pointer(ctrlParams))
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	origParams := ioctlParams.Params
+	ioctlParams.Params = p64FromPtr(unsafe.Pointer(ctrlParams))
+	n, err := frontendIoctlInvoke(fi, ioctlParams)
+	ioctlParams.Params = origParams
 	if err != nil {
 		return n, err
 	}
-	outIoctlParams := sentryIoctlParams
-	outIoctlParams.Params = ioctlParams.Params
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 	return n, nil
 }
 
 func ctrlClientSystemGetBuildVersionInvoke(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters, ctrlParams *nvgpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_PARAMS, driverVersionBuf, versionBuf, titleBuf *byte) (uintptr, error) {
-	sentryCtrlParams := *ctrlParams
-	sentryCtrlParams.PDriverVersionBuffer = p64FromPtr(unsafe.Pointer(driverVersionBuf))
-	sentryCtrlParams.PVersionBuffer = p64FromPtr(unsafe.Pointer(versionBuf))
-	sentryCtrlParams.PTitleBuffer = p64FromPtr(unsafe.Pointer(titleBuf))
-	n, err := rmControlInvoke(fi, ioctlParams, &sentryCtrlParams)
+	// *Buf arguments don't need runtime.KeepAlive() since our caller
+	// ctrlClientSystemGetBuildVersion() copies them out, keeping them alive
+	// during this function.
+	origPDriverVersionBuffer := ctrlParams.PDriverVersionBuffer
+	origPVersionBuffer := ctrlParams.PVersionBuffer
+	origPTitleBuffer := ctrlParams.PTitleBuffer
+	ctrlParams.PDriverVersionBuffer = p64FromPtr(unsafe.Pointer(driverVersionBuf))
+	ctrlParams.PVersionBuffer = p64FromPtr(unsafe.Pointer(versionBuf))
+	ctrlParams.PTitleBuffer = p64FromPtr(unsafe.Pointer(titleBuf))
+	n, err := rmControlInvoke(fi, ioctlParams, ctrlParams)
+	ctrlParams.PDriverVersionBuffer = origPDriverVersionBuffer
+	ctrlParams.PVersionBuffer = origPVersionBuffer
+	ctrlParams.PTitleBuffer = origPTitleBuffer
 	if err != nil {
 		return n, err
 	}
-	outCtrlParams := sentryCtrlParams
-	outCtrlParams.PDriverVersionBuffer = ctrlParams.PDriverVersionBuffer
-	outCtrlParams.PVersionBuffer = ctrlParams.PVersionBuffer
-	outCtrlParams.PTitleBuffer = ctrlParams.PTitleBuffer
-	if _, err := outCtrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+	if _, err := ctrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
 		return n, err
 	}
 	return n, nil
 }
 
 func ctrlDevGpuGetClasslistInvoke(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters, ctrlParams *nvgpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS, classList []uint32) (uintptr, error) {
-	sentryCtrlParams := *ctrlParams
-	sentryCtrlParams.ClassList = p64FromPtr(unsafe.Pointer(&classList[0]))
-	n, err := rmControlInvoke(fi, ioctlParams, &sentryCtrlParams)
+	origClassList := ctrlParams.ClassList
+	ctrlParams.ClassList = p64FromPtr(unsafe.Pointer(&classList[0]))
+	n, err := rmControlInvoke(fi, ioctlParams, ctrlParams)
+	ctrlParams.ClassList = origClassList
 	if err != nil {
 		return n, err
 	}
-	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(ctrlParams.ClassList), classList); err != nil {
+	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(origClassList), classList); err != nil {
 		return 0, err
 	}
-	outCtrlParams := sentryCtrlParams
-	outCtrlParams.ClassList = ctrlParams.ClassList
-	if _, err := outCtrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+	if _, err := ctrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
 		return n, err
 	}
 	return n, nil
@@ -114,25 +108,25 @@ func ctrlDevFIFOGetChannelList(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54
 	if _, err := primitive.CopyUint32SliceIn(fi.t, addrFromP64(ctrlParams.PChannelList), channelList); err != nil {
 		return 0, err
 	}
-	sentryCtrlParams := ctrlParams
-	sentryCtrlParams.PChannelHandleList = p64FromPtr(unsafe.Pointer(&channelHandleList[0]))
-	sentryCtrlParams.PChannelList = p64FromPtr(unsafe.Pointer(&channelList[0]))
 
-	n, err := rmControlInvoke(fi, ioctlParams, &sentryCtrlParams)
+	origPChannelHandleList := ctrlParams.PChannelHandleList
+	origPChannelList := ctrlParams.PChannelList
+	ctrlParams.PChannelHandleList = p64FromPtr(unsafe.Pointer(&channelHandleList[0]))
+	ctrlParams.PChannelList = p64FromPtr(unsafe.Pointer(&channelList[0]))
+	n, err := rmControlInvoke(fi, ioctlParams, &ctrlParams)
+	ctrlParams.PChannelHandleList = origPChannelHandleList
+	ctrlParams.PChannelList = origPChannelList
 	if err != nil {
 		return n, err
 	}
 
-	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(ctrlParams.PChannelHandleList), channelHandleList); err != nil {
+	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(origPChannelHandleList), channelHandleList); err != nil {
 		return 0, err
 	}
-	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(ctrlParams.PChannelList), channelList); err != nil {
+	if _, err := primitive.CopyUint32SliceOut(fi.t, addrFromP64(origPChannelList), channelList); err != nil {
 		return 0, err
 	}
-	outCtrlParams := sentryCtrlParams
-	outCtrlParams.PChannelHandleList = ctrlParams.PChannelHandleList
-	outCtrlParams.PChannelList = ctrlParams.PChannelList
-	if _, err := outCtrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+	if _, err := ctrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
 		return n, err
 	}
 
@@ -156,54 +150,66 @@ func ctrlSubdevGRGetInfo(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parame
 	if _, err := fi.t.CopyInBytes(addrFromP64(ctrlParams.GRInfoList), infoList); err != nil {
 		return 0, err
 	}
-	sentryCtrlParams := ctrlParams
-	sentryCtrlParams.GRInfoList = p64FromPtr(unsafe.Pointer(&infoList[0]))
 
-	n, err := rmControlInvoke(fi, ioctlParams, &sentryCtrlParams)
+	origGRInfoList := ctrlParams.GRInfoList
+	ctrlParams.GRInfoList = p64FromPtr(unsafe.Pointer(&infoList[0]))
+	n, err := rmControlInvoke(fi, ioctlParams, &ctrlParams)
+	ctrlParams.GRInfoList = origGRInfoList
 	if err != nil {
 		return n, err
 	}
 
-	if _, err := fi.t.CopyOutBytes(addrFromP64(ctrlParams.GRInfoList), infoList); err != nil {
+	if _, err := fi.t.CopyOutBytes(addrFromP64(origGRInfoList), infoList); err != nil {
 		return n, err
 	}
-	outCtrlParams := sentryCtrlParams
-	outCtrlParams.GRInfoList = ctrlParams.GRInfoList
-	if _, err := outCtrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
+	if _, err := ctrlParams.CopyOut(fi.t, addrFromP64(ioctlParams.Params)); err != nil {
 		return n, err
 	}
 
 	return n, nil
 }
 
-func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params, isNVOS64 bool) (uintptr, error) {
+func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params, isNVOS64 bool, addObjLocked func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params)) (uintptr, error) {
 	defer runtime.KeepAlive(allocParams) // since we convert to non-pointer-typed P64
 
-	sentryIoctlParams := nvgpu.GetRmAllocParamObj(isNVOS64)
-	sentryIoctlParams.FromOS64(*ioctlParams)
-	sentryIoctlParams.SetPAllocParms(p64FromPtr(unsafe.Pointer(allocParams)))
+	// Temporarily replace application pointers with sentry pointers.
+	origPAllocParms := ioctlParams.PAllocParms
+	origPRightsRequested := ioctlParams.PRightsRequested
 	var rightsRequested nvgpu.RS_ACCESS_MASK
 	if ioctlParams.PRightsRequested != 0 {
 		if _, err := rightsRequested.CopyIn(fi.t, addrFromP64(ioctlParams.PRightsRequested)); err != nil {
 			return 0, err
 		}
-		sentryIoctlParams.SetPRightsRequested(p64FromPtr(unsafe.Pointer(&rightsRequested)))
+		ioctlParams.PRightsRequested = p64FromPtr(unsafe.Pointer(&rightsRequested))
 	}
-	n, err := frontendIoctlInvokePtr(fi, sentryIoctlParams.GetPointer())
-	if err != nil {
-		return n, err
+	ioctlParams.PAllocParms = p64FromPtr(unsafe.Pointer(allocParams))
+
+	// Invoke the driver ioctl and restore application pointers. We always pass
+	// NVOS64Parameters to the driver even if !isNVOS64, as this is handled
+	// identically to the equivalent NVOS21Parameters; compare
+	// src/nvidia/src/kernel/rmapi/entry_points.c:_nv04AllocWithSecInfo() and
+	// _nv04AllocWithAccessSecInfo().
+	fi.fd.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)
 	}
+	fi.fd.nvp.objsUnlock()
+	ioctlParams.PAllocParms = origPAllocParms
+	ioctlParams.PRightsRequested = origPRightsRequested
+	if errno != 0 {
+		return n, errno
+	}
+
+	// Copy updated params out to the application.
+	outIoctlParams := nvgpu.GetRmAllocParamObj(isNVOS64)
+	outIoctlParams.FromOS64(*ioctlParams)
 	if ioctlParams.PRightsRequested != 0 {
 		if _, err := rightsRequested.CopyOut(fi.t, addrFromP64(ioctlParams.PRightsRequested)); err != nil {
 			return n, err
 		}
 	}
-	// Reuse sentryIoctlParams to write out params.
-	sentryIoctlParams.SetPAllocParms(ioctlParams.PAllocParms)
-	if ioctlParams.PRightsRequested != 0 {
-		sentryIoctlParams.SetPRightsRequested(ioctlParams.PRightsRequested)
-	}
-	if _, err := sentryIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 	return n, nil
@@ -211,31 +217,33 @@ func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64
 
 func rmVidHeapControlAllocSize(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS32Parameters) (uintptr, error) {
 	allocSizeParams := (*nvgpu.NVOS32AllocSize)(unsafe.Pointer(&ioctlParams.Data))
-
-	sentryIoctlParams := *ioctlParams
-	sentryAllocSizeParams := (*nvgpu.NVOS32AllocSize)(unsafe.Pointer(&sentryIoctlParams.Data))
+	origAddress := allocSizeParams.Address
 	var addr uint64
 	if allocSizeParams.Address != 0 {
 		if _, err := primitive.CopyUint64In(fi.t, addrFromP64(allocSizeParams.Address), &addr); err != nil {
 			return 0, err
 		}
-		sentryAllocSizeParams.Address = p64FromPtr(unsafe.Pointer(&addr))
+		allocSizeParams.Address = p64FromPtr(unsafe.Pointer(&addr))
 	}
 
-	n, err := frontendIoctlInvoke(fi, &sentryIoctlParams)
+	fi.fd.nvp.objsLock()
+	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.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, nvgpu.NV50_MEMORY_VIRTUAL, newSimpleObject(), ioctlParams.HObjectParent, ioctlParams.HVASpace)
+	}
+	fi.fd.nvp.objsUnlock()
+	allocSizeParams.Address = origAddress
 	if err != nil {
 		return n, err
 	}
 
-	outIoctlParams := sentryIoctlParams
-	outAllocSizeParams := (*nvgpu.NVOS32AllocSize)(unsafe.Pointer(&outIoctlParams.Data))
 	if allocSizeParams.Address != 0 {
 		if _, err := primitive.CopyUint64Out(fi.t, addrFromP64(allocSizeParams.Address), addr); err != nil {
 			return n, err
 		}
-		outAllocSizeParams.Address = allocSizeParams.Address
 	}
-	if _, err := outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
+	if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil {
 		return n, err
 	}
 

pkg/sentry/devices/nvproxy/nvproxy.go

@@ -22,11 +22,9 @@ import (
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/abi/nvgpu"
-	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/marshal"
-	"gvisor.dev/gvisor/pkg/sentry/mm"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
@@ -44,9 +42,10 @@ func Register(vfsObj *vfs.VirtualFilesystem, versionStr string, uvmDevMajor uint
 		return fmt.Errorf("unsupported Nvidia driver version: %s", versionStr)
 	}
 	nvp := &nvproxy{
-		objsLive: make(map[nvgpu.Handle]*object),
-		abi:      abiCons.cons(),
-		version:  version,
+		abi:         abiCons.cons(),
+		version:     version,
+		clients:     make(map[nvgpu.Handle]*rootClient),
+		objsFreeSet: make(map[*object]struct{}),
 	}
 	for minor := uint32(0); minor <= nvgpu.NV_CONTROL_DEVICE_MINOR; minor++ {
 		if err := vfsObj.RegisterDevice(vfs.CharDevice, nvgpu.NV_MAJOR_DEVICE_NUMBER, minor, &frontendDevice{
@@ -70,44 +69,18 @@ func Register(vfsObj *vfs.VirtualFilesystem, versionStr string, uvmDevMajor uint
 
 // +stateify savable
 type nvproxy struct {
-	objsMu   objsMutex                `state:"nosave"`
-	objsLive map[nvgpu.Handle]*object `state:"nosave"`
-	abi      *driverABI               `state:"nosave"`
-	version  DriverVersion
-}
+	abi     *driverABI `state:"nosave"`
+	version DriverVersion
 
-// object tracks an object allocated through the driver.
-//
-// +stateify savable
-type object struct {
-	impl objectImpl
-}
-
-func (o *object) init(impl objectImpl) {
-	o.impl = impl
-}
-
-// Release is called after the represented object is freed.
-func (o *object) Release(ctx context.Context) {
-	o.impl.Release(ctx)
-}
-
-type objectImpl interface {
-	Release(ctx context.Context)
-}
-
-// osDescMem is an objectImpl tracking an OS descriptor.
-//
-// +stateify savable
-type osDescMem struct {
-	object
-	pinnedRanges []mm.PinnedRange
-}
-
-// Release implements objectImpl.Release.
-func (o *osDescMem) Release(ctx context.Context) {
-	ctx.Infof("nvproxy: unpinning pages for released OS descriptor")
-	mm.Unpin(o.pinnedRanges)
+	// See object.go.
+	// Users should call nvproxy.objsLock/Unlock() rather than locking objsMu
+	// directly.
+	objsMu objsMutex `state:"nosave"`
+	// These fields are protected by objsMu.
+	clients      map[nvgpu.Handle]*rootClient
+	objsCleanup  []func()             `state:"nosave"`
+	objsFreeList objectFreeList       `state:"nosave"`
+	objsFreeSet  map[*object]struct{} `state:"nosave"`
 }
 
 type marshalPtr[T any] interface {

pkg/sentry/devices/nvproxy/object.go

@@ -0,0 +1,290 @@
+// Copyright 2024 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 nvproxy
+
+import (
+	"gvisor.dev/gvisor/pkg/abi/nvgpu"
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/mm"
+)
+
+// object tracks a driver object.
+//
+// +stateify savable
+type object struct {
+	// These fields are initialized by nvproxy.objAdd() and are immutable thereafter.
+	nvp    *nvproxy
+	client *rootClient // may be == impl
+	class  nvgpu.ClassID
+	handle nvgpu.Handle // in client.resources, and also nvp.clients if impl is rootClient
+	impl   objectImpl
+
+	// The driver tracks parent/child relationships and "arbitrary dependency"
+	// relationships between objects separately; we treat parent/child
+	// relationships as equivalent to other dependencies. These fields are
+	// protected by nvp.objsMu.
+	deps  map[*object]struct{} // objects that this object depends on
+	rdeps map[*object]struct{} // objects that depend on this object
+	objectFreeEntry
+}
+
+type objectImpl interface {
+	// Object returns the object embedded in this objectImpl.
+	Object() *object
+
+	// Release is called when the driver object represented by this objectImpl
+	// is freed.
+	//
+	// Preconditions: nvproxy.objsMu must be locked.
+	Release(ctx context.Context)
+}
+
+// Object implements objectImpl.Object.
+func (o *object) Object() *object {
+	return o
+}
+
+func (nvp *nvproxy) objsLock() {
+	nvp.objsMu.Lock()
+}
+
+func (nvp *nvproxy) objsUnlock() {
+	cleanup := nvp.objsCleanup
+	nvp.objsCleanup = nil
+	nvp.objsMu.Unlock()
+	for _, f := range cleanup {
+		f()
+	}
+}
+
+// objAdd records the allocation of a driver object with class c and handle h,
+// in the client with handle clientH, represented by oi. Each non-zero handle
+// in deps is a dependency of the created object, such that the freeing of any
+// of those objects also results in the freeing of the recorded object.
+func (nvp *nvproxy) objAdd(ctx context.Context, clientH, h nvgpu.Handle, c nvgpu.ClassID, oi objectImpl, deps ...nvgpu.Handle) {
+	if h.Val == 0 {
+		log.Traceback("nvproxy: new object (class %v) has invalid handle 0", c)
+		return
+	}
+	var client *rootClient
+	// The driver forced NV01_ROOT and NV01_ROOT_NON_PRIV to NV01_ROOT_CLIENT,
+	// so we only need to check for the latter.
+	if c == nvgpu.NV01_ROOT_CLIENT {
+		clientH = h
+		client = oi.(*rootClient)
+		if _, ok := nvp.clients[h]; ok {
+			ctx.Warningf("nvproxy: client handle %v already in use", h)
+		}
+		nvp.clients[h] = client
+	} else {
+		var ok bool
+		client, ok = nvp.clients[clientH]
+		if !ok {
+			log.Traceback("nvproxy: new object %v (class %v) has invalid client handle %v", h, c, clientH)
+			return
+		}
+	}
+	o := oi.Object()
+	o.nvp = nvp
+	o.client = client
+	o.class = c
+	o.handle = h
+	o.impl = oi
+	if _, ok := client.resources[h]; ok {
+		ctx.Warningf("nvproxy: handle %v:%v already in use", clientH, h)
+	}
+	client.resources[h] = o
+	for _, depH := range deps {
+		if depH.Val == 0 /* aka NV01_NULL_OBJECT */ {
+			continue
+		}
+		dep, ok := client.resources[depH]
+		if !ok {
+			log.Traceback("nvproxy: new object %v:%v (class %v) has invalid dependency handle %v", clientH, h, c, depH)
+			continue
+		}
+		nvp.objDep(o, dep)
+	}
+	if ctx.IsLogging(log.Debug) {
+		ctx.Debugf("nvproxy: added object %v:%v (class %v) with dependencies %v", clientH, h, c, deps)
+	}
+}
+
+// objAddDep records a dependency between the existing object with handle h1 on
+// the existing object with handle h2, such that the freeing of the object with
+// handle h2 results in the freeing of object h1. Both h1 and h2 are handles in
+// the client with handle clientH.
+func (nvp *nvproxy) objAddDep(clientH, h1, h2 nvgpu.Handle) {
+	if h1.Val == 0 || h2.Val == 0 {
+		return
+	}
+	client, ok := nvp.clients[clientH]
+	if !ok {
+		log.Traceback("nvproxy: invalid client handle %v", clientH)
+		return
+	}
+	o1, ok := client.resources[h1]
+	if !ok {
+		log.Traceback("nvproxy: invalid handle %v:%v", clientH, h1)
+		return
+	}
+	o2, ok := client.resources[h2]
+	if !ok {
+		log.Traceback("nvproxy: invalid handle %v:%v", clientH, h2)
+		return
+	}
+	nvp.objDep(o1, o2)
+}
+
+func (nvp *nvproxy) objDep(o1, o2 *object) {
+	if o1.deps == nil {
+		o1.deps = make(map[*object]struct{})
+	}
+	o1.deps[o2] = struct{}{}
+	if o2.rdeps == nil {
+		o2.rdeps = make(map[*object]struct{})
+	}
+	o2.rdeps[o1] = struct{}{}
+}
+
+// objFree marks an object and its transitive dependents as freed.
+//
+// Compare
+// src/nvidia/src/libraries/resserv/src/rs_server.c:serverFreeResourceTree().
+func (nvp *nvproxy) objFree(ctx context.Context, clientH, h nvgpu.Handle) {
+	// Check for recursive calls to objFree() (via objectImpl.Release()).
+	// serverFreeResourceTree() permits this; we currently don't for
+	// simplicity.
+	if !nvp.objsFreeList.Empty() {
+		panic("nvproxy.objFree called with non-empty free list (possible recursion?)")
+	}
+
+	client, ok := nvp.clients[clientH]
+	if !ok {
+		ctx.Warningf("nvproxy: freeing object handle %v with unknown client handle %v", h, clientH)
+		return
+	}
+	o, ok := client.resources[h]
+	if !ok {
+		// When RS_COMPATABILITY_MODE is defined as true in the driver (as it
+		// is in Linux), the driver permits NV_ESC_RM_FREE on nonexistent
+		// handles as a no-op, and applications do this, so log at level INFO
+		// rather than WARNING.
+		ctx.Infof("nvproxy: freeing object with unknown handle %v:%v", clientH, h)
+		return
+	}
+	nvp.prependFreedLockedRecursive(o)
+	for !nvp.objsFreeList.Empty() {
+		o2 := nvp.objsFreeList.Front()
+		o2.impl.Release(ctx)
+		for o3 := range o2.deps {
+			delete(o3.rdeps, o2)
+		}
+		delete(o2.client.resources, o2.handle)
+		if o2.class == nvgpu.NV01_ROOT_CLIENT {
+			delete(nvp.clients, o2.handle)
+		}
+		nvp.objsFreeList.Remove(o2)
+		delete(nvp.objsFreeSet, o2)
+		if ctx.IsLogging(log.Debug) {
+			ctx.Debugf("nvproxy: freed object %v:%v (class %v)", o2.client.handle, o2.handle, o2.class)
+		}
+	}
+}
+
+func (nvp *nvproxy) prependFreedLockedRecursive(o *object) {
+	if _, ok := nvp.objsFreeSet[o]; ok {
+		// o is already on the free list; move it to the front so that it
+		// remains freed before our caller's o.
+		nvp.objsFreeList.Remove(o)
+	} else {
+		nvp.objsFreeSet[o] = struct{}{}
+	}
+	nvp.objsFreeList.PushFront(o)
+
+	// In the driver, freeing an object causes its children and dependents to
+	// be freed first; see
+	// src/nvidia/src/libraries/resserv/src/rs_server.c:serverFreeResourceTree()
+	// => clientUpdatePendingFreeList_IMPL(). Replicate this freeing order.
+	for o2 := range o.rdeps {
+		nvp.prependFreedLockedRecursive(o2)
+	}
+}
+
+// enqueueCleanup enqueues a cleanup function that will run after nvp.objsMu is
+// unlocked.
+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
+}
+
+func newSimpleObject() *simpleObject {
+	return &simpleObject{}
+}
+
+// Release implements objectImpl.Release.
+func (o *simpleObject) Release(ctx context.Context) {
+	// no-op
+}
+
+// rootClient is an objectImpl tracking a NV01_ROOT_CLIENT.
+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
+}
+
+func newRootClient(fd *frontendFD) *rootClient {
+	return &rootClient{
+		resources: make(map[nvgpu.Handle]*object),
+		fd:        fd,
+	}
+}
+
+// Release implements objectImpl.Release.
+func (o *rootClient) Release(ctx context.Context) {
+	delete(o.fd.clients, o.handle)
+}
+
+// osDescMem is an objectImpl tracking a NV01_MEMORY_SYSTEM_OS_DESCRIPTOR.
+type osDescMem struct {
+	object
+	pinnedRanges []mm.PinnedRange
+}
+
+// Release implements objectImpl.Release.
+func (o *osDescMem) Release(ctx context.Context) {
+	// Unpin pages (which takes MM locks) without holding nvproxy locks.
+	o.nvp.enqueueCleanup(func() {
+		mm.Unpin(o.pinnedRanges)
+		if ctx.IsLogging(log.Debug) {
+			total := uint64(0)
+			for _, pr := range o.pinnedRanges {
+				total += uint64(pr.Source.Length())
+			}
+			ctx.Debugf("nvproxy: unpinned %d bytes for released OS descriptor", total)
+		}
+	})
+}

pkg/sentry/devices/nvproxy/save_restore.go

@@ -17,18 +17,14 @@ package nvproxy
 import (
 	goContext "context"
 	"fmt"
-
-	"gvisor.dev/gvisor/pkg/abi/nvgpu"
-	"gvisor.dev/gvisor/pkg/context"
 )
 
 func (n *nvproxy) beforeSave() {
-	n.objsMu.Lock()
-	defer n.objsMu.Unlock()
-	for _, o := range n.objsLive {
-		o.Release(context.Background())
+	n.objsLock()
+	defer n.objsUnlock()
+	if len(n.clients) != 0 {
+		panic("can't save with live nvproxy clients")
 	}
-	n.objsLive = nil
 }
 
 func (n *nvproxy) afterLoad(goContext.Context) {
@@ -38,5 +34,4 @@ func (n *nvproxy) afterLoad(goContext.Context) {
 		panic(fmt.Sprintf("driver version %q not found in abis map", n.version))
 	}
 	n.abi = abiCons.cons()
-	n.objsLive = make(map[nvgpu.Handle]*object)
 }

pkg/sentry/devices/nvproxy/seccomp_filters.go

@@ -70,10 +70,6 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_CONTROL, nvgpu.SizeofNVOS54Parameters)),
 			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS21Parameters)),
-			},
 			seccomp.PerArg{
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS64Parameters)),

pkg/sentry/devices/nvproxy/version.go

@@ -261,7 +261,6 @@ func Init() {
 					nvgpu.NV2080_CTRL_CMD_RC_RELEASE_WATCHDOG_REQUESTS:                     rmControlSimple,
 					nvgpu.NV2080_CTRL_CMD_RC_SOFT_DISABLE_WATCHDOG:                         rmControlSimple,
 					nvgpu.NV2080_CTRL_CMD_TIMER_GET_GPU_CPU_TIME_CORRELATION_INFO:          rmControlSimple,
-					nvgpu.NV503C_CTRL_CMD_REGISTER_VA_SPACE:                                rmControlSimple,
 					nvgpu.NV503C_CTRL_CMD_REGISTER_VIDMEM:                                  rmControlSimple,
 					nvgpu.NV503C_CTRL_CMD_UNREGISTER_VIDMEM:                                rmControlSimple,
 					nvgpu.NV83DE_CTRL_CMD_DEBUG_SET_EXCEPTION_MASK:                         rmControlSimple,
@@ -284,13 +283,14 @@ func Init() {
 					nvgpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:                                ctrlDevGpuGetClasslist,
 					nvgpu.NV2080_CTRL_CMD_FIFO_DISABLE_CHANNELS:                            ctrlSubdevFIFODisableChannels,
 					nvgpu.NV2080_CTRL_CMD_GR_GET_INFO:                                      ctrlSubdevGRGetInfo,
+					nvgpu.NV503C_CTRL_CMD_REGISTER_VA_SPACE:                                ctrlRegisterVASpace,
 				},
 				allocationClass: map[nvgpu.ClassID]allocationClassHandler{
-					nvgpu.NV01_ROOT:                 rmAllocSimple[nvgpu.Handle],
-					nvgpu.NV01_ROOT_NON_PRIV:        rmAllocSimple[nvgpu.Handle],
+					nvgpu.NV01_ROOT:                 rmAllocRootClient,
+					nvgpu.NV01_ROOT_NON_PRIV:        rmAllocRootClient,
 					nvgpu.NV01_MEMORY_SYSTEM:        rmAllocSimple[nvgpu.NV_MEMORY_ALLOCATION_PARAMS],
 					nvgpu.NV01_MEMORY_LOCAL_USER:    rmAllocSimple[nvgpu.NV_MEMORY_ALLOCATION_PARAMS],
-					nvgpu.NV01_ROOT_CLIENT:          rmAllocSimple[nvgpu.Handle],
+					nvgpu.NV01_ROOT_CLIENT:          rmAllocRootClient,
 					nvgpu.NV01_EVENT_OS_EVENT:       rmAllocEventOSEvent,
 					nvgpu.NV2081_BINAPI:             rmAllocSimple[nvgpu.NV2081_ALLOC_PARAMETERS],
 					nvgpu.NV01_DEVICE_0:             rmAllocSimple[nvgpu.NV0080_ALLOC_PARAMETERS],
@@ -299,13 +299,13 @@ func Init() {
 					nvgpu.NV50_MEMORY_VIRTUAL:       rmAllocSimple[nvgpu.NV_MEMORY_ALLOCATION_PARAMS],
 					nvgpu.NV50_P2P:                  rmAllocSimple[nvgpu.NV503B_ALLOC_PARAMETERS],
 					nvgpu.NV50_THIRD_PARTY_P2P:      rmAllocSimple[nvgpu.NV503C_ALLOC_PARAMETERS],
-					nvgpu.GT200_DEBUGGER:            rmAllocSimple[nvgpu.NV83DE_ALLOC_PARAMETERS],
-					nvgpu.FERMI_CONTEXT_SHARE_A:     rmAllocSimple[nvgpu.NV_CTXSHARE_ALLOCATION_PARAMETERS],
+					nvgpu.GT200_DEBUGGER:            rmAllocSMDebuggerSession,
+					nvgpu.FERMI_CONTEXT_SHARE_A:     rmAllocContextShare,
 					nvgpu.FERMI_VASPACE_A:           rmAllocSimple[nvgpu.NV_VASPACE_ALLOCATION_PARAMETERS],
-					nvgpu.KEPLER_CHANNEL_GROUP_A:    rmAllocSimple[nvgpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS],
-					nvgpu.TURING_CHANNEL_GPFIFO_A:   rmAllocSimple[nvgpu.NV_CHANNEL_ALLOC_PARAMS],
-					nvgpu.AMPERE_CHANNEL_GPFIFO_A:   rmAllocSimple[nvgpu.NV_CHANNEL_ALLOC_PARAMS],
-					nvgpu.HOPPER_CHANNEL_GPFIFO_A:   rmAllocSimple[nvgpu.NV_CHANNEL_ALLOC_PARAMS],
+					nvgpu.KEPLER_CHANNEL_GROUP_A:    rmAllocChannelGroup,
+					nvgpu.TURING_CHANNEL_GPFIFO_A:   rmAllocChannel,
+					nvgpu.AMPERE_CHANNEL_GPFIFO_A:   rmAllocChannel,
+					nvgpu.HOPPER_CHANNEL_GPFIFO_A:   rmAllocChannel,
 					nvgpu.TURING_DMA_COPY_A:         rmAllocSimple[nvgpu.NVB0B5_ALLOCATION_PARAMETERS],
 					nvgpu.AMPERE_DMA_COPY_A:         rmAllocSimple[nvgpu.NVB0B5_ALLOCATION_PARAMETERS],
 					nvgpu.AMPERE_DMA_COPY_B:         rmAllocSimple[nvgpu.NVB0B5_ALLOCATION_PARAMETERS],