Disallow task creation after Kernel.WaitExited() returns.

Otherwise tasks can be created via the control server between when
Kernel.WaitExited() returns and when the control server is stopped, resulting
in task goroutines running when Kernel.Release() is called.

PiperOrigin-RevId: 658891833

pkg/sentry/kernel/kernel.go

@@ -1370,9 +1370,15 @@ func (k *Kernel) decRunningTasks() {
 	// active without an expensive transition.
 }
 
-// WaitExited blocks until all tasks in k have exited.
+// WaitExited blocks until all tasks in k have exited. No tasks can be created
+// after WaitExited returns.
 func (k *Kernel) WaitExited() {
-	k.tasks.liveGoroutines.Wait()
+	k.tasks.mu.Lock()
+	defer k.tasks.mu.Unlock()
+	k.tasks.noNewTasksIfZeroLive = true
+	for k.tasks.liveTasks != 0 {
+		k.tasks.zeroLiveTasksCond.Wait()
+	}
 }
 
 // Kill requests that all tasks in k immediately exit as if group exiting with

pkg/sentry/kernel/kernel_state.go

@@ -20,6 +20,11 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
+// afterLoad is invoked by stateify.
+func (ts *TaskSet) afterLoad(_ context.Context) {
+	ts.zeroLiveTasksCond.L = &ts.mu
+}
+
 // saveDanglingEndpoints is invoked by stateify.
 func (k *Kernel) saveDanglingEndpoints() []tcpip.Endpoint {
 	return tcpip.GetDanglingEndpoints()

pkg/sentry/kernel/task_run.go

@@ -100,10 +100,17 @@ func (t *Task) run(threadID uintptr) {
 			t.accountTaskGoroutineEnter(TaskGoroutineNonexistent)
 			t.goroutineStopped.Done()
 			t.tg.liveGoroutines.Done()
-			t.tg.pidns.owner.liveGoroutines.Done()
-			t.tg.pidns.owner.runningGoroutines.Done()
 			t.p.Release()
 
+			ts := t.tg.pidns.owner
+			ts.mu.Lock()
+			ts.liveTasks--
+			if ts.liveTasks == 0 {
+				ts.zeroLiveTasksCond.Broadcast()
+			}
+			ts.mu.Unlock()
+			ts.runningGoroutines.Done()
+
 			// Deferring this store triggers a false positive in the race
 			// detector (https://github.com/golang/go/issues/42599).
 			t.goid.Store(0)

pkg/sentry/kernel/task_start.go

@@ -226,12 +226,20 @@ func (ts *TaskSet) newTask(ctx context.Context, cfg *TaskConfig) (*Task, error)
 		// we're in uncharted territory and can return whatever we want.
 		return nil, linuxerr.EINTR
 	}
+	if ts.liveTasks == 0 && ts.noNewTasksIfZeroLive {
+		// Since liveTasks == 0, our caller cannot be a task goroutine invoking
+		// a syscall, so it's safe to return a non-errno error that is more
+		// explanatory.
+		return nil, fmt.Errorf("task creation disabled after Kernel.WaitExited() may have returned")
+	}
 	if err := ts.assignTIDsLocked(t); err != nil {
 		return nil, err
 	}
 	// Below this point, newTask is expected not to fail (there is no rollback
 	// of assignTIDsLocked or any of the following).
 
+	ts.liveTasks++
+
 	// Logging on t's behalf will panic if t.logPrefix hasn't been
 	// initialized. This is the earliest point at which we can do so
 	// (since t now has thread IDs).
@@ -383,7 +391,6 @@ func (t *Task) Start(tid ThreadID) {
 	}
 	t.goroutineStopped.Add(1)
 	t.tg.liveGoroutines.Add(1)
-	t.tg.pidns.owner.liveGoroutines.Add(1)
 	t.tg.pidns.owner.runningGoroutines.Add(1)
 
 	// Task is now running in system mode.

pkg/sentry/kernel/threads.go

@@ -76,12 +76,17 @@ type TaskSet struct {
 	// always reset to zero after restore.
 	stopCount int32 `state:"nosave"`
 
-	// liveGoroutines is the number of non-exited task goroutines in the
-	// TaskSet.
-	//
-	// liveGoroutines is not saved; it is reset as task goroutines are
-	// restarted by Task.Start.
-	liveGoroutines sync.WaitGroup `state:"nosave"`
+	// liveTasks is the number of tasks in the TaskSet whose goroutines have
+	// not exited. liveTasks is protected by mu.
+	liveTasks uint32
+
+	// If noNewTasksIfZeroLive is true and liveTasks is zero, calls to
+	// Kernel.NewTask() will fail. noNewTasksIfZeroLive is protected by mu.
+	noNewTasksIfZeroLive bool
+
+	// zeroLiveTasksCond is broadcast when liveTasks transitions from non-zero
+	// to zero.
+	zeroLiveTasksCond sync.Cond `state:"nosave"`
 
 	// runningGoroutines is the number of running task goroutines in the
 	// TaskSet.
@@ -102,6 +107,7 @@ type TaskSet struct {
 // newTaskSet returns a new, empty TaskSet.
 func newTaskSet(pidns *PIDNamespace) *TaskSet {
 	ts := &TaskSet{Root: pidns}
+	ts.zeroLiveTasksCond.L = &ts.mu
 	pidns.owner = ts
 	return ts
 }