Extract low level UFFD handler to a separate package

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/contiguous_map.go

@@ -0,0 +1,24 @@
+package testutils
+
+// 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
+}
+
+func NewContiguousMap(start uintptr, size, pagesize uint64) *ContiguousMap {
+	return &ContiguousMap{
+		start:    start,
+		size:     size,
+		pagesize: pagesize,
+	}
+}
+
+func (m *ContiguousMap) GetOffset(addr uintptr) (int64, uint64, error) {
+	offset := addr - m.start
+	pagesize := m.pagesize
+
+	return int64(offset), pagesize, nil
+}

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

@@ -0,0 +1,48 @@
+package testutils
+
+import (
+	"context"
+	"math"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+type mockSlicer struct {
+	content []byte
+}
+
+func newMockSlicer(content []byte) *mockSlicer {
+	return &mockSlicer{content: content}
+}
+
+func (s *mockSlicer) Slice(_ context.Context, offset, size int64) ([]byte, error) {
+	return s.content[offset : offset+size], nil
+}
+
+func New4KPageMmap(size uint64) ([]byte, uintptr) {
+	return newMockMmap(size, header.PageSize, 0)
+}
+
+func New2MPageMmap(size uint64) ([]byte, uintptr) {
+	return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB)
+}
+
+func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr) {
+	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
+	}
+
+	return b, uintptr(unsafe.Pointer(&b[0]))
+}

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

@@ -0,0 +1,31 @@
+package testutils
+
+func repeatToSize(src []byte, size uint64) []byte {
+	if len(src) == 0 || size <= 0 {
+		return nil
+	}
+
+	dst := make([]byte, size)
+	for i := uint64(0); i < size; i += uint64(len(src)) {
+		end := i + uint64(len(src))
+		if end > size {
+			end = size
+		}
+		copy(dst[i:end], src[:end-i])
+	}
+
+	return dst
+}
+
+func PrepareTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size uint64) {
+	size = pagesize * pagesInTestData
+
+	data = newMockSlicer(
+		repeatToSize(
+			[]byte("Hello from userfaultfd! This is our test content that should be readable after the page fault."),
+			size,
+		),
+	)
+
+	return data, size
+}

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"
 )
@@ -39,6 +40,8 @@ type Uffd struct {
 
 	memfile    *block.TrackedSliceDevice
 	socketPath string
+
+	writeProtection bool
 }
 
 var _ MemoryBackend = (*Uffd)(nil)
@@ -55,11 +58,12 @@ func New(memfile block.ReadonlyDevice, socketPath string, blockSize int64) (*Uff
 	}
 
 	return &Uffd{
-		exit:       utils.NewErrorOnce(),
-		readyCh:    make(chan struct{}, 1),
-		fdExit:     fdExit,
-		memfile:    trackedMemfile,
-		socketPath: socketPath,
+		exit:            utils.NewErrorOnce(),
+		readyCh:         make(chan struct{}, 1),
+		fdExit:          fdExit,
+		memfile:         trackedMemfile,
+		socketPath:      socketPath,
+		writeProtection: true,
 	}, nil
 }
 
@@ -116,7 +120,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 +145,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 +182,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

@@ -31,17 +31,18 @@ const (
 
 	UFFDIO_WRITEPROTECT_MODE_WP = C.UFFDIO_WRITEPROTECT_MODE_WP
 	UFFDIO_COPY_MODE_WP         = C.UFFDIO_COPY_MODE_WP
+	UFFDIO_COPY_MODE_DONTWAKE   = C.UFFDIO_COPY_MODE_DONTWAKE
 
 	UFFDIO_API          = C.UFFDIO_API
 	UFFDIO_REGISTER     = C.UFFDIO_REGISTER
 	UFFDIO_WRITEPROTECT = C.UFFDIO_WRITEPROTECT
 	UFFDIO_COPY         = C.UFFDIO_COPY
+	UFFDIO_WAKE         = C.UFFDIO_WAKE
 
 	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 (
@@ -79,7 +80,7 @@ func NewUffdioRegister(start, length, mode CULong) UffdioRegister {
 func NewUffdioCopy(b []byte, address CULong, pagesize CULong, mode CULong, copy CLong) UffdioCopy {
 	return UffdioCopy{
 		src:  CULong(uintptr(unsafe.Pointer(&b[0]))),
-		dst:  address &^ (pagesize - 1),
+		dst:  address,
 		len:  pagesize,
 		mode: mode,
 		copy: copy,
@@ -104,14 +105,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)
 }