From 0b59173cff801c54d8056a48ae2b2a0dff0dd898 Mon Sep 17 00:00:00 2001
From: Jamie Liu <jamieliu@google.com>
Date: Fri, 18 Oct 2024 14:14:03 -0700
Subject: [PATCH] kernel: only lock TaskSet.mu in Task.unstopVforkParent() if
 necessary

PiperOrigin-RevId: 687421394
---
 pkg/sentry/kernel/task.go       | 10 +++++++++-
 pkg/sentry/kernel/task_clone.go | 22 ++++++++++++++--------
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go
index 35b5fc63e..0ebc15759 100644
--- a/pkg/sentry/kernel/task.go
+++ b/pkg/sentry/kernel/task.go
@@ -286,7 +286,7 @@ type Task struct {
 	// this TaskImage is released.
 	//
 	// vforkParent is protected by the TaskSet mutex.
-	vforkParent *Task
+	vforkParent atomic.Pointer[Task] `state:".(*Task)"`
 
 	// exitState is the task's progress through the exit path.
 	//
@@ -631,6 +631,14 @@ var (
 		})
 )
 
+func (t *Task) saveVforkParent() *Task {
+	return t.vforkParent.Load()
+}
+
+func (t *Task) loadVforkParent(_ gocontext.Context, vforkParent *Task) {
+	t.vforkParent.Store(vforkParent)
+}
+
 func (t *Task) savePtraceTracer() *Task {
 	return t.ptraceTracer.Load()
 }
diff --git a/pkg/sentry/kernel/task_clone.go b/pkg/sentry/kernel/task_clone.go
index 3a644c57f..1837b3e32 100644
--- a/pkg/sentry/kernel/task_clone.go
+++ b/pkg/sentry/kernel/task_clone.go
@@ -332,7 +332,7 @@ func (t *Task) Clone(args *linux.CloneArgs) (ThreadID, *SyscallControl, error) {
 		nt.seccomp.Store(nil)
 	}
 	if args.Flags&linux.CLONE_VFORK != 0 {
-		nt.vforkParent = t
+		nt.vforkParent.Store(t)
 	}
 
 	if args.Flags&linux.CLONE_CHILD_CLEARTID != 0 {
@@ -396,30 +396,36 @@ func getCloneSeccheckInfo(t, nt *Task, flags uint64) (seccheck.FieldSet, *pb.Clo
 //
 // Preconditions: The caller must be running on t's task goroutine.
 func (t *Task) maybeBeginVforkStop(child *Task) {
+	if child.vforkParent.Load() != t {
+		return
+	}
 	t.tg.pidns.owner.mu.Lock()
 	defer t.tg.pidns.owner.mu.Unlock()
 	t.tg.signalHandlers.mu.Lock()
 	defer t.tg.signalHandlers.mu.Unlock()
-	if t.killedLocked() {
-		child.vforkParent = nil
+	if child.vforkParent.Load() != t {
 		return
 	}
-	if child.vforkParent == t {
-		t.beginInternalStopLocked((*vforkStop)(nil))
+	if t.killedLocked() {
+		child.vforkParent.Store(nil)
+		return
 	}
+	t.beginInternalStopLocked((*vforkStop)(nil))
 }
 
 func (t *Task) unstopVforkParent() {
+	if t.vforkParent.Load() == nil {
+		return
+	}
 	t.tg.pidns.owner.mu.Lock()
 	defer t.tg.pidns.owner.mu.Unlock()
-	if p := t.vforkParent; p != nil {
+	if p := t.vforkParent.Load(); p != nil {
+		t.vforkParent.Store(nil)
 		p.tg.signalHandlers.mu.Lock()
 		defer p.tg.signalHandlers.mu.Unlock()
 		if _, ok := p.stop.(*vforkStop); ok {
 			p.endInternalStopLocked()
 		}
-		// Parent no longer needs to be unstopped.
-		t.vforkParent = nil
 	}
 }