Use /proc/PID/mem to access sandbox memory

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

@@ -0,0 +1,17 @@
+package block
+
+type Range struct {
+	// Start is the start address of the range in bytes.
+	// Start is inclusive.
+	Start int64
+	// Size is the size of the range in bytes.
+	Size int64
+}
+
+// NewRange creates a new range from a start address and size in bytes.
+func NewRange(start, size int64) Range {
+	return Range{
+		Start: start,
+		Size:  size,
+	}
+}

packages/orchestrator/internal/sandbox/diffcreator.go

@@ -3,14 +3,11 @@ package sandbox
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
-	"os"
 
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/rootfs"
-	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
@@ -28,7 +25,7 @@ func (r *RootfsDiffCreator) process(ctx context.Context, out io.Writer) (*header
 }
 
 type MemoryDiffCreator struct {
-	memfile    *storage.TemporaryMemfile
+	memory     io.ReaderAt
 	dirtyPages *bitset.BitSet
 	blockSize  int64
 	doneHook   func(context.Context) error
@@ -42,15 +39,9 @@ func (r *MemoryDiffCreator) process(ctx context.Context, out io.Writer) (h *head
 		}
 	}()
 
-	memfileSource, err := os.Open(r.memfile.Path())
-	if err != nil {
-		return nil, fmt.Errorf("failed to open memfile: %w", err)
-	}
-	defer memfileSource.Close()
-
 	return header.WriteDiffWithTrace(
 		ctx,
-		memfileSource,
+		r.memory,
 		r.blockSize,
 		r.dirtyPages,
 		out,

packages/orchestrator/internal/sandbox/sandbox.go

@@ -24,6 +24,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/memory"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/metadata"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
@@ -703,10 +704,25 @@ func (s *Sandbox) Pause(
 	err = s.process.CreateSnapshot(
 		ctx,
 		snapfile.Path(),
-		memfile.Path(),
+		"/dev/null",
 	)
 	if err != nil {
-		return nil, fmt.Errorf("error creating snapshot: %w", err)
+		// If the snapshot creation with /dev/null fails (the FC is not our custom version),
+		// we try to create a snapshot with the memfile path.
+		createSnapshotFallbackErr := s.process.CreateSnapshot(
+			ctx,
+			snapfile.Path(),
+			memfile.Path(),
+		)
+		if createSnapshotFallbackErr != nil {
+			return nil, fmt.Errorf("error creating snapshot fallback: %w", createSnapshotFallbackErr)
+		}
+
+		// Delete the memfile as it is not used—we always use the process memory view.
+		closeErr := memfile.Close()
+		if closeErr != nil {
+			return nil, fmt.Errorf("error closing memfile: %w", closeErr)
+		}
 	}
 
 	// Gather data for postprocessing
@@ -719,17 +735,32 @@ func (s *Sandbox) Pause(
 		return nil, fmt.Errorf("failed to get original rootfs: %w", err)
 	}
 
+	pid, err := s.process.Pid()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get process pid: %w", err)
+	}
+
+	memoryMapping, err := s.memory.Mapping(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get memory mapping: %w", err)
+	}
+
+	memoryView, err := memory.NewView(pid, memoryMapping)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create memory view: %w", err)
+	}
+
 	// Start POSTPROCESSING
 	memfileDiff, memfileDiffHeader, err := pauseProcessMemory(
 		ctx,
 		buildID,
 		originalMemfile.Header(),
 		&MemoryDiffCreator{
-			memfile:    memfile,
+			memory:     memoryView,
 			dirtyPages: dirty.BitSet(),
 			blockSize:  originalMemfile.BlockSize(),
 			doneHook: func(context.Context) error {
-				return memfile.Close()
+				return memoryView.Close()
 			},
 		},
 	)

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

@@ -2,6 +2,8 @@ package memory
 
 import (
 	"fmt"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
 )
 
 type Mapping struct {
@@ -22,3 +24,37 @@ func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uint64, error) {
 
 	return 0, 0, fmt.Errorf("address %d not found in any mapping", hostVirtAddr)
 }
+
+// GetHostVirtRanges returns the host virtual addresses and sizes (ranges) that cover exactly the given [offset, offset+length) range in the host virtual address space.
+func (m *Mapping) GetHostVirtRanges(off int64, size int64) (hostVirtRanges []block.Range, err error) {
+	for n := int64(0); n < size; {
+		currentOff := off + n
+
+		region, err := m.getHostVirtRegion(currentOff)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get host virt mapping: %w", err)
+		}
+
+		start := region.shiftedHostVirtAddr(currentOff)
+		s := min(int64(region.endHostVirtAddr()-start), size-n)
+
+		r := block.NewRange(int64(start), s)
+
+		hostVirtRanges = append(hostVirtRanges, r)
+
+		n += r.Size
+	}
+
+	return hostVirtRanges, nil
+}
+
+// getHostVirtRegion returns the region that contains the given offset.
+func (m *Mapping) getHostVirtRegion(off int64) (*Region, error) {
+	for _, r := range m.Regions {
+		if off >= int64(r.Offset) && off < r.endOffset() {
+			return &r, nil
+		}
+	}
+
+	return nil, fmt.Errorf("offset %d not found in any mapping", off)
+}

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

@@ -11,6 +11,12 @@ type Region struct {
 	PageSize uintptr `json:"page_size_kib"`
 }
 
+// endOffset returns the end offset of the region in bytes.
+// The end offset is exclusive.
+func (r *Region) endOffset() int64 {
+	return int64(r.Offset + r.Size)
+}
+
 // endHostVirtAddr returns the end address of the region in host virtual address.
 // The end address is exclusive.
 func (r *Region) endHostVirtAddr() uintptr {
@@ -21,3 +27,8 @@ func (r *Region) endHostVirtAddr() uintptr {
 func (r *Region) shiftedOffset(addr uintptr) int64 {
 	return int64(addr - r.BaseHostVirtAddr + r.Offset)
 }
+
+// shiftedHostVirtAddr returns the host virtual address of the given offset in the region.
+func (r *Region) shiftedHostVirtAddr(off int64) uintptr {
+	return uintptr(off) + r.BaseHostVirtAddr - r.Offset
+}

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

@@ -0,0 +1,54 @@
+package memory
+
+import (
+	"fmt"
+	"io"
+	"os"
+)
+
+var _ io.ReaderAt = (*View)(nil)
+var _ io.Closer = (*View)(nil)
+
+// View exposes memory of the underlying process, with the mappings applied.
+type View struct {
+	m       *Mapping
+	procMem *os.File
+}
+
+func NewView(pid int, m *Mapping) (*View, error) {
+	fd, err := os.Open(fmt.Sprintf("/proc/%d/mem", pid))
+	if err != nil {
+		return nil, fmt.Errorf("failed to open memory file: %w", err)
+	}
+
+	return &View{
+		procMem: fd,
+		m:       m,
+	}, nil
+}
+
+// ReadAt reads data from the memory view at the given offset.
+// If this operation crosses a page boundary, it will read the data from the next page.
+//
+// If you try to read missing pages that are not yet faulted in via UFFD, this will return an error.
+func (v *View) ReadAt(d []byte, off int64) (n int, err error) {
+	ranges, err := v.m.GetHostVirtRanges(off, int64(len(d)))
+	if err != nil {
+		return 0, fmt.Errorf("failed to get host virt regions: %w", err)
+	}
+
+	for _, r := range ranges {
+		written, err := v.procMem.ReadAt(d[n:r.Size], r.Start)
+		if err != nil {
+			return 0, fmt.Errorf("failed to read at: %w", err)
+		}
+
+		n += written
+	}
+
+	return n, nil
+}
+
+func (v *View) Close() error {
+	return v.procMem.Close()
+}

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

@@ -0,0 +1,55 @@
+package memory
+
+import (
+	"bytes"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/testutils"
+)
+
+func TestView(t *testing.T) {
+	pagesize := uint64(4096)
+
+	data := testutils.RandomPages(pagesize, 128)
+
+	size, err := data.Size()
+	require.NoError(t, err)
+
+	memoryArea, memoryStart, unmap, err := testutils.NewPageMmap(uint64(size), pagesize)
+	require.NoError(t, err)
+
+	defer unmap()
+
+	n := copy(memoryArea[0:size], data.Content())
+	require.Equal(t, int(size), n)
+
+	m := NewMapping([]Region{
+		{
+			BaseHostVirtAddr: memoryStart,
+			Size:             uintptr(size),
+			Offset:           uintptr(0),
+			PageSize:         uintptr(pagesize),
+		},
+	})
+
+	pc, err := NewView(os.Getpid(), m)
+	require.NoError(t, err)
+
+	defer pc.Close()
+
+	for i := 0; i < int(size); i += int(pagesize) {
+		readBytes := make([]byte, pagesize)
+		_, err := pc.ReadAt(readBytes, int64(i))
+		require.NoError(t, err)
+
+		expectedBytes := data.Content()[i : i+int(pagesize)]
+
+		if !bytes.Equal(readBytes, expectedBytes) {
+			idx, want, got := testutils.FirstDifferentByte(readBytes, expectedBytes)
+			t.Fatalf("content mismatch: want %x, got %x at index %d", want, got, idx)
+		}
+	}
+}

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

@@ -4,13 +4,15 @@ import (
 	"context"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/memory"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 type MemoryBackend interface {
 	Dirty(ctx context.Context) (*block.Tracker, error)
 	// Disable switch the uffd to start serving empty pages.
 	Disable(ctx context.Context) error
+	Mapping(ctx context.Context) (*memory.Mapping, error)
 
 	Start(ctx context.Context, sandboxId string) error
 	Stop() error

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

@@ -6,6 +6,7 @@ import (
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/memory"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
@@ -47,6 +48,10 @@ func (m *NoopMemory) Start(context.Context, string) error {
 	return nil
 }
 
+func (m *NoopMemory) Mapping(context.Context) (*memory.Mapping, error) {
+	return nil, nil
+}
+
 func (m *NoopMemory) Stop() error {
 	return m.exit.SetSuccess()
 }

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

@@ -216,3 +216,12 @@ func (u *Uffd) Dirty(ctx context.Context) (*block.Tracker, error) {
 
 	return uffd.Dirty(ctx)
 }
+
+func (u *Uffd) Mapping(ctx context.Context) (*memory.Mapping, error) {
+	uffd, err := u.handler.WaitWithContext(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get uffd: %w", err)
+	}
+
+	return uffd.Mapping(), nil
+}

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

@@ -69,3 +69,7 @@ func (u *Userfaultfd) Dirty(ctx context.Context) (*block.Tracker, error) {
 
 	return u.dirty.Clone(), nil
 }
+
+func (u *Userfaultfd) Mapping() *memory.Mapping {
+	return u.ma
+}