Add UFFD write protection

packages/orchestrator/internal/sandbox/block/device.go

@@ -16,6 +16,7 @@ func (BytesNotAvailableError) Error() string {
 
 type Slicer interface {
 	Slice(ctx context.Context, off, length int64) ([]byte, error)
+	BlockSize() int64
 }
 
 type ReadonlyDevice interface {

packages/orchestrator/internal/sandbox/block/tracker.go

@@ -1,66 +1,83 @@
 package block
 
 import (
-	"context"
-	"fmt"
+	"iter"
 	"sync"
-	"sync/atomic"
 
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
-type TrackedSliceDevice struct {
-	data      ReadonlyDevice
-	blockSize int64
+type Tracker struct {
+	b  *bitset.BitSet
+	mu sync.RWMutex
 
-	nilTracking atomic.Bool
-	dirty       *bitset.BitSet
-	dirtyMu     sync.Mutex
-	empty       []byte
+	blockSize int64
 }
 
-func NewTrackedSliceDevice(blockSize int64, device ReadonlyDevice) (*TrackedSliceDevice, error) {
-	return &TrackedSliceDevice{
-		data:      device,
-		empty:     make([]byte, blockSize),
+func NewTracker(blockSize int64) *Tracker {
+	return &Tracker{
+		// The bitset resizes automatically based on the maximum set bit.
+		b:         bitset.New(0),
 		blockSize: blockSize,
-	}, nil
+	}
 }
 
-func (t *TrackedSliceDevice) Disable() error {
-	size, err := t.data.Size()
-	if err != nil {
-		return fmt.Errorf("failed to get device size: %w", err)
+func NewTrackerFromBitset(b *bitset.BitSet, blockSize int64) *Tracker {
+	return &Tracker{
+		b:         b,
+		blockSize: blockSize,
 	}
+}
+
+func (t *Tracker) Has(off int64) bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
 
-	t.dirty = bitset.New(uint(header.TotalBlocks(size, t.blockSize)))
-	// We are starting with all being dirty.
-	t.dirty.FlipRange(0, t.dirty.Len())
+	return t.b.Test(uint(header.BlockIdx(off, t.blockSize)))
+}
 
-	t.nilTracking.Store(true)
+func (t *Tracker) Add(off int64) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
 
-	return nil
+	t.b.Set(uint(header.BlockIdx(off, t.blockSize)))
 }
 
-func (t *TrackedSliceDevice) Slice(ctx context.Context, off int64, length int64) ([]byte, error) {
-	if t.nilTracking.Load() {
-		t.dirtyMu.Lock()
-		t.dirty.Clear(uint(header.BlockIdx(off, t.blockSize)))
-		t.dirtyMu.Unlock()
+func (t *Tracker) Reset() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
 
-		return t.empty, nil
-	}
+	t.b.ClearAll()
+}
+
+// BitSet returns a clone of the bitset and the block size.
+func (t *Tracker) BitSet() *bitset.BitSet {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.b.Clone()
+}
 
-	return t.data.Slice(ctx, off, length)
+func (t *Tracker) BlockSize() int64 {
+	return t.blockSize
 }
 
-// Return which bytes were not read since Disable.
-// This effectively returns the bytes that have been requested after paused vm and are not dirty.
-func (t *TrackedSliceDevice) Dirty() *bitset.BitSet {
-	t.dirtyMu.Lock()
-	defer t.dirtyMu.Unlock()
+func (t *Tracker) Clone() *Tracker {
+	return &Tracker{
+		b:         t.BitSet(),
+		blockSize: t.BlockSize(),
+	}
+}
+
+func (t *Tracker) Offsets() iter.Seq[int64] {
+	return bitsetOffsets(t.BitSet(), t.BlockSize())
+}
 
-	return t.dirty.Clone()
+func bitsetOffsets(b *bitset.BitSet, blockSize int64) iter.Seq[int64] {
+	return utils.TransformTo(b.EachSet(), func(idx uint) int64 {
+		return header.BlockOffset(int64(idx), blockSize)
+	})
 }

packages/orchestrator/internal/sandbox/block/tracker_test.go

@@ -0,0 +1,167 @@
+package block
+
+import (
+	"maps"
+	"math/rand"
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTracker_AddAndHas(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offset := int64(pageSize * 4)
+
+	// Initially should not be marked
+	assert.False(t, tr.Has(offset), "Expected offset %d not to be marked initially", offset)
+
+	// After adding, should be marked
+	tr.Add(offset)
+	assert.True(t, tr.Has(offset), "Expected offset %d to be marked after Add", offset)
+
+	// Other offsets should not be marked
+	otherOffsets := []int64{
+		0, pageSize, 2 * pageSize, 3 * pageSize, 5 * pageSize, 10 * pageSize,
+	}
+	for _, other := range otherOffsets {
+		if other == offset {
+			continue
+		}
+		assert.False(t, tr.Has(other), "Did not expect offset %d to be marked (only %d should be marked)", other, offset)
+	}
+}
+
+func TestTracker_Reset(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offset := int64(pageSize * 4)
+
+	// Add offset and verify it's marked
+	tr.Add(offset)
+	assert.True(t, tr.Has(offset), "Expected offset %d to be marked after Add", offset)
+
+	// After reset, should not be marked
+	tr.Reset()
+	assert.False(t, tr.Has(offset), "Expected offset %d to be cleared after Reset", offset)
+
+	// Offsets that were never set should also remain unset
+	otherOffsets := []int64{0, pageSize, 2 * pageSize, pageSize * 10}
+	for _, other := range otherOffsets {
+		assert.False(t, tr.Has(other), "Expected offset %d to not be marked after Reset", other)
+	}
+}
+
+func TestTracker_MultipleOffsets(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offsets := []int64{0, pageSize, 2 * pageSize, 10 * pageSize}
+
+	// Add multiple offsets
+	for _, o := range offsets {
+		tr.Add(o)
+	}
+
+	// Verify all offsets are marked
+	for _, o := range offsets {
+		assert.True(t, tr.Has(o), "Expected offset %d to be marked", o)
+	}
+
+	// Check offsets in between added offsets are not set
+	// (Offsets that aren't inside any marked block should not be marked)
+	nonSetOffsets := []int64{
+		3 * pageSize,
+		4 * pageSize,
+		5 * pageSize,
+		6 * pageSize,
+		7 * pageSize,
+		8 * pageSize,
+		9 * pageSize,
+		11 * pageSize,
+	}
+	for _, off := range nonSetOffsets {
+		assert.False(t, tr.Has(off), "Expected offset %d to not be marked (only explicit blocks added)", off)
+	}
+}
+
+func TestTracker_ResetClearsAll(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offsets := []int64{0, pageSize, 2 * pageSize, 10 * pageSize}
+
+	// Add multiple offsets
+	for _, o := range offsets {
+		tr.Add(o)
+	}
+
+	// Reset should clear all
+	tr.Reset()
+
+	// Verify all offsets are cleared
+	for _, o := range offsets {
+		assert.False(t, tr.Has(o), "Expected offset %d to be cleared after Reset", o)
+	}
+	// Check unrelated offsets also not marked
+	moreOffsets := []int64{3 * pageSize, 7 * pageSize, 100, 4095}
+	for _, o := range moreOffsets {
+		assert.False(t, tr.Has(o), "Expected offset %d to not be marked after Reset", o)
+	}
+}
+
+func TestTracker_MisalignedOffset(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	// Test with misaligned offset
+	misalignedOffset := int64(123)
+	tr.Add(misalignedOffset)
+
+	// Should be set for the block containing the offset—that is, block 0 (0..4095)
+	assert.True(t, tr.Has(misalignedOffset), "Expected misaligned offset %d to be marked (should mark its containing block)", misalignedOffset)
+
+	// Now check that any offset in the same block is also considered marked
+	anotherOffsetInSameBlock := int64(1000)
+	assert.True(t, tr.Has(anotherOffsetInSameBlock), "Expected offset %d to be marked as in same block as %d", anotherOffsetInSameBlock, misalignedOffset)
+
+	// But not for a different block
+	offsetInNextBlock := int64(pageSize)
+	assert.False(t, tr.Has(offsetInNextBlock), "Did not expect offset %d to be marked", offsetInNextBlock)
+
+	// And not far outside any set block
+	offsetFar := int64(2 * pageSize)
+	assert.False(t, tr.Has(offsetFar), "Did not expect offset %d to be marked", offsetFar)
+}
+
+func TestTracker_Offsets(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	numOffsets := 300
+
+	offsetsMap := map[int64]struct{}{}
+
+	for range numOffsets {
+		select {
+		case <-t.Context().Done():
+			t.FailNow()
+		default:
+		}
+
+		base := int64(rand.Intn(121)) // 0..120
+		offset := base * pageSize
+
+		offsetsMap[offset] = struct{}{}
+		tr.Add(offset)
+	}
+
+	expectedOffsets := slices.Collect(maps.Keys(offsetsMap))
+	actualOffsets := slices.Collect(tr.Offsets())
+
+	assert.Len(t, actualOffsets, len(expectedOffsets))
+	assert.ElementsMatch(t, expectedOffsets, actualOffsets)
+}

packages/orchestrator/internal/sandbox/build/diff.go

@@ -62,3 +62,7 @@ func (n *NoDiff) CacheKey() DiffStoreKey {
 func (n *NoDiff) Init(context.Context) error {
 	return NoDiffError{}
 }
+
+func (n *NoDiff) BlockSize() int64 {
+	return 0
+}

packages/orchestrator/internal/sandbox/build/local_diff.go

@@ -137,3 +137,7 @@ func (b *localDiff) CacheKey() DiffStoreKey {
 func (b *localDiff) Init(context.Context) error {
 	return nil
 }
+
+func (b *localDiff) BlockSize() int64 {
+	return b.blockSize
+}

packages/orchestrator/internal/sandbox/build/storage_diff.go

@@ -158,3 +158,7 @@ func (b *StorageDiff) FileSize() (int64, error) {
 
 	return c.FileSize()
 }
+
+func (b *StorageDiff) BlockSize() int64 {
+	return b.blockSize
+}

packages/orchestrator/internal/sandbox/sandbox.go

@@ -672,7 +672,7 @@ func (s *Sandbox) Shutdown(ctx context.Context) error {
 		return fmt.Errorf("failed to pause VM: %w", err)
 	}
 
-	if err := s.memory.Disable(); err != nil {
+	if _, err := s.memory.Disable(ctx); err != nil {
 		return fmt.Errorf("failed to disable uffd: %w", err)
 	}
 
@@ -750,8 +750,12 @@ func (s *Sandbox) Pause(
 		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}
 
-	if err := s.memory.Disable(); err != nil {
-		return nil, fmt.Errorf("failed to disable uffd: %w", err)
+	// This disables the uffd and returns the dirty pages.
+	// With FC async io engine, there can be some further writes to the memory during the actual create snapshot process,
+	// but as we are still including even read pages as dirty so this should not introduce more bugs right now.
+	dirty, err := s.memory.Disable(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dirty pages: %w", err)
 	}
 
 	// Snapfile is not closed as it's returned and cached for later use (like resume)
@@ -798,7 +802,7 @@ func (s *Sandbox) Pause(
 		originalMemfile.Header(),
 		&MemoryDiffCreator{
 			memfile:    memfile,
-			dirtyPages: s.memory.Dirty(),
+			dirtyPages: dirty.BitSet(),
 			blockSize:  originalMemfile.BlockSize(),
 			doneHook: func(context.Context) error {
 				return memfile.Close()
@@ -979,7 +983,7 @@ func serveMemory(
 	ctx, span := tracer.Start(ctx, "serve-memory")
 	defer span.End()
 
-	fcUffd, err := uffd.New(memfile, socketPath, memfile.BlockSize())
+	fcUffd, err := uffd.New(memfile, socketPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create uffd: %w", err)
 	}

packages/orchestrator/internal/sandbox/uffd/fdexit/fdexit.go

@@ -7,6 +7,8 @@ import (
 	"sync"
 )
 
+var ErrFdExit = errors.New("fd exit signal")
+
 // FdExit is a wrapper around a pipe that allows to signal the exit of the uffd.
 type FdExit struct {
 	r    *os.File

packages/orchestrator/internal/sandbox/uffd/memory/mapping.go

@@ -21,10 +21,10 @@ func NewMapping(regions []Region) *Mapping {
 }
 
 // GetOffset returns the relative offset and the page size of the mapped range for a given address.
-func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uint64, error) {
+func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uintptr, error) {
 	for _, r := range m.Regions {
 		if hostVirtAddr >= r.BaseHostVirtAddr && hostVirtAddr < r.endHostVirtAddr() {
-			return r.shiftedOffset(hostVirtAddr), uint64(r.PageSize), nil
+			return r.shiftedOffset(hostVirtAddr), r.PageSize, nil
 		}
 	}