uffd: add support for UFFD_EVENT_REMOVE events

Enabling Firecracker free-page-reporting feature requires us to handle
remove events (UFFD_EVENT_REMOVE) in our userfaultfd handler. These
events are triggered whenever Firecracker calls madvise(MADV_DONTNEED)
(or similar) on a range of guest memory addresses.

The main thing that changes on our logic is that page faults in a page
that has previously been removed need to be served with a zero page
rather than a page from the snapshot file.

This commit changes the page fault serving logic to:
1. Introduce tracking of the state of every page in the guest's memory
   mappings.
2. Add logic to handle the new UFFD_EVENT_REMOVE event
3. Modify existing logic to take into account current state when
   deciding how to handle each page fault

Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>

Author: bchalios

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

@@ -15,6 +15,15 @@ func (e AddressNotFoundError) Error() string {
 	return fmt.Sprintf("host virtual address %d not found in any mapping", e.hostVirtAddr)
 }
 
+type AddressNotPageAlignedError struct {
+	hostVirtAddr uintptr
+	pageSize     uintptr
+}
+
+func (e AddressNotPageAlignedError) Error() string {
+	return fmt.Sprintf("host virtual address %d not aligned in regions page size %d", e.hostVirtAddr, e.pageSize)
+}
+
 type OffsetNotFoundError struct {
 	offset int64
 }
@@ -104,3 +113,49 @@ func (m *Mapping) GetHostVirtAddr(off int64) (uintptr, uintptr, error) {
 
 	return region.shiftedHostVirtAddr(off), region.PageSize, nil
 }
+
+// GetPageState returns the state of the page with the given start address.
+func (m *Mapping) GetPageState(hostVirtAddr uintptr) (PageState, error) {
+	for _, r := range m.Regions {
+		if r.addrInRegion(hostVirtAddr) {
+			if hostVirtAddr%r.PageSize != 0 {
+				return "", AddressNotPageAlignedError{hostVirtAddr: hostVirtAddr, pageSize: r.PageSize}
+			}
+
+			return r.pageState(hostVirtAddr), nil
+		}
+	}
+
+	return "", AddressNotFoundError{hostVirtAddr: hostVirtAddr}
+}
+
+// SetPageState sets a new state for a range of pages.
+func (m *Mapping) SetPageState(startHostVirtAddr uintptr, endHostVirtAddr uintptr, pageState PageState) error {
+	for i := range m.Regions {
+		if m.Regions[i].addrInRegion(startHostVirtAddr) {
+			// The entirety of the range needs to be in the same region. In other words: we can't have here a region
+			// [startHostVirtAddr, endHostVirtAddr) that spans across more than one guest memory regions
+			if !m.Regions[i].addrInRegion(endHostVirtAddr) {
+				return AddressNotFoundError{hostVirtAddr: endHostVirtAddr}
+			}
+
+			// Both start and end addresses need to be page aligned.
+			pageSize := m.Regions[i].PageSize
+			if startHostVirtAddr%pageSize != 0 {
+				return AddressNotPageAlignedError{hostVirtAddr: startHostVirtAddr, pageSize: pageSize}
+			}
+
+			if endHostVirtAddr%pageSize != 0 {
+				return AddressNotPageAlignedError{hostVirtAddr: endHostVirtAddr, pageSize: pageSize}
+			}
+
+			for addr := startHostVirtAddr; addr < endHostVirtAddr; addr += pageSize {
+				m.Regions[i].setPageState(addr, pageState)
+			}
+
+			return nil
+		}
+	}
+
+	return AddressNotFoundError{hostVirtAddr: startHostVirtAddr}
+}

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

@@ -1,5 +1,17 @@
 package memory
 
+// PageState describes the states that guest memory page can be in
+type PageState string
+
+const (
+	// Page has not been faulted in yet
+	Unfaulted PageState = "unfaulted"
+	// Page has been faulted in
+	Faulted PageState = "faulted"
+	// Page has been removed
+	Removed PageState = "removed"
+)
+
 // Region is a mapping of a region of memory of the guest to a region of memory on the host.
 // The serialization is based on the Firecracker UFFD protocol communication.
 // https://github.com/firecracker-microvm/firecracker/blob/ceeca6a14284537ae0b2a192cd2ffef10d3a81e2/src/vmm/src/persist.rs#L96
@@ -9,6 +21,9 @@ type Region struct {
 	Offset           uintptr `json:"offset"`
 	// This field is deprecated in the newer version of the Firecracker with a new field `page_size`.
 	PageSize uintptr `json:"page_size_kib"` // This is actually in bytes in the deprecated version.
+	// This field is not used in the serialization. It's metadata we're keeping to track the state of the
+	// pages within this region
+	PageState []PageState `json:"-"`
 }
 
 // endOffset returns the end offset of the region in bytes.
@@ -32,3 +47,27 @@ func (r *Region) shiftedOffset(addr uintptr) int64 {
 func (r *Region) shiftedHostVirtAddr(off int64) uintptr {
 	return uintptr(off) + r.BaseHostVirtAddr - r.Offset
 }
+
+// pageState returns the current PageState of a page in the region.
+//
+// This assumes that the caller has already checked that `addr` lies
+// within this region.
+func (r *Region) pageState(addr uintptr) PageState {
+	page_idx := (addr - r.BaseHostVirtAddr) / r.PageSize
+
+	return r.PageState[page_idx]
+}
+
+// setPageState sets the current PageState of a page in the region.
+//
+// This assumes that the caller has already checked that `addr` lies
+// within this region.
+func (r *Region) setPageState(addr uintptr, state PageState) {
+	page_idx := (addr - r.BaseHostVirtAddr) / r.PageSize
+	r.PageState[page_idx] = state
+}
+
+// addrInRegion returns true if an address is included in the region
+func (r *Region) addrInRegion(addr uintptr) bool {
+	return addr >= r.BaseHostVirtAddr && addr < r.endHostVirtAddr()
+}

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

@@ -144,6 +144,15 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string, fdExit *fdexit.FdEx
 		return fmt.Errorf("failed parsing memory mapping data: %w", err)
 	}
 
+	// Initialize PageState for each region (not included in JSON serialization)
+	for i := range regions {
+		numPages := regions[i].Size / regions[i].PageSize
+		regions[i].PageState = make([]memory.PageState, numPages)
+		for j := range regions[i].PageState {
+			regions[i].PageState[j] = memory.Unfaulted
+		}
+	}
+
 	controlMsgs, err := syscall.ParseSocketControlMessage(uffdBuf[:numBytesFd])
 	if err != nil {
 		return fmt.Errorf("failed parsing control messages: %w", err)

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

@@ -16,6 +16,12 @@ struct uffd_pagefault {
 	__u64	address;
 	__u32 ptid;
 };
+
+struct uffd_remove {
+	__u64 start;
+	__u64 end;
+};
+
 */
 import "C"
 
@@ -30,14 +36,18 @@ const (
 
 	UFFD_API             = C.UFFD_API
 	UFFD_EVENT_PAGEFAULT = C.UFFD_EVENT_PAGEFAULT
+	UFFD_EVENT_REMOVE    = C.UFFD_EVENT_REMOVE
 
 	UFFDIO_REGISTER_MODE_MISSING = C.UFFDIO_REGISTER_MODE_MISSING
 
 	UFFDIO_API      = C.UFFDIO_API
 	UFFDIO_REGISTER = C.UFFDIO_REGISTER
 	UFFDIO_COPY     = C.UFFDIO_COPY
+	UFFDIO_ZEROPAGE = C.UFFDIO_ZEROPAGE
 
 	UFFD_PAGEFAULT_FLAG_WRITE = C.UFFD_PAGEFAULT_FLAG_WRITE
+	UFFD_PAGEFAULT_FLAG_MINOR = C.UFFD_PAGEFAULT_FLAG_MINOR
+	UFFD_PAGEFAULT_FLAG_WP    = C.UFFD_PAGEFAULT_FLAG_WP
 
 	UFFD_FEATURE_MISSING_HUGETLBFS = C.UFFD_FEATURE_MISSING_HUGETLBFS
 )
@@ -49,11 +59,13 @@ type (
 
 	UffdMsg       = C.struct_uffd_msg
 	UffdPagefault = C.struct_uffd_pagefault
+	UffdRemove    = C.struct_uffd_remove
 
 	UffdioAPI          = C.struct_uffdio_api
 	UffdioRegister     = C.struct_uffdio_register
 	UffdioRange        = C.struct_uffdio_range
 	UffdioCopy         = C.struct_uffdio_copy
+	UffdioZero         = C.struct_uffdio_zeropage
 	UffdioWriteProtect = C.struct_uffdio_writeprotect
 )
 
@@ -88,6 +100,14 @@ func newUffdioCopy(b []byte, address CULong, pagesize CULong, mode CULong, bytes
 	}
 }
 
+func newUffdioZero(address, pagesize, mode CULong) UffdioZero {
+	return UffdioZero{
+		_range:   newUffdioRange(address, pagesize),
+		mode:     mode,
+		zeropage: 0,
+	}
+}
+
 func getMsgEvent(msg *UffdMsg) CUChar {
 	return msg.event
 }
@@ -120,6 +140,21 @@ func (f Fd) copy(addr, pagesize uintptr, data []byte, mode CULong) error {
 	return nil
 }
 
+func (f Fd) zero(addr, pagesize uintptr, mode CULong) error {
+	zero := newUffdioZero(CULong(addr)&^CULong(pagesize-1), CULong(pagesize), mode)
+
+	if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(f), UFFDIO_ZEROPAGE, uintptr(unsafe.Pointer(&zero))); errno != 0 {
+		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)
+	}
+
+	return nil
+}
+
 func (f Fd) close() error {
 	return syscall.Close(int(f))
 }