Merge pull request #12055 from tianyuzhou95:albert/refcount

PiperOrigin-RevId: 800675291

Author: gvisor-bot

pkg/sentry/control/lifecycle.go

@@ -21,6 +21,7 @@ import (
 
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/eventchannel"
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
@@ -288,6 +289,8 @@ func (l *Lifecycle) StartContainer(args *StartContainerArgs, _ *uint32) error {
 
 	initArgs.MountNamespace = mntns
 	initArgs.MountNamespace.IncRef()
+	mntnsCu := cleanup.Make(func() { initArgs.MountNamespace.DecRef(ctx) })
+	defer mntnsCu.Clean()
 
 	if args.ResolveBinaryPath {
 		resolved, err := user.ResolveExecutablePath(ctx, &initArgs)
@@ -327,6 +330,7 @@ func (l *Lifecycle) StartContainer(args *StartContainerArgs, _ *uint32) error {
 	}
 	initArgs.InitialCgroups = initialCgroups
 
+	mntnsCu.Release() // mntns ref is transferred to Kernel.CreateProcess()
 	tg, _, err := l.Kernel.CreateProcess(initArgs)
 	if err != nil {
 		return err

pkg/sentry/control/proc.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/fdimport"
@@ -217,13 +218,18 @@ func (proc *Proc) execAsync(args *ExecArgs) (*kernel.ThreadGroup, kernel.ThreadI
 		PIDNamespace:         pidns,
 		Origin:               kernel.OriginExec,
 	}
-	if initArgs.MountNamespace != nil {
-		// initArgs must hold a reference on MountNamespace, which will
-		// be donated to the new process in CreateProcess.
-		initArgs.MountNamespace.IncRef()
-	}
 	ctx := initArgs.NewContext(proc.Kernel)
 
+	if initArgs.MountNamespace == nil {
+		// Set initArgs so that 'ctx' returns the namespace.
+		initArgs.MountNamespace = proc.Kernel.GlobalInit().Leader().MountNamespace()
+	}
+	// initArgs must hold a reference on MountNamespace, which will
+	// be donated to the new process in CreateProcess.
+	initArgs.MountNamespace.IncRef()
+	mntnsCu := cleanup.Make(func() { initArgs.MountNamespace.DecRef(ctx) })
+	defer mntnsCu.Clean()
+
 	// Import file descriptors.
 	var fdTable *kernel.FDTable
 	if args.FDTable != nil {
@@ -235,15 +241,6 @@ func (proc *Proc) execAsync(args *ExecArgs) (*kernel.ThreadGroup, kernel.ThreadI
 	}
 	initArgs.FDTable = fdTable
 
-	// Get the full path to the filename from the PATH env variable.
-	if initArgs.MountNamespace == nil {
-		// Set initArgs so that 'ctx' returns the namespace.
-		//
-		// Add a reference to the namespace, which is transferred to the new process.
-		initArgs.MountNamespace = proc.Kernel.GlobalInit().Leader().MountNamespace()
-		initArgs.MountNamespace.IncRef()
-	}
-
 	fdMap, execFD, err := args.unpackFiles()
 	if err != nil {
 		return nil, 0, nil, fmt.Errorf("creating fd map: %w", err)
@@ -311,6 +308,7 @@ func (proc *Proc) execAsync(args *ExecArgs) (*kernel.ThreadGroup, kernel.ThreadI
 		initArgs.InitialCgroups = initialCgrps
 	}
 
+	mntnsCu.Release() // mntns ref is transferred to Kernel.CreateProcess()
 	tg, tid, err := proc.Kernel.CreateProcess(initArgs)
 	if err != nil {
 		return nil, 0, nil, err

pkg/sentry/kernel/kernel.go

@@ -1038,6 +1038,9 @@ func (ctx *createProcessContext) getMemoryCgroupID() uint32 {
 //
 // CreateProcess has no analogue in Linux; it is used to create the initial
 // application task, as well as processes started by the control server.
+//
+// Precondition: Caller must take a ref on args.MountNamespace, which is
+// transferred to CreateProcess.
 func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID, error) {
 	k.extMu.Lock()
 	defer k.extMu.Unlock()
@@ -1056,6 +1059,8 @@ func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID,
 	// Get the root directory from the MountNamespace.
 	root := mntns.Root(ctx)
 	defer root.DecRef(ctx)
+	refcountCu := cleanup.Make(func() { mntns.DecRef(ctx) })
+	defer refcountCu.Clean()
 
 	// Grab the working directory.
 	wd := root // Default.
@@ -1081,6 +1086,7 @@ func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID,
 		defer wd.DecRef(ctx)
 	}
 	fsContext := NewFSContext(root, wd, args.Umask)
+	refcountCu.Add(func() { fsContext.DecRef(ctx) })
 
 	tg := k.NewThreadGroup(args.PIDNamespace, NewSignalHandlers(), linux.SIGCHLD, args.Limits)
 	cu := cleanup.Make(func() {
@@ -1156,6 +1162,7 @@ func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID,
 	config.UTSNamespace.IncRef()
 	config.IPCNamespace.IncRef()
 	config.NetworkNamespace.IncRef()
+	refcountCu.Release() // refs(mntns, fsContext) are transferred to NewTask()
 	t, err := k.tasks.NewTask(ctx, config)
 	if err != nil {
 		return nil, 0, err