feat(uffd): handle UFFD_EVENT_REMOVE; track per-page state; race-safe COPY

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

@@ -289,38 +289,42 @@ func TestAsyncWriteProtection(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			h, err := configureCrossProcessTest(t.Context(), t, testConfig{
+			runMatrix(t, testConfig{
 				pagesize:      tt.pagesize,
 				numberOfPages: tt.numberOfPages,
 				alwaysWP:      tt.alwaysWP,
-			})
-			require.NoError(t, err)
+			}, func(t *testing.T, cfg testConfig) {
+				t.Helper()
+
+				h, err := configureCrossProcessTest(t.Context(), t, cfg)
+				require.NoError(t, err)
 
-			h.executeAll(t, tt.operations)
+				h.executeAll(t, tt.operations)
 
-			pagemap, err := testutils.NewPagemapReader()
-			require.NoError(t, err)
-			defer pagemap.Close()
+				pagemap, err := testutils.NewPagemapReader()
+				require.NoError(t, err)
+				defer pagemap.Close()
 
-			memStart := uintptr(unsafe.Pointer(&(*h.memoryArea)[0]))
+				memStart := uintptr(unsafe.Pointer(&(*h.memoryArea)[0]))
 
-			for _, p := range tt.expectedDirty {
-				addr := memStart + uintptr(p)*uintptr(tt.pagesize)
-				entry, err := pagemap.ReadEntry(addr)
-				require.NoError(t, err, "pagemap read for dirty page %d", p)
+				for _, p := range tt.expectedDirty {
+					addr := memStart + uintptr(p)*uintptr(tt.pagesize)
+					entry, err := pagemap.ReadEntry(addr)
+					require.NoError(t, err, "pagemap read for dirty page %d", p)
 
-				assert.True(t, entry.IsPresent(), "dirty page %d should be present", p)
-				assert.False(t, entry.IsWriteProtected(), "dirty page %d should have WP cleared", p)
-			}
+					assert.True(t, entry.IsPresent(), "dirty page %d should be present", p)
+					assert.False(t, entry.IsWriteProtected(), "dirty page %d should have WP cleared", p)
+				}
 
-			for _, p := range tt.expectedClean {
-				addr := memStart + uintptr(p)*uintptr(tt.pagesize)
-				entry, err := pagemap.ReadEntry(addr)
-				require.NoError(t, err, "pagemap read for clean page %d", p)
+				for _, p := range tt.expectedClean {
+					addr := memStart + uintptr(p)*uintptr(tt.pagesize)
+					entry, err := pagemap.ReadEntry(addr)
+					require.NoError(t, err, "pagemap read for clean page %d", p)
 
-				assert.True(t, entry.IsPresent(), "clean page %d should be present", p)
-				assert.True(t, entry.IsWriteProtected(), "clean page %d should have WP set", p)
-			}
+					assert.True(t, entry.IsPresent(), "clean page %d should be present", p)
+					assert.True(t, entry.IsWriteProtected(), "clean page %d should have WP set", p)
+				}
+			})
 		})
 	}
 }

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

@@ -0,0 +1,25 @@
+package userfaultfd
+
+import "sync"
+
+// deferredFaults collects pagefaults that returned EAGAIN so they get
+// retried on the next poll iteration. Safe for concurrent push.
+type deferredFaults struct {
+	mu sync.Mutex
+	pf []*UffdPagefault
+}
+
+func (d *deferredFaults) push(pf *UffdPagefault) {
+	d.mu.Lock()
+	d.pf = append(d.pf, pf)
+	d.mu.Unlock()
+}
+
+func (d *deferredFaults) drain() []*UffdPagefault {
+	d.mu.Lock()
+	out := d.pf
+	d.pf = nil
+	d.mu.Unlock()
+
+	return out
+}

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

@@ -146,18 +146,30 @@ func getPagefaultAddress(pagefault *UffdPagefault) uintptr {
 // Fd is a helper type that wraps uffd fd.
 type Fd uintptr
 
-// mode: UFFDIO_COPY_MODE_WP
-// When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page
+// copy requires UFFDIO_COPY_MODE_WP when both MISSING and WP tracking are active.
 func (f Fd) copy(addr, pagesize uintptr, data []byte, mode CULong) error {
 	cpy := newUffdioCopy(data, CULong(addr)&^CULong(pagesize-1), CULong(pagesize), mode, 0)
 
 	if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(f), UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 {
 		return errno
 	}
 
-	// Check if the copied size matches the requested pagesize
-	if cpy.copy != CLong(pagesize) {
-		return fmt.Errorf("UFFDIO_COPY copied %d bytes, expected %d", cpy.copy, pagesize)
+	return classifyCopyResult(int64(cpy.copy), int64(pagesize))
+}
+
+// classifyCopyResult turns the UFFDIO_COPY cpy.copy field into a Go error.
+// The kernel encodes a negated errno on failure (most commonly -EAGAIN when
+// mmap_changing is set), and a short positive count when the copy was
+// preempted mid-page (e.g. hugetlb). Both partial outcomes surface as EAGAIN;
+// the caller defers the fault for retry on the next poll iteration (via the
+// deferred queue + wakeup pipe), the kernel does not auto-redeliver.
+func classifyCopyResult(bytesCopied, pagesize int64) error {
+	if bytesCopied < 0 {
+		return syscall.Errno(-bytesCopied)
+	}
+
+	if bytesCopied != pagesize {
+		return syscall.EAGAIN
 	}
 
 	return nil
@@ -170,7 +182,6 @@ func (f Fd) zero(addr, pagesize uintptr, mode CULong) error {
 		return errno
 	}
 
-	// Check if the bytes actually zeroed out by the kernel match the page size
 	if zero.zeropage != CLong(pagesize) {
 		return fmt.Errorf("UFFDIO_ZEROPAGE copied %d bytes, expected %d", zero.zeropage, pagesize)
 	}

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

@@ -8,8 +8,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
-// Used for testing.
-// flags: syscall.O_CLOEXEC|syscall.O_NONBLOCK
+// Used for testing (flags: syscall.O_CLOEXEC|syscall.O_NONBLOCK).
 func newFd(flags uintptr) (Fd, error) {
 	uffd, _, errno := syscall.Syscall(NR_userfaultfd, flags, 0, 0)
 	if errno != 0 {
@@ -19,18 +18,20 @@ func newFd(flags uintptr) (Fd, error) {
 	return Fd(uffd), nil
 }
 
-// Used for testing
-// features: UFFD_FEATURE_MISSING_HUGETLBFS
-// This is already called by the FC
-func configureApi(f Fd, pagesize uint64) error {
+// configureApi is used for testing. The caller (FC in production) sets
+// UFFD_FEATURE_MISSING_HUGETLBFS only when hugepages are in use.
+func configureApi(f Fd, pagesize uint64, removeEnabled bool) error {
 	var features CULong
 
-	// Only set the hugepage feature if we're using hugepages
+	// Only set the hugepage feature if we're using hugepages.
 	if pagesize == header.HugepageSize {
 		features |= UFFD_FEATURE_MISSING_HUGETLBFS
 	}
 
 	features |= UFFD_FEATURE_WP_ASYNC
+	if removeEnabled {
+		features |= UFFD_FEATURE_EVENT_REMOVE
+	}
 
 	api := newUffdioAPI(UFFD_API, features)
 	ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(f), UFFDIO_API, uintptr(unsafe.Pointer(&api)))
@@ -42,9 +43,9 @@ func configureApi(f Fd, pagesize uint64) error {
 }
 
 // unregister tears down the UFFD registration over [addr, addr+size).
-// Used in test cleanup so that any in-flight REMOVE events the kernel
-// may have queued (once UFFD_FEATURE_EVENT_REMOVE is enabled in a
-// follow-up) don't keep munmap blocked on un-acked events.
+// Used in test cleanup so in-flight REMOVE events queued by the kernel
+// (configureApi enables UFFD_FEATURE_EVENT_REMOVE when removeEnabled)
+// don't keep munmap blocked on un-acked events.
 func unregister(f Fd, addr uintptr, size uint64) error {
 	r := newUffdioRange(CULong(addr), CULong(size))
 
@@ -56,9 +57,8 @@ func unregister(f Fd, addr uintptr, size uint64) error {
 	return nil
 }
 
-// mode: UFFDIO_REGISTER_MODE_WP|UFFDIO_REGISTER_MODE_MISSING
-// This is already called by the FC, but only with the UFFDIO_REGISTER_MODE_MISSING
-// We need to call it with UFFDIO_REGISTER_MODE_WP when we use both missing and wp
+// register is used for testing. FC uses UFFDIO_REGISTER_MODE_MISSING; we add
+// UFFDIO_REGISTER_MODE_WP when both missing and write-protection are needed.
 func register(f Fd, addr uintptr, size uint64, mode CULong) error {
 	register := newUffdioRegister(CULong(addr), CULong(size), mode)
 

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

@@ -64,6 +64,10 @@ func configureCrossProcessTest(ctx context.Context, t *testing.T, tt testConfig)
 
 	data := RandomPages(tt.pagesize, tt.numberOfPages)
 
+	if tt.sourcePatcher != nil {
+		tt.sourcePatcher(data.Content())
+	}
+
 	size, err := data.Size()
 	require.NoError(t, err)
 
@@ -74,7 +78,7 @@ func configureCrossProcessTest(ctx context.Context, t *testing.T, tt testConfig)
 	require.NoError(t, err)
 	t.Cleanup(func() { uffdFd.close() })
 
-	require.NoError(t, configureApi(uffdFd, tt.pagesize))
+	require.NoError(t, configureApi(uffdFd, tt.pagesize, tt.removeEnabled))
 	require.NoError(t, register(uffdFd, memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP))
 	t.Cleanup(func() {
 		// Unregister before close (LIFO): a future test enabling