From 96e773f1f13a0a8e4d7bd2ce6d8b4e847a76dc7a Mon Sep 17 00:00:00 2001
From: AdaAibaby <shaolila@buaa.edu.cn>
Date: Thu, 13 Aug 2026 16:50:39 +0800
Subject: [PATCH] fix(orchestrator): close TOCTOU race in
 Cleanup.Add/AddPriority
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Cleanup.run() set hasRun before acquiring mu, while Add/AddPriority
read hasRun before acquiring mu. This created a window where a goroutine
could observe hasRun==false, lose the race to run(), and then append a
cleanup function into an already-drained slice — silently discarding it.

Fix:
- Move hasRun.Store(true) inside mu in run(), so the flag is only set
  while the slice is still accessible under the lock.
- Add a double-check of hasRun after acquiring mu in Add/AddPriority,
  so any function that loses the race is executed immediately rather
  than silently dropped.

The optimistic fast-path Load() before the lock is kept for the common
post-run case (avoids lock contention when the sandbox is fully torn
down), but is no longer the only gate.

Fixes https://github.com/e2b-dev/infra/issues/3557
---
 packages/orchestrator/pkg/sandbox/cleanup.go | 36 +++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/packages/orchestrator/pkg/sandbox/cleanup.go b/packages/orchestrator/pkg/sandbox/cleanup.go
index 2de71f2db9..d01a9abd9f 100644
--- a/packages/orchestrator/pkg/sandbox/cleanup.go
+++ b/packages/orchestrator/pkg/sandbox/cleanup.go
@@ -37,7 +37,7 @@ func (c *Cleanup) AddNoContext(ctx context.Context, f func() error) {
 }
 
 func (c *Cleanup) Add(ctx context.Context, f func(ctx context.Context) error) {
-	if c.hasRun.Load() == true {
+	if c.hasRun.Load() {
 		err := f(context.WithoutCancel(ctx))
 		if err != nil {
 			logger.L().Error(ctx, "failed to run function after cleanup has run", zap.Error(err))
@@ -47,13 +47,25 @@ func (c *Cleanup) Add(ctx context.Context, f func(ctx context.Context) error) {
 	}
 
 	c.mu.Lock()
+	// Double-check under lock: run() may have completed between the Load above
+	// and acquiring mu, silently dropping f into an already-drained slice.
+	if c.hasRun.Load() {
+		c.mu.Unlock()
+
+		err := f(context.WithoutCancel(ctx))
+		if err != nil {
+			logger.L().Error(ctx, "failed to run function after cleanup has run", zap.Error(err))
+		}
+
+		return
+	}
 	defer c.mu.Unlock()
 
 	c.cleanup = append(c.cleanup, f)
 }
 
 func (c *Cleanup) AddPriority(ctx context.Context, f func(ctx context.Context) error) {
-	if c.hasRun.Load() == true {
+	if c.hasRun.Load() {
 		err := f(context.WithoutCancel(ctx))
 		if err != nil {
 			logger.L().Error(ctx, "failed to run priority function after cleanup has run", zap.Error(err))
@@ -63,6 +75,18 @@ func (c *Cleanup) AddPriority(ctx context.Context, f func(ctx context.Context) e
 	}
 
 	c.mu.Lock()
+	// Double-check under lock: run() may have completed between the Load above
+	// and acquiring mu, silently dropping f into an already-drained slice.
+	if c.hasRun.Load() {
+		c.mu.Unlock()
+
+		err := f(context.WithoutCancel(ctx))
+		if err != nil {
+			logger.L().Error(ctx, "failed to run priority function after cleanup has run", zap.Error(err))
+		}
+
+		return
+	}
 	defer c.mu.Unlock()
 
 	c.priorityCleanup = append(c.priorityCleanup, f)
@@ -77,11 +101,15 @@ func (c *Cleanup) Run(ctx context.Context) error {
 }
 
 func (c *Cleanup) run(ctx context.Context) {
-	c.hasRun.Store(true)
-
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	// Set hasRun inside the lock so Add/AddPriority's double-check (also under
+	// mu) sees the updated value before the slice is drained. Without this,
+	// Add can observe hasRun==false, lose the race to run(), and append f to an
+	// already-drained slice where it will never be executed.
+	c.hasRun.Store(true)
+
 	var errs []error
 
 	for _, cleanup := range slices.Backward(c.priorityCleanup) {