Extract low level UFFD handler to a separate package

.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/template/mask_template.go

@@ -38,7 +38,7 @@ func NewMaskTemplate(
 	return t
 }
 
-func (c *MaskTemplate) Close(_ context.Context) error {
+func (c *MaskTemplate) Close(context.Context) error {
 	if c.memfile != nil {
 		return (*c.memfile).Close()
 	}

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/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
+	}
+
+	closeMmap := func() error {
+		return syscall.Munmap(b)
+	}
+
+	return b, uintptr(unsafe.Pointer(&b[0])), closeMmap, 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
+}

packages/orchestrator/internal/sandbox/uffd/handler.go → packages/orchestrator/internal/sandbox/uffd/uffd.go

@@ -16,7 +16,8 @@ import (
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit"
-	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/mapping"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/memory"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/userfaultfd"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
@@ -116,7 +117,7 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error {
 
 	mappingsBuf = mappingsBuf[:numBytesMappings]
 
-	var m mapping.FcMappings
+	var m memory.MemfileMap
 
 	err = json.Unmarshal(mappingsBuf, &m)
 	if err != nil {
@@ -141,20 +142,19 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error {
 		return fmt.Errorf("expected 1 fd: found %d", len(fds))
 	}
 
-	uffd := fds[0]
+	uffd := userfaultfd.NewUserfaultfdFromFd(uintptr(fds[0]))
 
 	defer func() {
-		closeErr := syscall.Close(uffd)
+		closeErr := uffd.Close()
 		if closeErr != nil {
 			zap.L().Error("failed to close uffd", logger.WithSandboxID(sandboxId), zap.String("socket_path", u.socketPath), zap.Error(closeErr))
 		}
 	}()
 
 	u.readyCh <- struct{}{}
 
-	err = Serve(
+	err = uffd.Serve(
 		ctx,
-		uffd,
 		m,
 		u.memfile,
 		u.fdExit,
@@ -179,10 +179,6 @@ func (u *Uffd) Exit() *utils.ErrorOnce {
 	return u.exit
 }
 
-func (u *Uffd) TrackAndReturnNil() error {
-	return u.lis.Close()
-}
-
 func (u *Uffd) Disable() error {
 	return u.memfile.Disable()
 }

packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go

@@ -27,21 +27,12 @@ const (
 	UFFD_EVENT_PAGEFAULT = C.UFFD_EVENT_PAGEFAULT
 
 	UFFDIO_REGISTER_MODE_MISSING = C.UFFDIO_REGISTER_MODE_MISSING
-	UFFDIO_REGISTER_MODE_WP      = C.UFFDIO_REGISTER_MODE_WP
 
-	UFFDIO_WRITEPROTECT_MODE_WP = C.UFFDIO_WRITEPROTECT_MODE_WP
-	UFFDIO_COPY_MODE_WP         = C.UFFDIO_COPY_MODE_WP
+	UFFDIO_API      = C.UFFDIO_API
+	UFFDIO_REGISTER = C.UFFDIO_REGISTER
+	UFFDIO_COPY     = C.UFFDIO_COPY
 
-	UFFDIO_API          = C.UFFDIO_API
-	UFFDIO_REGISTER     = C.UFFDIO_REGISTER
-	UFFDIO_WRITEPROTECT = C.UFFDIO_WRITEPROTECT
-	UFFDIO_COPY         = C.UFFDIO_COPY
-
-	UFFD_PAGEFAULT_FLAG_WP    = C.UFFD_PAGEFAULT_FLAG_WP
-	UFFD_PAGEFAULT_FLAG_WRITE = C.UFFD_PAGEFAULT_FLAG_WRITE
-
-	UFFD_FEATURE_MISSING_HUGETLBFS  = C.UFFD_FEATURE_MISSING_HUGETLBFS
-	UFFD_FEATURE_WP_HUGETLBFS_SHMEM = C.UFFD_FEATURE_WP_HUGETLBFS_SHMEM
+	UFFD_FEATURE_MISSING_HUGETLBFS = C.UFFD_FEATURE_MISSING_HUGETLBFS
 )
 
 type (
@@ -52,11 +43,10 @@ type (
 	UffdMsg       = C.struct_uffd_msg
 	UffdPagefault = C.struct_uffd_pagefault
 
-	UffdioAPI          = C.struct_uffdio_api
-	UffdioRegister     = C.struct_uffdio_register
-	UffdioRange        = C.struct_uffdio_range
-	UffdioCopy         = C.struct_uffdio_copy
-	UffdioWriteProtect = C.struct_uffdio_writeprotect
+	UffdioAPI      = C.struct_uffdio_api
+	UffdioRegister = C.struct_uffdio_register
+	UffdioRange    = C.struct_uffdio_range
+	UffdioCopy     = C.struct_uffdio_copy
 )
 
 func NewUffdioAPI(api, features CULong) UffdioAPI {
@@ -79,23 +69,13 @@ func NewUffdioRegister(start, length, mode CULong) UffdioRegister {
 func NewUffdioCopy(b []byte, address CULong, pagesize CULong, mode CULong, bytesCopied CLong) UffdioCopy {
 	return UffdioCopy{
 		src:  CULong(uintptr(unsafe.Pointer(&b[0]))),
-		dst:  address &^ (pagesize - 1),
+		dst:  address,
 		len:  pagesize,
 		mode: mode,
 		copy: bytesCopied,
 	}
 }
 
-func NewUffdioWriteProtect(start, length, mode CULong) UffdioWriteProtect {
-	return UffdioWriteProtect{
-		_range: UffdioRange{
-			start: start,
-			len:   length,
-		},
-		mode: mode,
-	}
-}
-
 func GetMsgEvent(msg *UffdMsg) CUChar {
 	return msg.event
 }
@@ -104,14 +84,6 @@ func GetMsgArg(msg *UffdMsg) [24]byte {
 	return msg.arg
 }
 
-func GetPagefaultAddress(pagefault *UffdPagefault) CULong {
-	return pagefault.address
-}
-
-func IsWritePageFault(pagefault *UffdPagefault) bool {
-	return pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0
-}
-
-func IsWriteProtectPageFault(pagefault *UffdPagefault) bool {
-	return pagefault.flags&UFFD_PAGEFAULT_FLAG_WP != 0
+func GetPagefaultAddress(pagefault *UffdPagefault) uintptr {
+	return uintptr(pagefault.address)
 }