Add dirty page tracking via UFFD write protect

packages/orchestrator/go.mod

@@ -20,7 +20,6 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.0
 	github.com/hashicorp/consul/api v1.30.0
 	github.com/jellydator/ttlcache/v3 v3.3.1-0.20250207140243-aefc35918359
-	github.com/loopholelabs/userfaultfd-go v0.1.2
 	github.com/ngrok/firewall_toolkit v0.0.18
 	github.com/pkg/errors v0.9.1
 	github.com/pojntfx/go-nbd v0.3.2

packages/orchestrator/go.sum

@@ -739,8 +739,6 @@ github.com/launchdarkly/go-test-helpers/v2 v2.2.0 h1:L3kGILP/6ewikhzhdNkHy1b5y4z
 github.com/launchdarkly/go-test-helpers/v2 v2.2.0/go.mod h1:L7+th5govYp5oKU9iN7To5PgznBuIjBPn+ejqKR0avw=
 github.com/launchdarkly/go-test-helpers/v3 v3.0.2 h1:rh0085g1rVJM5qIukdaQ8z1XTWZztbJ49vRZuveqiuU=
 github.com/launchdarkly/go-test-helpers/v3 v3.0.2/go.mod h1:u2ZvJlc/DDJTFrshWW50tWMZHLVYXofuSHUfTU/eIwM=
-github.com/loopholelabs/userfaultfd-go v0.1.2 h1:HwXFNoQ+/eWNgYIcIyrqn54gDVVJk+TmszYxMGnJVu4=
-github.com/loopholelabs/userfaultfd-go v0.1.2/go.mod h1:6+5c50Ji7MUXuWUSrPUhAttECwmEeLAmR33FlP7Fn4o=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/mailru/easyjson v0.0.0-20190614124828-94de47d64c63/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=
 github.com/mailru/easyjson v0.0.0-20190626092158-b2ccc519800e/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=

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

@@ -12,10 +12,14 @@ func (ErrBytesNotAvailable) Error() string {
 	return "The requested bytes are not available on the device"
 }
 
+type Slicer interface {
+	Slice(off, length int64) ([]byte, error)
+}
+
 type ReadonlyDevice interface {
 	io.ReaderAt
 	io.Closer
-	Slice(off, length int64) ([]byte, error)
+	Slicer
 	Size() (int64, error)
 	BlockSize() int64
 	Header() *header.Header

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

@@ -4,6 +4,7 @@ import (
 	"context"
 	"io"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 )
 
@@ -23,7 +24,7 @@ const (
 type Diff interface {
 	io.Closer
 	io.ReaderAt
-	Slice(off, length int64) ([]byte, error)
+	block.Slicer
 	CacheKey() DiffStoreKey
 	CachePath() (string, error)
 	FileSize() (int64, error)

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

@@ -0,0 +1,39 @@
+package firecracker
+
+import "fmt"
+
+type Mappings interface {
+	GetRange(addr uintptr) (offset int64, pagesize int64, err error)
+}
+
+type GuestRegionUffdMapping struct {
+	BaseHostVirtAddr uintptr `json:"base_host_virt_addr"`
+	Size             uintptr `json:"size"`
+	Offset           uintptr `json:"offset"`
+	// This is actually in bytes—it is deprecated and they introduced "page_size"
+	PageSize uintptr `json:"page_size_kib"`
+}
+
+type FcMappings []GuestRegionUffdMapping
+
+func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) int64 {
+	return int64(m.Offset + addr - m.BaseHostVirtAddr)
+}
+
+func (m FcMappings) GetRange(addr uintptr) (int64, int64, error) {
+	return getMappedRange(addr, m)
+}
+
+// Returns the relative offset and the page size of the mapped range for a given address
+func getMappedRange(addr uintptr, mappings []GuestRegionUffdMapping) (int64, int64, error) {
+	for _, m := range mappings {
+		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)
+}