Add UFFD write protection for tracking dirty pages

.github/workflows/pr-tests.yml

@@ -33,6 +33,16 @@ jobs:
         with:
           job_key: unit-tests-${{ matrix.package }}
 
+      - name: Enable unprivileged uffd mode
+        run: |
+          echo 1 | sudo tee /proc/sys/vm/unprivileged_userfaultfd
+
+      - name: Enable hugepages
+        run: |
+          sudo mkdir -p /mnt/hugepages
+          sudo mount -t hugetlbfs none /mnt/hugepages
+          echo 128 | sudo tee /proc/sys/vm/nr_hugepages
+
       - name: Run tests
         working-directory: ${{ matrix.package }}
         run: go test -v ${{ matrix.test_path }}

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

@@ -1,7 +1,20 @@
-package mapping
+package memory
 
 import "fmt"
 
+type MemfileMap []GuestRegionUffdMapping
+
+// GetOffset returns the relative offset and the page size of the mapped range for a given address
+func (m MemfileMap) GetOffset(hostVirtAddr uintptr) (int64, uint64, error) {
+	for _, m := range m {
+		if hostVirtAddr >= m.BaseHostVirtAddr && hostVirtAddr < m.BaseHostVirtAddr+m.Size {
+			return m.relativeOffset(hostVirtAddr), uint64(m.PageSize), nil
+		}
+	}
+
+	return 0, 0, fmt.Errorf("address %d not found in any mapping", hostVirtAddr)
+}
+
 type GuestRegionUffdMapping struct {
 	BaseHostVirtAddr uintptr `json:"base_host_virt_addr"`
 	Size             uintptr `json:"size"`
@@ -14,19 +27,3 @@ type GuestRegionUffdMapping struct {
 func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) int64 {
 	return int64(m.Offset + addr - m.BaseHostVirtAddr)
 }
-
-type FcMappings []GuestRegionUffdMapping
-
-// Returns the relative offset and the page size of the mapped range for a given address
-func (m FcMappings) GetRange(addr uintptr) (int64, int64, error) {
-	for _, m := range m {
-		if addr < m.BaseHostVirtAddr || m.BaseHostVirtAddr+m.Size <= addr {
-			// Outside of this mapping
-			continue
-		}
-
-		return m.relativeOffset(addr), int64(m.PageSize), nil
-	}
-
-	return 0, 0, fmt.Errorf("address %d not found in any mapping", addr)
-}

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

@@ -0,0 +1,5 @@
+package memory
+
+type MemoryMap interface {
+	GetOffset(hostVirtAddr uintptr) (offset int64, pagesize uint64, err error)
+}

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

@@ -0,0 +1,43 @@
+package memory
+
+import (
+	"sync"
+
+	"github.com/bits-and-blooms/bitset"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+type Tracker struct {
+	bitset    *bitset.BitSet
+	blockSize int64
+	mu        sync.RWMutex
+}
+
+func NewTracker(size, blockSize int64) *Tracker {
+	return &Tracker{
+		bitset:    bitset.New(uint(header.TotalBlocks(size, blockSize))),
+		blockSize: blockSize,
+		mu:        sync.RWMutex{},
+	}
+}
+
+func (t *Tracker) Mark(offset int64) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	t.bitset.Set(uint(header.BlockIdx(offset, t.blockSize)))
+}
+
+func (t *Tracker) BitSet() *bitset.BitSet {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.bitset.Clone()
+}
+
+func (t *Tracker) Check(offset int64) bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.bitset.Test(uint(header.BlockIdx(offset, t.blockSize)))
+}

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

@@ -9,7 +9,6 @@ import (
 )
 
 type MemoryBackend interface {
-	Disable() error
 	Dirty() *bitset.BitSet
 
 	Start(ctx context.Context, sandboxId string) error

packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go

@@ -0,0 +1,17 @@
+package testutils
+
+import "context"
+
+// contentSlicer exposes byte slice via the Slicer interface.
+// This is used for testing purposes.
+type contentSlicer struct {
+	content []byte
+}
+
+func newContentSlicer(content []byte) *contentSlicer {
+	return &contentSlicer{content: content}
+}
+
+func (s *contentSlicer) Slice(_ context.Context, offset, size int64) ([]byte, error) {
+	return s.content[offset : offset+size], nil
+}

packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go

@@ -0,0 +1,46 @@
+package testutils
+
+import (
+	"sync"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// ContiguousMap is a mapping that is contiguous in the host virtual address space.
+// This is used for testing purposes.
+type ContiguousMap struct {
+	start    uintptr
+	size     uint64
+	pagesize uint64
+
+	accessedOffsets map[uint64]struct{}
+	mu              sync.RWMutex
+}
+
+func NewContiguousMap(start uintptr, size, pagesize uint64) *ContiguousMap {
+	return &ContiguousMap{
+		start:           start,
+		size:            size,
+		pagesize:        pagesize,
+		accessedOffsets: make(map[uint64]struct{}),
+	}
+}
+
+func (m *ContiguousMap) GetOffset(addr uintptr) (int64, uint64, error) {
+	offset := addr - m.start
+	pagesize := m.pagesize
+
+	m.mu.Lock()
+	m.accessedOffsets[uint64(offset)] = struct{}{}
+	m.mu.Unlock()
+
+	return int64(offset), pagesize, nil
+}
+
+func (m *ContiguousMap) Map() map[uint64]struct{} {
+	return m.accessedOffsets
+}
+
+func (m *ContiguousMap) Keys() []uint64 {
+	return utils.MapKeys(m.accessedOffsets)
+}

packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go

@@ -0,0 +1,44 @@
+package testutils
+
+import (
+	"fmt"
+	"math"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+func NewPageMmap(size, pagesize uint64) ([]byte, uintptr, func() error, error) {
+	if pagesize == header.PageSize {
+		return newMmap(size, header.PageSize, 0)
+	}
+
+	if pagesize == header.HugepageSize {
+		return newMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB)
+	}
+
+	return nil, 0, nil, fmt.Errorf("unsupported page size: %d", pagesize)
+}
+
+func newMmap(size, pagesize uint64, flags int) ([]byte, uintptr, func() error, error) {
+	l := int(math.Ceil(float64(size)/float64(pagesize)) * float64(pagesize))
+	b, err := syscall.Mmap(
+		-1,
+		0,
+		l,
+		syscall.PROT_READ|syscall.PROT_WRITE,
+		syscall.MAP_PRIVATE|syscall.MAP_ANONYMOUS|flags,
+	)
+	if err != nil {
+		return nil, 0, nil, err
+	}
+
+	close := func() error {
+		return syscall.Munmap(b)
+	}
+
+	return b, uintptr(unsafe.Pointer(&b[0])), close, nil
+}

packages/orchestrator/internal/sandbox/uffd/testutils/random_data.go

@@ -0,0 +1,41 @@
+package testutils
+
+import (
+	"crypto/rand"
+)
+
+func RandomPages(pagesize, numberOfPages uint64) (data *contentSlicer, size uint64) {
+	size = pagesize * numberOfPages
+
+	n := int(size)
+	buf := make([]byte, n)
+	if _, err := rand.Read(buf); err != nil {
+		panic(err)
+	}
+
+	data = newContentSlicer(buf)
+
+	return data, size
+}
+
+// DiffByte returns the first byte index where a and b differ.
+// It also returns the differing byte values (want, got).
+// If slices are identical, it returns -1.
+func DiffByte(a, b []byte) (idx int, want, got byte) {
+	minLen := len(a)
+	if len(b) < minLen {
+		minLen = len(b)
+	}
+
+	for i := range minLen {
+		if a[i] != b[i] {
+			return i, b[i], a[i]
+		}
+	}
+
+	if len(a) != len(b) {
+		return minLen, 0, 0
+	}
+
+	return -1, 0, 0
+}