From 4298980325950f64c1dca0e6a11902a27a48cf09 Mon Sep 17 00:00:00 2001
From: Jamie Liu <jamieliu@google.com>
Date: Fri, 2 Aug 2024 13:37:22 -0700
Subject: [PATCH] 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       | 10 ++++++++--
 pkg/sentry/kernel/kernel_state.go |  5 +++++
 pkg/sentry/kernel/task_run.go     | 11 +++++++++--
 pkg/sentry/kernel/task_start.go   |  9 ++++++++-
 pkg/sentry/kernel/threads.go      | 18 ++++++++++++------
 5 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 2f4f6f410..964730384 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/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
diff --git a/pkg/sentry/kernel/kernel_state.go b/pkg/sentry/kernel/kernel_state.go
index 65247f3a7..a47d164e9 100644
--- a/pkg/sentry/kernel/kernel_state.go
+++ b/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()
diff --git a/pkg/sentry/kernel/task_run.go b/pkg/sentry/kernel/task_run.go
index b39fe1998..10b616f02 100644
--- a/pkg/sentry/kernel/task_run.go
+++ b/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)
diff --git a/pkg/sentry/kernel/task_start.go b/pkg/sentry/kernel/task_start.go
index 5b6b6e1c6..ded77133f 100644
--- a/pkg/sentry/kernel/task_start.go
+++ b/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.
diff --git a/pkg/sentry/kernel/threads.go b/pkg/sentry/kernel/threads.go
index f25ca7076..6c7d7f782 100644
--- a/pkg/sentry/kernel/threads.go
+++ b/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
 }