orch: fix uffd prefault/close race condition

packages/orchestrator/pkg/sandbox/uffd/testutils/testharness/barriers.go

@@ -15,6 +15,9 @@ const (
 	BeforeRLock Point = iota
 	// BeforeFaultPage parks after settleRequests.RLock, before UFFDIO_COPY.
 	BeforeFaultPage
+	// BeforePrefaultRLock parks inside Prefault(), before settleRequests.RLock.
+	// Value must stay in sync with faultPhaseBeforePrefaultRLock in the parent package.
+	BeforePrefaultRLock
 )
 
 // Registry is the child-side barrier store consulted by the per-fault hook.

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/close_race_test.go

@@ -0,0 +1,121 @@
+//go:build linux
+
+package userfaultfd
+
+import (
+	"context"
+	"syscall"
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/memory"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+// TestPrefaultConcurrentWithClose is a regression test for the race between
+// Prefault() and Close() that produced ENOTTY/EBADF in production:
+//
+//	Prefault                           Close
+//	 (about to acquire RLock)
+//	                                    acquires Lock
+//	                                    closed = true
+//	                                    fd.close()          ← fd freed/recycled
+//	                                    releases Lock
+//	 acquires RLock
+//	 sees closed == true → return nil  ✓
+//
+// Without the fix, Prefault had no closed check and would call UFFDIO_COPY
+// on the closed (and potentially recycled) fd, returning EBADF or ENOTTY.
+//
+// The test uses faultPhaseBeforePrefaultRLock to park Prefault before the
+// RLock acquisition, lets Close() run to completion, then releases the park
+// and asserts Prefault returns nil.
+func TestPrefaultConcurrentWithClose(t *testing.T) {
+	t.Parallel()
+
+	withRaceContext(t, func(ctx context.Context) {
+		// Create a real userfaultfd so Close() can close a valid fd.  We do
+		// NOT call configureApi/register because the test short-circuits via
+		// the closed flag before any UFFDIO_COPY is attempted.
+		uffdFd, err := newFd(syscall.O_CLOEXEC | syscall.O_NONBLOCK)
+		require.NoError(t, err)
+
+		// Minimal page-sized backing store (content doesn't matter; Prefault
+		// won't reach ReadAt in the short-circuit path).
+		pageData := make([]byte, header.PageSize)
+		src := NewMemorySlicer(pageData, int64(header.PageSize))
+
+		// A single-region mapping anchored at a real heap address.
+		// GetHostVirtAddr is called after the closed check, so the address
+		// only needs to satisfy NewUserfaultfdFromFd's region validation.
+		regionBuf := make([]byte, header.PageSize)
+		mapping := memory.NewMapping([]memory.Region{{
+			BaseHostVirtAddr: uintptr(unsafe.Pointer(&regionBuf[0])),
+			Size:             uintptr(header.PageSize),
+			Offset:           0,
+			PageSize:         uintptr(header.PageSize),
+		}})
+
+		log, err := logger.NewDevelopmentLogger()
+		require.NoError(t, err)
+
+		u, err := NewUserfaultfdFromFd(uintptr(uffdFd), src, mapping, log)
+		require.NoError(t, err)
+
+		// Barrier channels: park Prefault at faultPhaseBeforePrefaultRLock so
+		// Close() can run to completion before the RLock is acquired.
+		arrived := make(chan struct{})
+		release := make(chan struct{})
+
+		u.SetTestFaultHook(func(_ uintptr, phase faultPhase) {
+			if phase != faultPhaseBeforePrefaultRLock {
+				return
+			}
+			// Signal arrival exactly once (guard against re-entry).
+			select {
+			case <-arrived:
+			default:
+				close(arrived)
+			}
+			// Wait for the test to release us.
+			select {
+			case <-release:
+			case <-ctx.Done():
+			}
+		})
+
+		prefaultErrs := make(chan error, 1)
+		go func() {
+			prefaultErrs <- u.Prefault(ctx, 0, pageData)
+		}()
+
+		// Wait for Prefault to park at the pre-RLock hook.
+		select {
+		case <-arrived:
+		case <-ctx.Done():
+			t.Fatal("Prefault goroutine did not reach pre-RLock hook within budget")
+		}
+
+		// Close the uffd while Prefault is parked.
+		// Pre-fix: fd gets closed and may be recycled; Prefault later calls
+		//   UFFDIO_COPY on it and returns EBADF or ENOTTY.
+		// Post-fix: Close() holds Lock(); Prefault acquires RLock, sees
+		//   closed==true, returns nil without touching the fd.
+		require.NoError(t, u.Close())
+
+		// Release the parked Prefault goroutine.
+		close(release)
+
+		select {
+		case err := <-prefaultErrs:
+			require.NoError(t, err,
+				"Prefault must return nil after concurrent Close() — "+
+					"a non-nil error means UFFDIO_COPY was attempted on the closed fd (EBADF/ENOTTY regression)")
+		case <-ctx.Done():
+			t.Fatal("Prefault goroutine did not complete after hook release")
+		}
+	})
+}

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/prefault.go

@@ -13,9 +13,23 @@ import (
 // Prefault proactively copies a page to guest memory at the given offset
 // to speed up sandbox starts. EEXIST (already mapped) is handled gracefully.
 func (u *Userfaultfd) Prefault(ctx context.Context, offset int64, data []byte) error {
+	// Test hook: fires before settleRequests.RLock so that a test can park
+	// the goroutine here, call Close() concurrently, then release and observe
+	// that the closed check below returns nil without calling UFFDIO_COPY.
+	if h := u.testFaultHook.Load(); h != nil {
+		(*h)(0, faultPhaseBeforePrefaultRLock)
+	}
+
 	u.settleRequests.RLock()
 	defer u.settleRequests.RUnlock()
 
+	// Close() sets closed under Lock(). Seeing it here under RLock means the
+	// fd is already closed and its number may have been recycled by the OS —
+	// skip silently instead of calling UFFDIO_COPY and getting EBADF/ENOTTY.
+	if u.closed {
+		return nil
+	}
+
 	ctx, span := tracer.Start(ctx, "prefault page")
 	defer span.End()
 

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go

@@ -99,6 +99,11 @@ type Userfaultfd struct {
 	// testFaultHook is set only by SetTestFaultHook in test builds.
 	testFaultHook atomic.Pointer[func(uintptr, faultPhase)]
 
+	// closed is set by Close() under settleRequests.Lock(). Prefault() checks
+	// it under settleRequests.RLock() so it never calls UFFDIO_COPY on a fd
+	// that has already been closed (and potentially recycled by the OS).
+	closed bool
+
 	logger logger.Logger
 }
 
@@ -108,6 +113,11 @@ type faultPhase uint8
 const (
 	faultPhaseBeforeRLock faultPhase = iota
 	faultPhaseBeforeFaultPage
+	// faultPhaseBeforePrefaultRLock fires inside Prefault(), before acquiring
+	// settleRequests.RLock. Used by TestPrefaultConcurrentWithClose to park
+	// Prefault here, call Close() concurrently, and verify the closed flag is
+	// checked after the RLock is acquired.
+	faultPhaseBeforePrefaultRLock
 )
 
 // faultOutcome is the terminal classification of a faultPage call.
@@ -655,6 +665,22 @@ func (u *Userfaultfd) PageSize() int64 {
 }
 
 func (u *Userfaultfd) Close() error {
+	// Hold the write lock for the entire close sequence so that:
+	//   (a) any Prefault() caller currently holding RLock finishes its
+	//       UFFDIO_COPY before the fd number is freed and potentially
+	//       recycled by the OS;
+	//   (b) Close() is idempotent — a second call sees closed==true and
+	//       returns immediately without touching already-freed fds.
+	// In production Serve() drains all workers (u.wg.Wait) before
+	// returning, so this Lock() is always uncontended when Close() fires.
+	u.settleRequests.Lock()
+	defer u.settleRequests.Unlock()
+
+	if u.closed {
+		return nil
+	}
+	u.closed = true
+
 	syscall.Close(u.wakeupPipe[0])
 	syscall.Close(u.wakeupPipe[1])