From 0dadb586f2210b46610b5de024faf1944dd0bb50 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Thu, 7 Aug 2025 14:36:56 -0700 Subject: [PATCH 01/49] Wrap uffd fd in a struct with corresponding methods; Move the low level uffd handling to a subpackage --- .../sandbox/uffd/{handler.go => uffd.go} | 5 +- .../sandbox/uffd/userfaultfd/mocks_test.go | 100 ++++++++++++++++++ .../sandbox/uffd/{ => userfaultfd}/serve.go | 65 ++++-------- .../sandbox/uffd/userfaultfd/userfaultfd.go | 98 +++++++++++++++++ 4 files changed, 222 insertions(+), 46 deletions(-) rename packages/orchestrator/internal/sandbox/uffd/{handler.go => uffd.go} (96%) create mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go rename packages/orchestrator/internal/sandbox/uffd/{ => userfaultfd}/serve.go (71%) create mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go diff --git a/packages/orchestrator/internal/sandbox/uffd/handler.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go similarity index 96% rename from packages/orchestrator/internal/sandbox/uffd/handler.go rename to packages/orchestrator/internal/sandbox/uffd/uffd.go index 947ff168cf..bb77b8dd42 100644 --- a/packages/orchestrator/internal/sandbox/uffd/handler.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -17,6 +17,7 @@ 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/userfaultfd" "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/utils" ) @@ -141,10 +142,10 @@ 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]), false) 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)) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go new file mode 100644 index 0000000000..92b5835649 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go @@ -0,0 +1,100 @@ +package userfaultfd + +import ( + "math" + "syscall" + "unsafe" + + "golang.org/x/sys/unix" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +const pagesInTestData = 32 + +type mockMappings struct { + start uintptr + size int64 + pagesize int64 +} + +func newMockMappings(start uintptr, size, pagesize int64) *mockMappings { + return &mockMappings{ + start: start, + size: size, + pagesize: pagesize, + } +} + +func (m *mockMappings) GetRange(addr uintptr) (int64, int64, error) { + offset := addr - m.start + pagesize := m.pagesize + + return int64(offset), pagesize, nil +} + +type mockSlicer struct { + content []byte +} + +func newMockSlicer(content []byte) *mockSlicer { + return &mockSlicer{content: content} +} + +func (s *mockSlicer) Slice(offset, size int64) ([]byte, error) { + return s.content[offset : offset+size], nil +} + +func newMock4KPageMmap(size int64) ([]byte, uintptr) { + return newMockMmap(size, header.PageSize, 0) +} + +func newMock2MPageMmap(size int64) ([]byte, uintptr) { + return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) +} + +func newMockMmap(size, pagesize int64, 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])) +} + +func repeatToSize(src []byte, size int64) []byte { + if len(src) == 0 || size <= 0 { + return nil + } + + dst := make([]byte, size) + for i := 0; i < int(size); i += len(src) { + end := i + len(src) + if end > int(size) { + end = int(size) + } + copy(dst[i:end], src[:end-i]) + } + + return dst +} + +func prepareTestData(pagesize int64) (data *mockSlicer, size int64) { + 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 +} diff --git a/packages/orchestrator/internal/sandbox/uffd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go similarity index 71% rename from packages/orchestrator/internal/sandbox/uffd/serve.go rename to packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index c5866b52dd..f1332b53fa 100644 --- a/packages/orchestrator/internal/sandbox/uffd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -1,4 +1,4 @@ -package uffd +package userfaultfd import ( "context" @@ -14,29 +14,18 @@ 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/userfaultfd" ) -var ErrUnexpectedEventType = errors.New("unexpected event type") - -type GuestRegionUffdMapping struct { - BaseHostVirtAddr uintptr `json:"base_host_virt_addr"` - Size uintptr `json:"size"` - Offset uintptr `json:"offset"` - PageSize uintptr `json:"page_size_kib"` -} - -func Serve( +func (u *userfaultfd) Serve( ctx context.Context, - uffd int, mappings mapping.Mappings, src block.Slicer, fdExit *fdexit.FdExit, logger *zap.Logger, ) error { pollFds := []unix.PollFd{ - {Fd: int32(uffd), Events: unix.POLLIN}, - {Fd: fdExit.Reader(), Events: unix.POLLIN}, + {Fd: int32(u.fd), Events: unix.POLLIN}, + {Fd: int32(fdExit.Reader()), Events: unix.POLLIN}, } var eg errgroup.Group @@ -94,10 +83,10 @@ outerLoop: continue } - buf := make([]byte, unsafe.Sizeof(userfaultfd.UffdMsg{})) + buf := make([]byte, unsafe.Sizeof(UffdMsg{})) for { - n, err := syscall.Read(uffd, buf) + n, err := syscall.Read(int(u.fd), buf) if err == syscall.EINTR { logger.Debug("uffd: interrupted read, reading again") @@ -121,17 +110,17 @@ outerLoop: return fmt.Errorf("failed to read: %w", err) } - msg := *(*userfaultfd.UffdMsg)(unsafe.Pointer(&buf[0])) - if userfaultfd.GetMsgEvent(&msg) != userfaultfd.UFFD_EVENT_PAGEFAULT { - logger.Error("UFFD serve unexpected event type", zap.Any("event_type", userfaultfd.GetMsgEvent(&msg))) + msg := *(*UffdMsg)(unsafe.Pointer(&buf[0])) + if GetMsgEvent(&msg) != UFFD_EVENT_PAGEFAULT { + logger.Error("UFFD serve unexpected event type", zap.Any("event_type", GetMsgEvent(&msg))) return ErrUnexpectedEventType } - arg := userfaultfd.GetMsgArg(&msg) - pagefault := (*(*userfaultfd.UffdPagefault)(unsafe.Pointer(&arg[0]))) + arg := GetMsgArg(&msg) + pagefault := (*(*UffdPagefault)(unsafe.Pointer(&arg[0]))) - addr := userfaultfd.GetPagefaultAddress(&pagefault) + addr := GetPagefaultAddress(&pagefault) offset, pagesize, err := mappings.GetRange(uintptr(addr)) if err != nil { @@ -164,30 +153,18 @@ outerLoop: return fmt.Errorf("failed to read from source: %w", joinedErr) } - cpy := userfaultfd.NewUffdioCopy( - b, - addr&^userfaultfd.CULong(pagesize-1), - userfaultfd.CULong(pagesize), - 0, - 0, - ) - - if _, _, errno := syscall.Syscall( - syscall.SYS_IOCTL, - uintptr(uffd), - userfaultfd.UFFDIO_COPY, - uintptr(unsafe.Pointer(&cpy)), - ); errno != 0 { - if errno == unix.EEXIST { - logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) - - // Page is already mapped - return nil - } + err = u.Copy(addr, b, pagesize) + if err == unix.EEXIST { + logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) + // Page is already mapped + return nil + } + + if err != nil { signalErr := fdExit.SignalExit() - joinedErr := errors.Join(errno, signalErr) + joinedErr := errors.Join(err, signalErr) logger.Error("UFFD serve uffdio copy error", zap.Error(joinedErr)) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go new file mode 100644 index 0000000000..1bef93c899 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -0,0 +1,98 @@ +package userfaultfd + +import ( + "errors" + "fmt" + "syscall" + "unsafe" +) + +var ErrUnexpectedEventType = errors.New("unexpected event type") + +type userfaultfd struct { + fd uintptr + copyMode CULong +} + +// flags: syscall.O_CLOEXEC|syscall.O_NONBLOCK +func NewUserfaultfd(flags uintptr, wp bool) (*userfaultfd, error) { + uffd, _, errno := syscall.Syscall(NR_userfaultfd, flags, 0, 0) + if errno != 0 { + return nil, fmt.Errorf("userfaultfd syscall failed: %v", errno) + } + + return NewUserfaultfdFromFd(uffd, wp), nil +} + +func NewUserfaultfdFromFd(fd uintptr, wp bool) *userfaultfd { + copyMode := CULong(0) + if wp { + copyMode = UFFDIO_COPY_MODE_WP + } + + return &userfaultfd{ + fd: fd, + copyMode: copyMode, + } +} + +// features: UFFD_FEATURE_MISSING_HUGETLBFS +// This is already called by the FC +func (u *userfaultfd) ConfigureApi(features CULong) error { + api := NewUffdioAPI(UFFD_API, features) + ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_API, uintptr(unsafe.Pointer(&api))) + if errno != 0 { + return fmt.Errorf("UFFDIO_API ioctl failed: %v (ret=%d)", errno, ret) + } + + return nil +} + +// mode: UFFDIO_REGISTER_MODE_WP|UFFDIO_REGISTER_MODE_MISSING +// This is already called by the FC, but only with the UFFDIO_REGISTER_MODE_MISSING +// We need to call it with UFFDIO_REGISTER_MODE_WP when we use both missing and wp +func (h *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { + register := NewUffdioRegister(CULong(addr), CULong(size), mode) + + ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_REGISTER, uintptr(unsafe.Pointer(®ister))) + if errno != 0 { + return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %v (ret=%d)", errno, ret) + } + + return nil +} + +func (h *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error { + register := NewUffdioWriteProtect(CULong(addr), CULong(size), mode) + + ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_WRITEPROTECT, uintptr(unsafe.Pointer(®ister))) + if errno != 0 { + return fmt.Errorf("UFFDIO_WRITEPROTECT ioctl failed: %v (ret=%d)", errno, ret) + } + + return nil +} + +func (h *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { + return h.writeProtect(addr, size, 0) +} + +func (h *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { + return h.writeProtect(addr, size, UFFDIO_WRITEPROTECT_MODE_WP) +} + +// mode: UFFDIO_COPY_MODE_WP +// When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page +func (h *userfaultfd) Copy(addr CULong, data []byte, pagesize int64) error { + cpy := NewUffdioCopy(data, addr&^CULong(pagesize-1), CULong(pagesize), h.copyMode, 0) + + if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { + return errno + } + + return nil +} + +func (h *userfaultfd) Close() error { + return syscall.Close(int(h.fd)) +} From 34ccfc5ad2ccd4f45267b94f192925ba1367e47c Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Thu, 7 Aug 2025 14:45:26 -0700 Subject: [PATCH 02/49] Unify receiver name --- .../sandbox/uffd/userfaultfd/userfaultfd.go | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 1bef93c899..2893e26925 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -51,10 +51,10 @@ func (u *userfaultfd) ConfigureApi(features CULong) error { // mode: UFFDIO_REGISTER_MODE_WP|UFFDIO_REGISTER_MODE_MISSING // This is already called by the FC, but only with the UFFDIO_REGISTER_MODE_MISSING // We need to call it with UFFDIO_REGISTER_MODE_WP when we use both missing and wp -func (h *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { +func (u *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { register := NewUffdioRegister(CULong(addr), CULong(size), mode) - ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_REGISTER, uintptr(unsafe.Pointer(®ister))) + ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_REGISTER, uintptr(unsafe.Pointer(®ister))) if errno != 0 { return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %v (ret=%d)", errno, ret) } @@ -62,10 +62,10 @@ func (h *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { return nil } -func (h *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error { +func (u *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error { register := NewUffdioWriteProtect(CULong(addr), CULong(size), mode) - ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_WRITEPROTECT, uintptr(unsafe.Pointer(®ister))) + ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_WRITEPROTECT, uintptr(unsafe.Pointer(®ister))) if errno != 0 { return fmt.Errorf("UFFDIO_WRITEPROTECT ioctl failed: %v (ret=%d)", errno, ret) } @@ -73,26 +73,26 @@ func (h *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error return nil } -func (h *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { - return h.writeProtect(addr, size, 0) +func (u *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { + return u.writeProtect(addr, size, 0) } -func (h *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { - return h.writeProtect(addr, size, UFFDIO_WRITEPROTECT_MODE_WP) +func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { + return u.writeProtect(addr, size, UFFDIO_WRITEPROTECT_MODE_WP) } // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page -func (h *userfaultfd) Copy(addr CULong, data []byte, pagesize int64) error { - cpy := NewUffdioCopy(data, addr&^CULong(pagesize-1), CULong(pagesize), h.copyMode, 0) +func (u *userfaultfd) Copy(addr CULong, data []byte, pagesize int64) error { + cpy := NewUffdioCopy(data, addr&^CULong(pagesize-1), CULong(pagesize), u.copyMode, 0) - if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, h.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { + if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { return errno } return nil } -func (h *userfaultfd) Close() error { - return syscall.Close(int(h.fd)) +func (u *userfaultfd) Close() error { + return syscall.Close(int(u.fd)) } From 7312bcacd2185769faac9b491b1b2d02842a496f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Thu, 7 Aug 2025 14:50:48 -0700 Subject: [PATCH 03/49] [WIP] Add uffd tests --- .../userfaultfd_cross_process_test.go | 138 +++++++++++ .../uffd/userfaultfd/userfaultfd_test.go | 230 ++++++++++++++++++ 2 files changed, 368 insertions(+) create mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go create mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go new file mode 100644 index 0000000000..9ab1589cb7 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -0,0 +1,138 @@ +package userfaultfd + +// This tests is creating uffd in a process and handling the page faults in another process. +// It also tests reregistering the uffd with the additional wp flag in the another process (in "orchestrator") after registering the missing handler already (in "FC"), +// simulating the case we have with the write protection being set up after FC already registered the uffd. + +import ( + "bytes" + "fmt" + "os" + "os/exec" + "os/signal" + "strconv" + "syscall" + "testing" + "time" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "go.uber.org/zap" +) + +var ( + // Data shared across the testing processes to check if the served and received data is the same + testCrossProcessPageSize = int64(header.HugepageSize) + testCrossProcessData, testCrossProcessSize = prepareTestData(testCrossProcessPageSize) +) + +// Main process, FC in our case +func TestCrossProcessDoubleRegistration(t *testing.T) { + memoryArea, memoryStart := newMock2MPageMmap(testCrossProcessSize) + + uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + defer uffd.Close() + + err = uffd.ConfigureApi(0) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + err = uffd.Register(memoryStart, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING) + if err != nil { + t.Fatal("failed to register memory", err) + } + + sigc := make(chan os.Signal, 1) + signal.Notify(sigc, syscall.SIGUSR1) + + cmd := exec.Command(os.Args[0], "-test.run=TestHelperProcess") + cmd.Env = append(os.Environ(), "GO_TEST_HELPER_PROCESS=1") + cmd.Env = append(cmd.Env, fmt.Sprintf("GO_MMAP_START=%d", memoryStart)) + + // Passing the fd to the child process + uffdFile := os.NewFile(uffd.fd, "userfaultfd") + + cmd.ExtraFiles = []*os.File{uffdFile} + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + if err := cmd.Start(); err != nil { + t.Fatalf("failed to start helper: %v", err) + } + defer func() { + cmd.Process.Kill() + cmd.Wait() + }() + + <-sigc + fmt.Println("child signaled ready") + + data := testCrossProcessData + + servedContent, err := data.Slice(0, testCrossProcessPageSize) + if err != nil { + t.Fatal("cannot read content", err) + } + + if !bytes.Equal(memoryArea[0:testCrossProcessPageSize], servedContent) { + t.Fatal("content mismatch", string(servedContent)) + } +} + +// Secondary process, orchestrator in our case +func TestHelperProcess(t *testing.T) { + if os.Getenv("GO_TEST_HELPER_PROCESS") != "1" { + return + } + + mmapStart := os.Getenv("GO_MMAP_START") + startRaw, err := strconv.Atoi(mmapStart) + if err != nil { + fmt.Print("exit parsing mmap start", err) + os.Exit(1) + } + + start := uintptr(startRaw) + + uffdFile := os.NewFile(uintptr(3), "userfaultfd") + uffd := NewUserfaultfdFromFd(uffdFile.Fd(), true) + + // done in the FC + // Check: The reregistration works + err = uffd.Register(start, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + if err != nil { + fmt.Print("exit registering uffd", err) + os.Exit(1) + } + + fmt.Println("after register") + + ppid := os.Getppid() + syscall.Kill(ppid, syscall.SIGUSR1) + + mappings := newMockMappings(start, testCrossProcessSize, testCrossProcessPageSize) + + fdExit, err := fdexit.New() + if err != nil { + fmt.Print("exit creating fd exit", err) + os.Exit(1) + } + defer fdExit.Close() + + data := testCrossProcessData + + go func() { + err := uffd.Serve(mappings, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) + } + }() + + time.Sleep(10 * time.Second) + + os.Exit(0) +} diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go new file mode 100644 index 0000000000..70d492d5f5 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -0,0 +1,230 @@ +package userfaultfd + +import ( + "bytes" + "fmt" + "syscall" + "testing" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "go.uber.org/zap" +) + +func TestUffdMissing(t *testing.T) { + pagesize := int64(header.PageSize) + data, size := prepareTestData(pagesize) + + uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, false) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + defer uffd.Close() + + err = uffd.ConfigureApi(0) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + memoryArea, memoryStart := newMock4KPageMmap(size) + + err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING) + if err != nil { + t.Fatal("failed to register memory", err) + } + + mappings := newMockMappings(memoryStart, size, pagesize) + + fdExit, err := fdexit.New() + if err != nil { + t.Fatal("failed to create fd exit", err) + } + defer fdExit.Close() + + go func() { + err := uffd.Serve(mappings, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdMissing] failed to serve uffd", err) + } + }() + + d, err := data.Slice(0, pagesize) + if err != nil { + t.Fatal("cannot read content", err) + } + + if !bytes.Equal(memoryArea[0:pagesize], d) { + t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) + } +} + +func TestUffdWriteProtect(t *testing.T) { + pagesize := int64(header.PageSize) + data, size := prepareTestData(pagesize) + + uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + defer uffd.Close() + + err = uffd.ConfigureApi(0) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + memoryArea, memoryStart := newMock4KPageMmap(size) + + err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_WP) + if err != nil { + t.Fatal("failed to register memory", err) + } + + err = uffd.AddWriteProtection(memoryStart, uint64(size)) + if err != nil { + t.Fatal("failed to write protect memory", err) + } + + fdExit, err := fdexit.New() + if err != nil { + t.Fatal("failed to create fd exit", err) + } + defer fdExit.Close() + + mappings := newMockMappings(memoryStart, size, pagesize) + + go func() { + err := uffd.Serve(mappings, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) + } + }() + + memoryArea[0] = 'A' + + // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd +} + +func TestUffdWriteProtectWithMissing(t *testing.T) { + pagesize := int64(header.PageSize) + + data, size := prepareTestData(pagesize) + + uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + defer uffd.Close() + + err = uffd.ConfigureApi(0) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + memoryArea, memoryStart := newMock4KPageMmap(size) + + err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + if err != nil { + t.Fatal("failed to register memory", err) + } + + err = uffd.AddWriteProtection(memoryStart, uint64(size)) + if err != nil { + t.Fatal("failed to write protect memory", err) + } + + mappings := newMockMappings(memoryStart, size, pagesize) + + fdExit, err := fdexit.New() + if err != nil { + t.Fatal("failed to create fd exit", err) + } + defer fdExit.Close() + + go func() { + err := uffd.Serve(mappings, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) + } + }() + + d, err := data.Slice(0, pagesize) + if err != nil { + t.Fatal("cannot read content", err) + } + + if !bytes.Equal(memoryArea[0:pagesize], d) { + t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) + } + + memoryArea[0] = 'A' + + // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd +} + +// We are trying to simulate registering the missing handler in the FC and then registering the missing+wp handler again in the orchestrator +func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { + pagesize := int64(header.PageSize) + data, size := prepareTestData(pagesize) + + uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + defer uffd.Close() + + err = uffd.ConfigureApi(0) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + memoryArea, memoryStart := newMock4KPageMmap(size) + + // done in the FC + err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING) + if err != nil { + t.Fatal("failed to register memory", err) + } + + // TODO: Can we reregister after triggering missing and still properly handle such a page later? + + // done little later in the orchestrator + // both flags needs to be present + err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + if err != nil { + t.Fatal("failed to register memory", err) + } + + err = uffd.AddWriteProtection(memoryStart, uint64(size)) + if err != nil { + t.Fatal("failed to write protect memory", err) + } + + mappings := newMockMappings(memoryStart, size, pagesize) + + fdExit, err := fdexit.New() + if err != nil { + t.Fatal("failed to create fd exit", err) + } + defer fdExit.Close() + + go func() { + err := uffd.Serve(mappings, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) + } + }() + + d, err := data.Slice(0, pagesize) + if err != nil { + t.Fatal("cannot read content", err) + } + + if !bytes.Equal(memoryArea[0:pagesize], d) { + t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) + } + + memoryArea[0] = 'A' + + // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd +} From 3c75576bb30511a424a98c3ca9c070c8ab407216 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Thu, 7 Aug 2025 16:49:20 -0700 Subject: [PATCH 04/49] Add disabled wp --- .../internal/sandbox/uffd/uffd.go | 24 ++++++++++++++++- .../sandbox/uffd/userfaultfd/serve.go | 2 +- .../sandbox/uffd/userfaultfd/userfaultfd.go | 26 +++++++++---------- .../userfaultfd_cross_process_test.go | 8 +++--- .../uffd/userfaultfd/userfaultfd_test.go | 16 ++++++------ 5 files changed, 49 insertions(+), 27 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index bb77b8dd42..40b66b2550 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -40,6 +40,8 @@ type Uffd struct { memfile *block.TrackedSliceDevice socketPath string + + writeProtection bool } var _ MemoryBackend = (*Uffd)(nil) @@ -142,7 +144,7 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { return fmt.Errorf("expected 1 fd: found %d", len(fds)) } - uffd := userfaultfd.NewUserfaultfdFromFd(uintptr(fds[0]), false) + uffd := userfaultfd.NewUserfaultfdFromFd(uintptr(fds[0])) defer func() { closeErr := uffd.Close() @@ -151,6 +153,26 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { } }() + if u.writeProtection { + for _, region := range m { + // Register the WP. It is possible that the memory region was already registered (with missing pages in FC), but registering it again with bigger subset should merge these. + err := uffd.Register( + region.BaseHostVirtAddr, + uint64(region.Size), + userfaultfd.UFFDIO_REGISTER_MODE_WP|userfaultfd.UFFDIO_REGISTER_MODE_MISSING, + ) + + // Add write protection to the regions provided by the UFFD + err = uffd.AddWriteProtection( + region.BaseHostVirtAddr, + uint64(region.Size), + ) + if err != nil { + return fmt.Errorf("failed to add write protection to region %d-%d", region.Offset, region.Offset+region.Size) + } + } + } + u.readyCh <- struct{}{} err = Serve( diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index f1332b53fa..44ba0e991b 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -153,7 +153,7 @@ outerLoop: return fmt.Errorf("failed to read from source: %w", joinedErr) } - err = u.Copy(addr, b, pagesize) + err = u.copy(addr, b, pagesize) if err == unix.EEXIST { logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 2893e26925..26ce999ac8 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -15,30 +15,25 @@ type userfaultfd struct { } // flags: syscall.O_CLOEXEC|syscall.O_NONBLOCK -func NewUserfaultfd(flags uintptr, wp bool) (*userfaultfd, error) { +func newUserfaultfd(flags uintptr) (*userfaultfd, error) { uffd, _, errno := syscall.Syscall(NR_userfaultfd, flags, 0, 0) if errno != 0 { return nil, fmt.Errorf("userfaultfd syscall failed: %v", errno) } - return NewUserfaultfdFromFd(uffd, wp), nil + return NewUserfaultfdFromFd(uffd), nil } -func NewUserfaultfdFromFd(fd uintptr, wp bool) *userfaultfd { - copyMode := CULong(0) - if wp { - copyMode = UFFDIO_COPY_MODE_WP - } - +// NewUserfaultfdFromFd creates a new userfaultfd instance with optional configuration. +func NewUserfaultfdFromFd(fd uintptr) *userfaultfd { return &userfaultfd{ - fd: fd, - copyMode: copyMode, + fd: fd, } } // features: UFFD_FEATURE_MISSING_HUGETLBFS // This is already called by the FC -func (u *userfaultfd) ConfigureApi(features CULong) error { +func (u *userfaultfd) configureApi(features CULong) error { api := NewUffdioAPI(UFFD_API, features) ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_API, uintptr(unsafe.Pointer(&api))) if errno != 0 { @@ -59,6 +54,11 @@ func (u *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %v (ret=%d)", errno, ret) } + // If we register with write protection automatically use the copy for missing pages without disabling the WP on that page. + if mode&UFFDIO_REGISTER_MODE_WP != 0 { + u.copyMode = UFFDIO_COPY_MODE_WP + } + return nil } @@ -73,7 +73,7 @@ func (u *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error return nil } -func (u *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { +func (u *userfaultfd) removeWriteProtection(addr uintptr, size uint64) error { return u.writeProtect(addr, size, 0) } @@ -83,7 +83,7 @@ func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page -func (u *userfaultfd) Copy(addr CULong, data []byte, pagesize int64) error { +func (u *userfaultfd) copy(addr CULong, data []byte, pagesize int64) error { cpy := NewUffdioCopy(data, addr&^CULong(pagesize-1), CULong(pagesize), u.copyMode, 0) if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 9ab1589cb7..594581757e 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -30,18 +30,18 @@ var ( func TestCrossProcessDoubleRegistration(t *testing.T) { memoryArea, memoryStart := newMock2MPageMmap(testCrossProcessSize) - uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) if err != nil { t.Fatal("failed to create userfaultfd", err) } defer uffd.Close() - err = uffd.ConfigureApi(0) + err = uffd.configureApi(0) if err != nil { t.Fatal("failed to configure uffd api", err) } - err = uffd.Register(memoryStart, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING) + err = uffd.register(memoryStart, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING) if err != nil { t.Fatal("failed to register memory", err) } @@ -103,7 +103,7 @@ func TestHelperProcess(t *testing.T) { // done in the FC // Check: The reregistration works - err = uffd.Register(start, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + err = uffd.register(start, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { fmt.Print("exit registering uffd", err) os.Exit(1) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 70d492d5f5..07f2c3eda7 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -15,13 +15,13 @@ func TestUffdMissing(t *testing.T) { pagesize := int64(header.PageSize) data, size := prepareTestData(pagesize) - uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, false) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { t.Fatal("failed to create userfaultfd", err) } defer uffd.Close() - err = uffd.ConfigureApi(0) + err = uffd.configureApi(0) if err != nil { t.Fatal("failed to configure uffd api", err) } @@ -62,13 +62,13 @@ func TestUffdWriteProtect(t *testing.T) { pagesize := int64(header.PageSize) data, size := prepareTestData(pagesize) - uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { t.Fatal("failed to create userfaultfd", err) } defer uffd.Close() - err = uffd.ConfigureApi(0) + err = uffd.configureApi(0) if err != nil { t.Fatal("failed to configure uffd api", err) } @@ -110,13 +110,13 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { data, size := prepareTestData(pagesize) - uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { t.Fatal("failed to create userfaultfd", err) } defer uffd.Close() - err = uffd.ConfigureApi(0) + err = uffd.configureApi(0) if err != nil { t.Fatal("failed to configure uffd api", err) } @@ -167,13 +167,13 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { pagesize := int64(header.PageSize) data, size := prepareTestData(pagesize) - uffd, err := NewUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { t.Fatal("failed to create userfaultfd", err) } defer uffd.Close() - err = uffd.ConfigureApi(0) + err = uffd.configureApi(0) if err != nil { t.Fatal("failed to configure uffd api", err) } From b67f0f25a87dd12fdcb4286e035ef251847a6c63 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 8 Aug 2025 00:14:21 -0700 Subject: [PATCH 05/49] [WIP] WP --- .../sandbox/uffd/mapping/firecracker.go | 8 +- .../internal/sandbox/uffd/mapping/mapping.go | 2 +- .../internal/sandbox/uffd/uffd.go | 54 +++++++----- .../sandbox/uffd/userfaultfd/constants.go | 12 +-- .../sandbox/uffd/userfaultfd/mocks_test.go | 28 +++--- .../sandbox/uffd/userfaultfd/serve.go | 88 ++++++++++++------- .../sandbox/uffd/userfaultfd/userfaultfd.go | 12 +-- .../userfaultfd_cross_process_test.go | 12 +-- .../uffd/userfaultfd/userfaultfd_test.go | 30 +++---- 9 files changed, 130 insertions(+), 116 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go index dbff8bbd7a..48ee4f2db4 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go @@ -11,21 +11,21 @@ type GuestRegionUffdMapping struct { PageSize uintptr `json:"page_size_kib"` } -func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) int64 { - return int64(m.Offset + addr - m.BaseHostVirtAddr) +func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) uint64 { + return uint64(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) { +func (m FcMappings) GetRange(addr uintptr) (uint64, uint64, 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 m.relativeOffset(addr), uint64(m.PageSize), nil } return 0, 0, fmt.Errorf("address %d not found in any mapping", addr) diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go index ffa8c4a6fa..33f9191f36 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go @@ -1,5 +1,5 @@ package mapping type Mappings interface { - GetRange(addr uintptr) (offset int64, pagesize int64, err error) + GetRange(addr uintptr) (offset uint64, pagesize uint64, err error) } diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index 40b66b2550..b1c2c4f13f 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -58,11 +58,13 @@ 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(), + exitCh: make(chan error, 1), + readyCh: make(chan struct{}, 1), + fdExit: fdExit, + memfile: trackedMemfile, + socketPath: socketPath, + writeProtection: true, }, nil } @@ -153,25 +155,28 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { } }() - if u.writeProtection { - for _, region := range m { - // Register the WP. It is possible that the memory region was already registered (with missing pages in FC), but registering it again with bigger subset should merge these. - err := uffd.Register( - region.BaseHostVirtAddr, - uint64(region.Size), - userfaultfd.UFFDIO_REGISTER_MODE_WP|userfaultfd.UFFDIO_REGISTER_MODE_MISSING, - ) - - // Add write protection to the regions provided by the UFFD - err = uffd.AddWriteProtection( - region.BaseHostVirtAddr, - uint64(region.Size), - ) - if err != nil { - return fmt.Errorf("failed to add write protection to region %d-%d", region.Offset, region.Offset+region.Size) - } - } - } + // if u.writeProtection { + // for _, region := range m { + // // Register the WP. It is possible that the memory region was already registered (with missing pages in FC), but registering it again with bigger subset should merge these. + // err := uffd.Register( + // region.Offset+region.BaseHostVirtAddr, + // uint64(region.Size), + // userfaultfd.UFFDIO_REGISTER_MODE_WP|userfaultfd.UFFDIO_REGISTER_MODE_MISSING, + // ) + // if err != nil { + // return fmt.Errorf("failed to reregister memory region with write protection %d-%d", region.Offset, region.Offset+region.Size) + // } + + // // Add write protection to the regions provided by the UFFD + // err = uffd.AddWriteProtection( + // region.Offset+region.BaseHostVirtAddr, + // uint64(region.Size), + // ) + // if err != nil { + // return fmt.Errorf("failed to add write protection to region %d-%d", region.Offset, region.Offset+region.Size) + // } + // } + // } u.readyCh <- struct{}{} @@ -184,6 +189,7 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { zap.L().With(logger.WithSandboxID(sandboxId)), ) if err != nil { + fmt.Fprintf(os.Stderr, "serve error %v", err) return fmt.Errorf("failed handling uffd: %w", err) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index 63a3b3e71e..070379ce7b 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -104,14 +104,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) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go index 92b5835649..5763a1c4cc 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go @@ -14,11 +14,11 @@ const pagesInTestData = 32 type mockMappings struct { start uintptr - size int64 - pagesize int64 + size uint64 + pagesize uint64 } -func newMockMappings(start uintptr, size, pagesize int64) *mockMappings { +func newMockMappings(start uintptr, size, pagesize uint64) *mockMappings { return &mockMappings{ start: start, size: size, @@ -26,11 +26,11 @@ func newMockMappings(start uintptr, size, pagesize int64) *mockMappings { } } -func (m *mockMappings) GetRange(addr uintptr) (int64, int64, error) { +func (m *mockMappings) GetRange(addr uintptr) (uint64, uint64, error) { offset := addr - m.start pagesize := m.pagesize - return int64(offset), pagesize, nil + return uint64(offset), pagesize, nil } type mockSlicer struct { @@ -45,15 +45,15 @@ func (s *mockSlicer) Slice(offset, size int64) ([]byte, error) { return s.content[offset : offset+size], nil } -func newMock4KPageMmap(size int64) ([]byte, uintptr) { +func newMock4KPageMmap(size uint64) ([]byte, uintptr) { return newMockMmap(size, header.PageSize, 0) } -func newMock2MPageMmap(size int64) ([]byte, uintptr) { +func newMock2MPageMmap(size uint64) ([]byte, uintptr) { return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) } -func newMockMmap(size, pagesize int64, flags int) ([]byte, uintptr) { +func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr) { l := int(math.Ceil(float64(size)/float64(pagesize)) * float64(pagesize)) b, err := syscall.Mmap( -1, @@ -69,16 +69,16 @@ func newMockMmap(size, pagesize int64, flags int) ([]byte, uintptr) { return b, uintptr(unsafe.Pointer(&b[0])) } -func repeatToSize(src []byte, size int64) []byte { +func repeatToSize(src []byte, size uint64) []byte { if len(src) == 0 || size <= 0 { return nil } dst := make([]byte, size) - for i := 0; i < int(size); i += len(src) { - end := i + len(src) - if end > int(size) { - end = int(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]) } @@ -86,7 +86,7 @@ func repeatToSize(src []byte, size int64) []byte { return dst } -func prepareTestData(pagesize int64) (data *mockSlicer, size int64) { +func prepareTestData(pagesize uint64) (data *mockSlicer, size uint64) { size = pagesize * pagesInTestData data = newMockSlicer( diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 44ba0e991b..bd014298e4 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "os" "syscall" "unsafe" @@ -23,6 +24,7 @@ func (u *userfaultfd) Serve( fdExit *fdexit.FdExit, logger *zap.Logger, ) error { + defer fmt.Fprintf(os.Stderr, "exiting serve >>>>>>>>>>>>") pollFds := []unix.PollFd{ {Fd: int32(u.fd), Events: unix.POLLIN}, {Fd: int32(fdExit.Reader()), Events: unix.POLLIN}, @@ -30,7 +32,8 @@ func (u *userfaultfd) Serve( var eg errgroup.Group - missingPagesBeingHandled := map[int64]struct{}{} + missingPagesBeingHandled := map[uint64]struct{}{} + writePagesBeingHandled := map[uint64]struct{}{} outerLoop: for { @@ -122,56 +125,75 @@ outerLoop: addr := GetPagefaultAddress(&pagefault) - offset, pagesize, err := mappings.GetRange(uintptr(addr)) + offset, pagesize, err := mappings.GetRange(addr) if err != nil { logger.Error("UFFD serve get mapping error", zap.Error(err)) return fmt.Errorf("failed to map: %w", err) } - if _, ok := missingPagesBeingHandled[offset]; ok { - continue - } + switch { + case pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0: + if _, ok := writePagesBeingHandled[offset]; ok { + continue + } - missingPagesBeingHandled[offset] = struct{}{} + writePagesBeingHandled[offset] = struct{}{} - eg.Go(func() error { - defer func() { - if r := recover(); r != nil { - logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) - } - }() + fallthrough + case pagefault.flags == 0: + if _, ok := missingPagesBeingHandled[offset]; ok { + continue + } - b, err := src.Slice(ctx, offset, pagesize) - if err != nil { - signalErr := fdExit.SignalExit() + missingPagesBeingHandled[offset] = struct{}{} - joinedErr := errors.Join(err, signalErr) + fmt.Fprintf(os.Stderr, "[%d] writes without protection: %d, missing read: %d\n", offset/pagesize, len(writePagesBeingHandled), len(missingPagesBeingHandled)) - logger.Error("UFFD serve slice error", zap.Error(joinedErr)) + // TODO: Copy the page ahead as we already fetch a chunk (increase it) + // TODO: Prefetch the top of the memory downward, the start of the memory up - return fmt.Errorf("failed to read from source: %w", joinedErr) - } + eg.Go(func() error { + defer func() { + if r := recover(); r != nil { + logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) + } + }() - err = u.copy(addr, b, pagesize) - if err == unix.EEXIST { - logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) + b, sliceErr := src.Slice(int64(offset), int64(pagesize)) + if sliceErr != nil { - // Page is already mapped - return nil - } + signalErr := fdExit.SignalExit() + + joinedErr := errors.Join(sliceErr, signalErr) - if err != nil { - signalErr := fdExit.SignalExit() + logger.Error("UFFD serve slice error", zap.Error(joinedErr)) - joinedErr := errors.Join(err, signalErr) + return fmt.Errorf("failed to read from source: %w", joinedErr) + } - logger.Error("UFFD serve uffdio copy error", zap.Error(joinedErr)) + copyErr := u.copy(addr, b, pagesize, 0) + if copyErr == unix.EEXIST { + logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) - return fmt.Errorf("failed uffdio copy %w", joinedErr) - } + // Page is already mapped + return nil + } - return nil - }) + if copyErr != nil { + signalErr := fdExit.SignalExit() + + joinedErr := errors.Join(copyErr, signalErr) + + logger.Error("UFFD serve uffdio copy error", zap.Error(joinedErr)) + + return fmt.Errorf("failed uffdio copy %w", joinedErr) + } + + return nil + }) + default: + return fmt.Errorf("invalid pagefault flags %d for address %d", pagefault.flags, addr) + } } } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 26ce999ac8..0efeb55745 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -10,8 +10,7 @@ import ( var ErrUnexpectedEventType = errors.New("unexpected event type") type userfaultfd struct { - fd uintptr - copyMode CULong + fd uintptr } // flags: syscall.O_CLOEXEC|syscall.O_NONBLOCK @@ -54,11 +53,6 @@ func (u *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %v (ret=%d)", errno, ret) } - // If we register with write protection automatically use the copy for missing pages without disabling the WP on that page. - if mode&UFFDIO_REGISTER_MODE_WP != 0 { - u.copyMode = UFFDIO_COPY_MODE_WP - } - return nil } @@ -83,8 +77,8 @@ func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page -func (u *userfaultfd) copy(addr CULong, data []byte, pagesize int64) error { - cpy := NewUffdioCopy(data, addr&^CULong(pagesize-1), CULong(pagesize), u.copyMode, 0) +func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULong) error { + cpy := NewUffdioCopy(data, CULong(addr)&^CULong(pagesize-1), CULong(pagesize), mode, 0) if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { return errno diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 594581757e..94eabc56f8 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -22,7 +22,7 @@ import ( var ( // Data shared across the testing processes to check if the served and received data is the same - testCrossProcessPageSize = int64(header.HugepageSize) + testCrossProcessPageSize = uint64(header.HugepageSize) testCrossProcessData, testCrossProcessSize = prepareTestData(testCrossProcessPageSize) ) @@ -30,7 +30,7 @@ var ( func TestCrossProcessDoubleRegistration(t *testing.T) { memoryArea, memoryStart := newMock2MPageMmap(testCrossProcessSize) - uffd, err := newUserfaultfd(syscall.O_CLOEXEC|syscall.O_NONBLOCK, true) + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { t.Fatal("failed to create userfaultfd", err) } @@ -41,7 +41,7 @@ func TestCrossProcessDoubleRegistration(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - err = uffd.register(memoryStart, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING) + err = uffd.Register(memoryStart, testCrossProcessSize, UFFDIO_REGISTER_MODE_MISSING) if err != nil { t.Fatal("failed to register memory", err) } @@ -73,7 +73,7 @@ func TestCrossProcessDoubleRegistration(t *testing.T) { data := testCrossProcessData - servedContent, err := data.Slice(0, testCrossProcessPageSize) + servedContent, err := data.Slice(0, int64(testCrossProcessPageSize)) if err != nil { t.Fatal("cannot read content", err) } @@ -99,11 +99,11 @@ func TestHelperProcess(t *testing.T) { start := uintptr(startRaw) uffdFile := os.NewFile(uintptr(3), "userfaultfd") - uffd := NewUserfaultfdFromFd(uffdFile.Fd(), true) + uffd := NewUserfaultfdFromFd(uffdFile.Fd()) // done in the FC // Check: The reregistration works - err = uffd.register(start, uint64(testCrossProcessSize), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + err = uffd.Register(start, testCrossProcessSize, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { fmt.Print("exit registering uffd", err) os.Exit(1) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 07f2c3eda7..79a265fd59 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -12,7 +12,7 @@ import ( ) func TestUffdMissing(t *testing.T) { - pagesize := int64(header.PageSize) + pagesize := uint64(header.PageSize) data, size := prepareTestData(pagesize) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) @@ -28,7 +28,7 @@ func TestUffdMissing(t *testing.T) { memoryArea, memoryStart := newMock4KPageMmap(size) - err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING) + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) if err != nil { t.Fatal("failed to register memory", err) } @@ -48,7 +48,7 @@ func TestUffdMissing(t *testing.T) { } }() - d, err := data.Slice(0, pagesize) + d, err := data.Slice(0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } @@ -59,7 +59,7 @@ func TestUffdMissing(t *testing.T) { } func TestUffdWriteProtect(t *testing.T) { - pagesize := int64(header.PageSize) + pagesize := uint64(header.PageSize) data, size := prepareTestData(pagesize) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) @@ -75,12 +75,12 @@ func TestUffdWriteProtect(t *testing.T) { memoryArea, memoryStart := newMock4KPageMmap(size) - err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_WP) + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_WP) if err != nil { t.Fatal("failed to register memory", err) } - err = uffd.AddWriteProtection(memoryStart, uint64(size)) + err = uffd.AddWriteProtection(memoryStart, size) if err != nil { t.Fatal("failed to write protect memory", err) } @@ -106,7 +106,7 @@ func TestUffdWriteProtect(t *testing.T) { } func TestUffdWriteProtectWithMissing(t *testing.T) { - pagesize := int64(header.PageSize) + pagesize := uint64(header.PageSize) data, size := prepareTestData(pagesize) @@ -123,12 +123,12 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { memoryArea, memoryStart := newMock4KPageMmap(size) - err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { t.Fatal("failed to register memory", err) } - err = uffd.AddWriteProtection(memoryStart, uint64(size)) + err = uffd.AddWriteProtection(memoryStart, size) if err != nil { t.Fatal("failed to write protect memory", err) } @@ -148,7 +148,7 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { } }() - d, err := data.Slice(0, pagesize) + d, err := data.Slice(0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } @@ -164,7 +164,7 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { // We are trying to simulate registering the missing handler in the FC and then registering the missing+wp handler again in the orchestrator func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { - pagesize := int64(header.PageSize) + pagesize := uint64(header.PageSize) data, size := prepareTestData(pagesize) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) @@ -181,7 +181,7 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { memoryArea, memoryStart := newMock4KPageMmap(size) // done in the FC - err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING) + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) if err != nil { t.Fatal("failed to register memory", err) } @@ -190,12 +190,12 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { // done little later in the orchestrator // both flags needs to be present - err = uffd.Register(memoryStart, uint64(size), UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { t.Fatal("failed to register memory", err) } - err = uffd.AddWriteProtection(memoryStart, uint64(size)) + err = uffd.AddWriteProtection(memoryStart, size) if err != nil { t.Fatal("failed to write protect memory", err) } @@ -215,7 +215,7 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { } }() - d, err := data.Slice(0, pagesize) + d, err := data.Slice(0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } From 9f1708cf11b691fe440687b64136394e5dbd3af7 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sat, 9 Aug 2025 18:09:24 -0700 Subject: [PATCH 06/49] [WIP] Check multicopy --- .../internal/orchestrator/create_instance.go | 4 +- .../template-manager/create_template.go | 11 +- .../internal/sandbox/build/build.go | 2 +- .../sandbox/uffd/mapping/firecracker.go | 6 +- .../internal/sandbox/uffd/mapping/mapping.go | 2 +- .../internal/sandbox/uffd/uffd.go | 2 +- .../sandbox/uffd/userfaultfd/constants.go | 3 + .../sandbox/uffd/userfaultfd/serve.go | 139 ++++++++++++------ .../sandbox/uffd/userfaultfd/userfaultfd.go | 24 +++ 9 files changed, 133 insertions(+), 60 deletions(-) diff --git a/packages/api/internal/orchestrator/create_instance.go b/packages/api/internal/orchestrator/create_instance.go index d6dce5e684..a732d391c5 100644 --- a/packages/api/internal/orchestrator/create_instance.go +++ b/packages/api/internal/orchestrator/create_instance.go @@ -85,7 +85,7 @@ func (o *Orchestrator) CreateSandbox( telemetry.ReportEvent(ctx, "Reserved sandbox for team") defer releaseTeamSandboxReservation() - features, err := sandbox.NewVersionInfo(build.FirecrackerVersion) + _, err = sandbox.NewVersionInfo(build.FirecrackerVersion) if err != nil { errMsg := fmt.Errorf("failed to get features for firecracker version '%s': %w", build.FirecrackerVersion, err) @@ -128,7 +128,7 @@ func (o *Orchestrator) CreateSandbox( EnvVars: envVars, EnvdAccessToken: envdAuthToken, MaxSandboxLength: team.Tier.MaxLengthHours, - HugePages: features.HasHugePages(), + HugePages: false, RamMb: build.RamMb, Vcpu: build.Vcpu, Snapshot: isResume, diff --git a/packages/api/internal/template-manager/create_template.go b/packages/api/internal/template-manager/create_template.go index 6ac3c10d1a..698f6143d4 100644 --- a/packages/api/internal/template-manager/create_template.go +++ b/packages/api/internal/template-manager/create_template.go @@ -12,7 +12,6 @@ import ( "google.golang.org/grpc/metadata" "github.com/e2b-dev/infra/packages/api/internal/api" - "github.com/e2b-dev/infra/packages/api/internal/sandbox" "github.com/e2b-dev/infra/packages/api/internal/utils" templatemanagergrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager" "github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild" @@ -81,10 +80,10 @@ func (tm *TemplateManager) CreateTemplate( } }() - features, err := sandbox.NewVersionInfo(firecrackerVersion) - if err != nil { - return fmt.Errorf("failed to get features for firecracker version '%s': %w", firecrackerVersion, err) - } + // features, err := sandbox.NewVersionInfo(firecrackerVersion) + // if err != nil { + // return fmt.Errorf("failed to get features for firecracker version '%s': %w", firecrackerVersion, err) + // } cli, err := tm.GetClusterBuildClient(clusterID, nodeID) if err != nil { @@ -114,7 +113,7 @@ func (tm *TemplateManager) CreateTemplate( DiskSizeMB: int32(diskSizeMB), KernelVersion: kernelVersion, FirecrackerVersion: firecrackerVersion, - HugePages: features.HasHugePages(), + HugePages: false, StartCommand: startCmd, ReadyCommand: readyCmd, Force: force, diff --git a/packages/orchestrator/internal/sandbox/build/build.go b/packages/orchestrator/internal/sandbox/build/build.go index 386577f848..6c9d307eb6 100644 --- a/packages/orchestrator/internal/sandbox/build/build.go +++ b/packages/orchestrator/internal/sandbox/build/build.go @@ -104,7 +104,7 @@ func (b *File) Slice(ctx context.Context, off, _ int64) ([]byte, error) { // Pass empty huge page when the build id is nil. if *buildID == uuid.Nil { - return header.EmptyHugePage, nil + return header.EmptyHugePage[:length], nil } build, err := b.getBuild(ctx, buildID) diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go index 48ee4f2db4..7bd6d344c7 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go @@ -18,15 +18,15 @@ func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) uint64 { 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) (uint64, uint64, error) { +func (m FcMappings) GetRange(addr uintptr) (uint64, uint64, uint64, error) { for _, m := range m { if addr < m.BaseHostVirtAddr || m.BaseHostVirtAddr+m.Size <= addr { // Outside of this mapping continue } - return m.relativeOffset(addr), uint64(m.PageSize), nil + return m.relativeOffset(addr), uint64(m.PageSize), uint64(m.Size), nil } - return 0, 0, fmt.Errorf("address %d not found in any mapping", addr) + return 0, 0, 0, fmt.Errorf("address %d not found in any mapping", addr) } diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go index 33f9191f36..7536e924e3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go @@ -1,5 +1,5 @@ package mapping type Mappings interface { - GetRange(addr uintptr) (offset uint64, pagesize uint64, err error) + GetRange(addr uintptr) (offset uint64, pagesize uint64, size uint64, err error) } diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index b1c2c4f13f..6aba9896c8 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -167,7 +167,7 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { // return fmt.Errorf("failed to reregister memory region with write protection %d-%d", region.Offset, region.Offset+region.Size) // } - // // Add write protection to the regions provided by the UFFD + // Add write protection to the regions provided by the UFFD // err = uffd.AddWriteProtection( // region.Offset+region.BaseHostVirtAddr, // uint64(region.Size), diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index 070379ce7b..10930086f3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -28,14 +28,17 @@ const ( UFFDIO_REGISTER_MODE_MISSING = C.UFFDIO_REGISTER_MODE_MISSING UFFDIO_REGISTER_MODE_WP = C.UFFDIO_REGISTER_MODE_WP + UFFDIO_REGISTER_MODE_MINOR = C.UFFDIO_REGISTER_MODE_MINOR 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 diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index bd014298e4..49070c98c6 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -1,5 +1,16 @@ package userfaultfd +/* +This is the flow of the UFFD events: + +```mermaid +flowchart TD +A[missing page] -- write (WRITE flag) --> B(COPY) --> C[mark as dirty] +A -- read (0 flag) --> D(COPY + WP protect) --> E[faulted page] +E -- write (WP|WRITE flag) --> F(remove WP) --> C +``` +*/ + import ( "context" "errors" @@ -32,8 +43,7 @@ func (u *userfaultfd) Serve( var eg errgroup.Group - missingPagesBeingHandled := map[uint64]struct{}{} - writePagesBeingHandled := map[uint64]struct{}{} + missingChunksBeingHandled := map[uint64]struct{}{} outerLoop: for { @@ -125,34 +135,14 @@ outerLoop: addr := GetPagefaultAddress(&pagefault) - offset, pagesize, err := mappings.GetRange(addr) + offset, pagesize, size, err := mappings.GetRange(addr) if err != nil { logger.Error("UFFD serve get mapping error", zap.Error(err)) return fmt.Errorf("failed to map: %w", err) } - switch { - case pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0: - if _, ok := writePagesBeingHandled[offset]; ok { - continue - } - - writePagesBeingHandled[offset] = struct{}{} - - fallthrough - case pagefault.flags == 0: - if _, ok := missingPagesBeingHandled[offset]; ok { - continue - } - - missingPagesBeingHandled[offset] = struct{}{} - - fmt.Fprintf(os.Stderr, "[%d] writes without protection: %d, missing read: %d\n", offset/pagesize, len(writePagesBeingHandled), len(missingPagesBeingHandled)) - - // TODO: Copy the page ahead as we already fetch a chunk (increase it) - // TODO: Prefetch the top of the memory downward, the start of the memory up - + if pagefault.flags&UFFD_PAGEFAULT_FLAG_WP != 0 { eg.Go(func() error { defer func() { if r := recover(); r != nil { @@ -160,40 +150,97 @@ outerLoop: } }() - b, sliceErr := src.Slice(int64(offset), int64(pagesize)) - if sliceErr != nil { + wpErr := u.removeWriteProtection(addr, pagesize) + if wpErr != nil { + return fmt.Errorf("error removing write protection from page %d", addr) + } - signalErr := fdExit.SignalExit() + return nil + }) - joinedErr := errors.Join(sliceErr, signalErr) + fmt.Fprintf(os.Stderr, "wp trigger %d %d\n", addr, offset/pagesize) - logger.Error("UFFD serve slice error", zap.Error(joinedErr)) + continue + } - return fmt.Errorf("failed to read from source: %w", joinedErr) - } + if pagefault.flags == 0 { + fmt.Fprintf(os.Stderr, "read trigger %d %d\n", addr, offset/pagesize) + } - copyErr := u.copy(addr, b, pagesize, 0) - if copyErr == unix.EEXIST { - logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) + if pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0 { + fmt.Fprintf(os.Stderr, "write trigger %d %d\n", addr, offset/pagesize) + } - // Page is already mapped - return nil - } + if _, ok := missingChunksBeingHandled[offset]; ok { + fmt.Fprintf(os.Stderr, "page already being handled %d %d\n", offset, addr) - if copyErr != nil { - signalErr := fdExit.SignalExit() + continue + } - joinedErr := errors.Join(copyErr, signalErr) + missingChunksBeingHandled[offset] = struct{}{} - logger.Error("UFFD serve uffdio copy error", zap.Error(joinedErr)) + newPagesize := pagesize + if offset+pagesize < size { + missingChunksBeingHandled[offset+pagesize] = struct{}{} - return fmt.Errorf("failed uffdio copy %w", joinedErr) + newPagesize = pagesize * 2 + } + + var copyMode CULong + + goAddress := addr + goOffset := offset + goPagesize := newPagesize + + eg.Go(func() error { + defer func() { + if r := recover(); r != nil { + logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) } + }() + + b, sliceErr := src.Slice(int64(goOffset), int64(goPagesize)) + if sliceErr != nil { + signalErr := fdExit.SignalExit() + + joinedErr := errors.Join(sliceErr, signalErr) + + logger.Error("UFFD serve slice error", zap.Error(joinedErr)) + + return fmt.Errorf("failed to read from source: %w", joinedErr) + } + + fmt.Fprintf( + os.Stderr, + "COPY: triggered by address %d, from offset %d, pagesize %d, index %d, total size %d, total offset %d\n", + goAddress+uintptr(goOffset), + goOffset, + goPagesize, + goOffset/goPagesize, + size, + goOffset, + ) + + copyErr := u.copy(goAddress, b[:goPagesize], goPagesize, copyMode) + if copyErr == unix.EEXIST { + logger.Debug("UFFD serve page already mapped", zap.Any("offset", goOffset), zap.Any("pagesize", goPagesize)) + + // Page is already mapped return nil - }) - default: - return fmt.Errorf("invalid pagefault flags %d for address %d", pagefault.flags, addr) - } + } + + if copyErr != nil { + signalErr := fdExit.SignalExit() + + joinedErr := errors.Join(copyErr, signalErr) + + logger.Error("UFFD serve uffdio copy error", zap.Error(joinedErr)) + + return fmt.Errorf("failed uffdio copy %w", joinedErr) + } + + return nil + }) } } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 0efeb55745..1f2f256c7a 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -84,6 +84,30 @@ func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULo return errno } + // Check if the copied size matches the requested pagesize + if uint64(cpy.copy) != pagesize { + return fmt.Errorf("UFFDIO_COPY copied %d bytes, expected %d", cpy.copy, pagesize) + } + + return nil +} + +// Wake wakes up the faulting thread(s) waiting on the given address range. +func (u *userfaultfd) Wake(addr uintptr, size uint64) error { + type uffdioWake struct { + addr CULong + len CULong + flags CULong + } + wake := uffdioWake{ + addr: CULong(addr), + len: CULong(size), + flags: 0, + } + + if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_WAKE, uintptr(unsafe.Pointer(&wake))); errno != 0 { + return errno + } return nil } From b07e13d5435db63bcaa50f67ba36cc49d1966734 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Tue, 12 Aug 2025 23:54:39 +0200 Subject: [PATCH 07/49] [WIP] Multicopy --- .../internal/sandbox/uffd/uffd.go | 2 ++ .../sandbox/uffd/userfaultfd/serve.go | 30 +++---------------- 2 files changed, 6 insertions(+), 26 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index 6aba9896c8..f99b1c994f 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -178,6 +178,8 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { // } // } + fmt.Fprintf(os.Stderr, "uffd: serving: %d\n", len(m)) + u.readyCh <- struct{}{} err = Serve( diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 49070c98c6..69e98d9a46 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -135,7 +135,7 @@ outerLoop: addr := GetPagefaultAddress(&pagefault) - offset, pagesize, size, err := mappings.GetRange(addr) + offset, pagesize, _, err := mappings.GetRange(addr) if err != nil { logger.Error("UFFD serve get mapping error", zap.Error(err)) @@ -179,19 +179,8 @@ outerLoop: missingChunksBeingHandled[offset] = struct{}{} - newPagesize := pagesize - if offset+pagesize < size { - missingChunksBeingHandled[offset+pagesize] = struct{}{} - - newPagesize = pagesize * 2 - } - var copyMode CULong - goAddress := addr - goOffset := offset - goPagesize := newPagesize - eg.Go(func() error { defer func() { if r := recover(); r != nil { @@ -199,7 +188,7 @@ outerLoop: } }() - b, sliceErr := src.Slice(int64(goOffset), int64(goPagesize)) + b, sliceErr := src.Slice(int64(offset), int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() @@ -210,20 +199,9 @@ outerLoop: return fmt.Errorf("failed to read from source: %w", joinedErr) } - fmt.Fprintf( - os.Stderr, - "COPY: triggered by address %d, from offset %d, pagesize %d, index %d, total size %d, total offset %d\n", - goAddress+uintptr(goOffset), - goOffset, - goPagesize, - goOffset/goPagesize, - size, - goOffset, - ) - - copyErr := u.copy(goAddress, b[:goPagesize], goPagesize, copyMode) + copyErr := u.copy(addr, b, pagesize, copyMode) if copyErr == unix.EEXIST { - logger.Debug("UFFD serve page already mapped", zap.Any("offset", goOffset), zap.Any("pagesize", goPagesize)) + logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) // Page is already mapped From ff23f18ac75806cbcbde8b683f65bcce2927a5d4 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:25:09 +0200 Subject: [PATCH 08/49] Remove experiments --- .../internal/orchestrator/create_instance.go | 4 +-- .../template-manager/create_template.go | 11 ++++---- .../internal/sandbox/build/build.go | 2 +- .../internal/sandbox/uffd/uffd.go | 4 --- .../sandbox/uffd/userfaultfd/constants.go | 2 +- .../sandbox/uffd/userfaultfd/serve.go | 27 ++++++++++--------- .../sandbox/uffd/userfaultfd/userfaultfd.go | 2 +- 7 files changed, 25 insertions(+), 27 deletions(-) diff --git a/packages/api/internal/orchestrator/create_instance.go b/packages/api/internal/orchestrator/create_instance.go index a732d391c5..d6dce5e684 100644 --- a/packages/api/internal/orchestrator/create_instance.go +++ b/packages/api/internal/orchestrator/create_instance.go @@ -85,7 +85,7 @@ func (o *Orchestrator) CreateSandbox( telemetry.ReportEvent(ctx, "Reserved sandbox for team") defer releaseTeamSandboxReservation() - _, err = sandbox.NewVersionInfo(build.FirecrackerVersion) + features, err := sandbox.NewVersionInfo(build.FirecrackerVersion) if err != nil { errMsg := fmt.Errorf("failed to get features for firecracker version '%s': %w", build.FirecrackerVersion, err) @@ -128,7 +128,7 @@ func (o *Orchestrator) CreateSandbox( EnvVars: envVars, EnvdAccessToken: envdAuthToken, MaxSandboxLength: team.Tier.MaxLengthHours, - HugePages: false, + HugePages: features.HasHugePages(), RamMb: build.RamMb, Vcpu: build.Vcpu, Snapshot: isResume, diff --git a/packages/api/internal/template-manager/create_template.go b/packages/api/internal/template-manager/create_template.go index 698f6143d4..6ac3c10d1a 100644 --- a/packages/api/internal/template-manager/create_template.go +++ b/packages/api/internal/template-manager/create_template.go @@ -12,6 +12,7 @@ import ( "google.golang.org/grpc/metadata" "github.com/e2b-dev/infra/packages/api/internal/api" + "github.com/e2b-dev/infra/packages/api/internal/sandbox" "github.com/e2b-dev/infra/packages/api/internal/utils" templatemanagergrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager" "github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild" @@ -80,10 +81,10 @@ func (tm *TemplateManager) CreateTemplate( } }() - // features, err := sandbox.NewVersionInfo(firecrackerVersion) - // if err != nil { - // return fmt.Errorf("failed to get features for firecracker version '%s': %w", firecrackerVersion, err) - // } + features, err := sandbox.NewVersionInfo(firecrackerVersion) + if err != nil { + return fmt.Errorf("failed to get features for firecracker version '%s': %w", firecrackerVersion, err) + } cli, err := tm.GetClusterBuildClient(clusterID, nodeID) if err != nil { @@ -113,7 +114,7 @@ func (tm *TemplateManager) CreateTemplate( DiskSizeMB: int32(diskSizeMB), KernelVersion: kernelVersion, FirecrackerVersion: firecrackerVersion, - HugePages: false, + HugePages: features.HasHugePages(), StartCommand: startCmd, ReadyCommand: readyCmd, Force: force, diff --git a/packages/orchestrator/internal/sandbox/build/build.go b/packages/orchestrator/internal/sandbox/build/build.go index 6c9d307eb6..386577f848 100644 --- a/packages/orchestrator/internal/sandbox/build/build.go +++ b/packages/orchestrator/internal/sandbox/build/build.go @@ -104,7 +104,7 @@ func (b *File) Slice(ctx context.Context, off, _ int64) ([]byte, error) { // Pass empty huge page when the build id is nil. if *buildID == uuid.Nil { - return header.EmptyHugePage[:length], nil + return header.EmptyHugePage, nil } build, err := b.getBuild(ctx, buildID) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index f99b1c994f..eb309bbe66 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -210,10 +210,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() } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index 10930086f3..540d848071 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -82,7 +82,7 @@ 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, diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 69e98d9a46..a9783daf8d 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -163,31 +163,32 @@ outerLoop: continue } - if pagefault.flags == 0 { - fmt.Fprintf(os.Stderr, "read trigger %d %d\n", addr, offset/pagesize) - } - - if pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0 { - fmt.Fprintf(os.Stderr, "write trigger %d %d\n", addr, offset/pagesize) - } - + // This prevents serving missing pages multiple times. + // For normal sized pages with swap on, the behavior seems not to be properly described in docs + // and it's not clear if the missing can be legitimately triggered multiple times. if _, ok := missingChunksBeingHandled[offset]; ok { - fmt.Fprintf(os.Stderr, "page already being handled %d %d\n", offset, addr) - continue } missingChunksBeingHandled[offset] = struct{}{} - var copyMode CULong + if pagefault.flags == 0 { + // fmt.Fprintf(os.Stderr, "read trigger %d %d\n", addr, offset/pagesize) + } + + if pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0 { + // fmt.Fprintf(os.Stderr, "write trigger %d %d\n", addr, offset/pagesize) + } eg.Go(func() error { defer func() { if r := recover(); r != nil { - logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) + logger.Error("UFFD serve panic", zap.Any("pagesize", pagesize), zap.Any("panic", r)) } }() + var copyMode CULong + b, sliceErr := src.Slice(int64(offset), int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() @@ -201,7 +202,7 @@ outerLoop: copyErr := u.copy(addr, b, pagesize, copyMode) if copyErr == unix.EEXIST { - logger.Debug("UFFD serve page already mapped", zap.Any("offset", offset), zap.Any("pagesize", pagesize)) + logger.Debug("UFFD serve page already mapped", zap.Any("offset", addr), zap.Any("pagesize", pagesize)) // Page is already mapped diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 1f2f256c7a..ecb2e737f7 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -78,7 +78,7 @@ func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULong) error { - cpy := NewUffdioCopy(data, CULong(addr)&^CULong(pagesize-1), CULong(pagesize), mode, 0) + cpy := NewUffdioCopy(data, CULong(addr), CULong(pagesize), mode, 0) if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { return errno From 389862e71ce12244710892d78bd0354c4b8bcb5f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:27:45 +0200 Subject: [PATCH 09/49] Add the page alignment check (only in method) --- .../internal/sandbox/uffd/userfaultfd/userfaultfd.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index ecb2e737f7..1f2f256c7a 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -78,7 +78,7 @@ func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULong) error { - cpy := NewUffdioCopy(data, CULong(addr), CULong(pagesize), mode, 0) + cpy := NewUffdioCopy(data, CULong(addr)&^CULong(pagesize-1), CULong(pagesize), mode, 0) if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_COPY, uintptr(unsafe.Pointer(&cpy))); errno != 0 { return errno From a8ec1bffeb8f2f1858f00247436686c38c964a64 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:28:37 +0200 Subject: [PATCH 10/49] Remove commented out part --- .../internal/sandbox/uffd/uffd.go | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index eb309bbe66..bd30bd3156 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -155,29 +155,6 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { } }() - // if u.writeProtection { - // for _, region := range m { - // // Register the WP. It is possible that the memory region was already registered (with missing pages in FC), but registering it again with bigger subset should merge these. - // err := uffd.Register( - // region.Offset+region.BaseHostVirtAddr, - // uint64(region.Size), - // userfaultfd.UFFDIO_REGISTER_MODE_WP|userfaultfd.UFFDIO_REGISTER_MODE_MISSING, - // ) - // if err != nil { - // return fmt.Errorf("failed to reregister memory region with write protection %d-%d", region.Offset, region.Offset+region.Size) - // } - - // Add write protection to the regions provided by the UFFD - // err = uffd.AddWriteProtection( - // region.Offset+region.BaseHostVirtAddr, - // uint64(region.Size), - // ) - // if err != nil { - // return fmt.Errorf("failed to add write protection to region %d-%d", region.Offset, region.Offset+region.Size) - // } - // } - // } - fmt.Fprintf(os.Stderr, "uffd: serving: %d\n", len(m)) u.readyCh <- struct{}{} From 5632f63b5273e8962028faf63a8a52c01c26fa1b Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:31:10 +0200 Subject: [PATCH 11/49] Remove unused return field --- .../internal/sandbox/uffd/mapping/firecracker.go | 6 +++--- .../orchestrator/internal/sandbox/uffd/mapping/mapping.go | 2 +- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go index 7bd6d344c7..48ee4f2db4 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go @@ -18,15 +18,15 @@ func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) uint64 { 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) (uint64, uint64, uint64, error) { +func (m FcMappings) GetRange(addr uintptr) (uint64, uint64, error) { for _, m := range m { if addr < m.BaseHostVirtAddr || m.BaseHostVirtAddr+m.Size <= addr { // Outside of this mapping continue } - return m.relativeOffset(addr), uint64(m.PageSize), uint64(m.Size), nil + return m.relativeOffset(addr), uint64(m.PageSize), nil } - return 0, 0, 0, fmt.Errorf("address %d not found in any mapping", addr) + return 0, 0, fmt.Errorf("address %d not found in any mapping", addr) } diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go index 7536e924e3..33f9191f36 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go @@ -1,5 +1,5 @@ package mapping type Mappings interface { - GetRange(addr uintptr) (offset uint64, pagesize uint64, size uint64, err error) + GetRange(addr uintptr) (offset uint64, pagesize uint64, err error) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index a9783daf8d..976be96f0b 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -135,7 +135,7 @@ outerLoop: addr := GetPagefaultAddress(&pagefault) - offset, pagesize, _, err := mappings.GetRange(addr) + offset, pagesize, err := mappings.GetRange(addr) if err != nil { logger.Error("UFFD serve get mapping error", zap.Error(err)) From ef33099c224764d0611a549207ea1cc7fae79a9f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:39:24 +0200 Subject: [PATCH 12/49] Cleanup --- .../sandbox/uffd/mapping/firecracker.go | 6 +++--- .../internal/sandbox/uffd/mapping/mapping.go | 2 +- .../sandbox/uffd/userfaultfd/constants.go | 1 - .../sandbox/uffd/userfaultfd/mocks_test.go | 4 ++-- .../sandbox/uffd/userfaultfd/serve.go | 6 ++---- .../sandbox/uffd/userfaultfd/userfaultfd.go | 19 ------------------- 6 files changed, 8 insertions(+), 30 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go index 48ee4f2db4..4cdad3fc42 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go @@ -11,14 +11,14 @@ type GuestRegionUffdMapping struct { PageSize uintptr `json:"page_size_kib"` } -func (m *GuestRegionUffdMapping) relativeOffset(addr uintptr) uint64 { - return uint64(m.Offset + addr - m.BaseHostVirtAddr) +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) (uint64, uint64, error) { +func (m FcMappings) GetRange(addr uintptr) (int64, uint64, error) { for _, m := range m { if addr < m.BaseHostVirtAddr || m.BaseHostVirtAddr+m.Size <= addr { // Outside of this mapping diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go index 33f9191f36..592925a695 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go +++ b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go @@ -1,5 +1,5 @@ package mapping type Mappings interface { - GetRange(addr uintptr) (offset uint64, pagesize uint64, err error) + GetRange(addr uintptr) (offset int64, pagesize uint64, err error) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index 540d848071..cec6935227 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -28,7 +28,6 @@ const ( UFFDIO_REGISTER_MODE_MISSING = C.UFFDIO_REGISTER_MODE_MISSING UFFDIO_REGISTER_MODE_WP = C.UFFDIO_REGISTER_MODE_WP - UFFDIO_REGISTER_MODE_MINOR = C.UFFDIO_REGISTER_MODE_MINOR UFFDIO_WRITEPROTECT_MODE_WP = C.UFFDIO_WRITEPROTECT_MODE_WP UFFDIO_COPY_MODE_WP = C.UFFDIO_COPY_MODE_WP diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go index 5763a1c4cc..9cce370eb2 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go @@ -26,11 +26,11 @@ func newMockMappings(start uintptr, size, pagesize uint64) *mockMappings { } } -func (m *mockMappings) GetRange(addr uintptr) (uint64, uint64, error) { +func (m *mockMappings) GetRange(addr uintptr) (int64, uint64, error) { offset := addr - m.start pagesize := m.pagesize - return uint64(offset), pagesize, nil + return int64(offset), pagesize, nil } type mockSlicer struct { diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 976be96f0b..41c44101da 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -43,7 +43,7 @@ func (u *userfaultfd) Serve( var eg errgroup.Group - missingChunksBeingHandled := map[uint64]struct{}{} + missingChunksBeingHandled := map[int64]struct{}{} outerLoop: for { @@ -158,8 +158,6 @@ outerLoop: return nil }) - fmt.Fprintf(os.Stderr, "wp trigger %d %d\n", addr, offset/pagesize) - continue } @@ -189,7 +187,7 @@ outerLoop: var copyMode CULong - b, sliceErr := src.Slice(int64(offset), int64(pagesize)) + b, sliceErr := src.Slice(offset, int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 1f2f256c7a..de6dcb4294 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -92,25 +92,6 @@ func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULo return nil } -// Wake wakes up the faulting thread(s) waiting on the given address range. -func (u *userfaultfd) Wake(addr uintptr, size uint64) error { - type uffdioWake struct { - addr CULong - len CULong - flags CULong - } - wake := uffdioWake{ - addr: CULong(addr), - len: CULong(size), - flags: 0, - } - - if _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_WAKE, uintptr(unsafe.Pointer(&wake))); errno != 0 { - return errno - } - return nil -} - func (u *userfaultfd) Close() error { return syscall.Close(int(u.fd)) } From 5384bec396a5c025cb439c7f4e90dfe6950b362e Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:40:16 +0200 Subject: [PATCH 13/49] Remove WP parts --- .../sandbox/uffd/userfaultfd/serve.go | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 41c44101da..b9b81ccf78 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -142,25 +142,6 @@ outerLoop: return fmt.Errorf("failed to map: %w", err) } - if pagefault.flags&UFFD_PAGEFAULT_FLAG_WP != 0 { - eg.Go(func() error { - defer func() { - if r := recover(); r != nil { - logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) - } - }() - - wpErr := u.removeWriteProtection(addr, pagesize) - if wpErr != nil { - return fmt.Errorf("error removing write protection from page %d", addr) - } - - return nil - }) - - continue - } - // This prevents serving missing pages multiple times. // For normal sized pages with swap on, the behavior seems not to be properly described in docs // and it's not clear if the missing can be legitimately triggered multiple times. @@ -170,14 +151,6 @@ outerLoop: missingChunksBeingHandled[offset] = struct{}{} - if pagefault.flags == 0 { - // fmt.Fprintf(os.Stderr, "read trigger %d %d\n", addr, offset/pagesize) - } - - if pagefault.flags&UFFD_PAGEFAULT_FLAG_WRITE != 0 { - // fmt.Fprintf(os.Stderr, "write trigger %d %d\n", addr, offset/pagesize) - } - eg.Go(func() error { defer func() { if r := recover(); r != nil { From f26cc5cdc24e79c609eac8d86b72d9be03f9b793 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 20 Aug 2025 00:49:32 +0200 Subject: [PATCH 14/49] Cleanup --- packages/orchestrator/internal/sandbox/uffd/uffd.go | 3 --- .../internal/sandbox/uffd/userfaultfd/serve.go | 8 +++----- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index bd30bd3156..ca569def95 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -155,8 +155,6 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { } }() - fmt.Fprintf(os.Stderr, "uffd: serving: %d\n", len(m)) - u.readyCh <- struct{}{} err = Serve( @@ -168,7 +166,6 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { zap.L().With(logger.WithSandboxID(sandboxId)), ) if err != nil { - fmt.Fprintf(os.Stderr, "serve error %v", err) return fmt.Errorf("failed handling uffd: %w", err) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index b9b81ccf78..1607367886 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -15,7 +15,6 @@ import ( "context" "errors" "fmt" - "os" "syscall" "unsafe" @@ -35,7 +34,6 @@ func (u *userfaultfd) Serve( fdExit *fdexit.FdExit, logger *zap.Logger, ) error { - defer fmt.Fprintf(os.Stderr, "exiting serve >>>>>>>>>>>>") pollFds := []unix.PollFd{ {Fd: int32(u.fd), Events: unix.POLLIN}, {Fd: int32(fdExit.Reader()), Events: unix.POLLIN}, @@ -43,7 +41,7 @@ func (u *userfaultfd) Serve( var eg errgroup.Group - missingChunksBeingHandled := map[int64]struct{}{} + handledPages := map[int64]struct{}{} outerLoop: for { @@ -145,11 +143,11 @@ outerLoop: // This prevents serving missing pages multiple times. // For normal sized pages with swap on, the behavior seems not to be properly described in docs // and it's not clear if the missing can be legitimately triggered multiple times. - if _, ok := missingChunksBeingHandled[offset]; ok { + if _, ok := handledPages[offset]; ok { continue } - missingChunksBeingHandled[offset] = struct{}{} + handledPages[offset] = struct{}{} eg.Go(func() error { defer func() { From 47e14b058ef45ade58ee0441fa3c17d0bfe8ed00 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 16:53:45 -0700 Subject: [PATCH 15/49] Fix compile errors --- packages/orchestrator/internal/sandbox/uffd/uffd.go | 4 +--- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index ca569def95..60f5ea2d38 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -59,7 +59,6 @@ func New(memfile block.ReadonlyDevice, socketPath string, blockSize int64) (*Uff return &Uffd{ exit: utils.NewErrorOnce(), - exitCh: make(chan error, 1), readyCh: make(chan struct{}, 1), fdExit: fdExit, memfile: trackedMemfile, @@ -157,9 +156,8 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error { u.readyCh <- struct{}{} - err = Serve( + err = uffd.Serve( ctx, - uffd, m, u.memfile, u.fdExit, diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 1607367886..a78a034c15 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -158,7 +158,7 @@ outerLoop: var copyMode CULong - b, sliceErr := src.Slice(offset, int64(pagesize)) + b, sliceErr := src.Slice(ctx, offset, int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() From 7ba06e79eeab9ee5fdd0a58d7a4f9cd4f3e61dee Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 17:13:12 -0700 Subject: [PATCH 16/49] Cleanup --- .../internal/sandbox/uffd/mapping/mapping.go | 5 - .../uffd/{mapping => memory}/firecracker.go | 35 +++--- .../internal/sandbox/uffd/memory/map.go | 5 + .../sandbox/uffd/testutils/contiguous_map.go | 24 +++++ .../internal/sandbox/uffd/testutils/mmap.go | 48 +++++++++ .../sandbox/uffd/testutils/test_data.go | 31 ++++++ .../internal/sandbox/uffd/uffd.go | 4 +- .../sandbox/uffd/userfaultfd/constants.go | 3 +- .../sandbox/uffd/userfaultfd/mocks_test.go | 100 ------------------ .../sandbox/uffd/userfaultfd/serve.go | 6 +- .../userfaultfd_cross_process_test.go | 20 ++-- .../uffd/userfaultfd/userfaultfd_test.go | 42 ++++---- 12 files changed, 167 insertions(+), 156 deletions(-) delete mode 100644 packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go rename packages/orchestrator/internal/sandbox/uffd/{mapping => memory}/firecracker.go (61%) create mode 100644 packages/orchestrator/internal/sandbox/uffd/memory/map.go create mode 100644 packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go create mode 100644 packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go create mode 100644 packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go delete mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go b/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go deleted file mode 100644 index 592925a695..0000000000 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/mapping.go +++ /dev/null @@ -1,5 +0,0 @@ -package mapping - -type Mappings interface { - GetRange(addr uintptr) (offset int64, pagesize uint64, err error) -} diff --git a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go similarity index 61% rename from packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go rename to packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go index 4cdad3fc42..67dbf0e8e8 100644 --- a/packages/orchestrator/internal/sandbox/uffd/mapping/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go @@ -1,7 +1,24 @@ -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 || m.BaseHostVirtAddr+m.Size <= hostVirtAddr { + // Outside of this mapping + + continue + } + + 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 +31,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, uint64, error) { - for _, m := range m { - if addr < m.BaseHostVirtAddr || m.BaseHostVirtAddr+m.Size <= addr { - // Outside of this mapping - continue - } - - return m.relativeOffset(addr), uint64(m.PageSize), nil - } - - return 0, 0, fmt.Errorf("address %d not found in any mapping", addr) -} diff --git a/packages/orchestrator/internal/sandbox/uffd/memory/map.go b/packages/orchestrator/internal/sandbox/uffd/memory/map.go new file mode 100644 index 0000000000..4782b571ae --- /dev/null +++ b/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) +} diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go b/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go new file mode 100644 index 0000000000..54377a2090 --- /dev/null +++ b/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 +} diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go b/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go new file mode 100644 index 0000000000..fc87df6bc9 --- /dev/null +++ b/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])) +} diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go b/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go new file mode 100644 index 0000000000..57991751b1 --- /dev/null +++ b/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 +} diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index 60f5ea2d38..0449161da0 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -16,7 +16,7 @@ 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" @@ -120,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 { diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index cec6935227..274fb7b77a 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -42,8 +42,7 @@ const ( 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 ( diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go deleted file mode 100644 index 9cce370eb2..0000000000 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/mocks_test.go +++ /dev/null @@ -1,100 +0,0 @@ -package userfaultfd - -import ( - "math" - "syscall" - "unsafe" - - "golang.org/x/sys/unix" - - "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" -) - -const pagesInTestData = 32 - -type mockMappings struct { - start uintptr - size uint64 - pagesize uint64 -} - -func newMockMappings(start uintptr, size, pagesize uint64) *mockMappings { - return &mockMappings{ - start: start, - size: size, - pagesize: pagesize, - } -} - -func (m *mockMappings) GetRange(addr uintptr) (int64, uint64, error) { - offset := addr - m.start - pagesize := m.pagesize - - return int64(offset), pagesize, nil -} - -type mockSlicer struct { - content []byte -} - -func newMockSlicer(content []byte) *mockSlicer { - return &mockSlicer{content: content} -} - -func (s *mockSlicer) Slice(offset, size int64) ([]byte, error) { - return s.content[offset : offset+size], nil -} - -func newMock4KPageMmap(size uint64) ([]byte, uintptr) { - return newMockMmap(size, header.PageSize, 0) -} - -func newMock2MPageMmap(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])) -} - -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 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 -} diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index a78a034c15..6635c46239 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -24,12 +24,12 @@ 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" ) func (u *userfaultfd) Serve( ctx context.Context, - mappings mapping.Mappings, + m memory.MemoryMap, src block.Slicer, fdExit *fdexit.FdExit, logger *zap.Logger, @@ -133,7 +133,7 @@ outerLoop: addr := GetPagefaultAddress(&pagefault) - offset, pagesize, err := mappings.GetRange(addr) + offset, pagesize, err := m.GetOffset(addr) if err != nil { logger.Error("UFFD serve get mapping error", zap.Error(err)) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 94eabc56f8..ac3774ca10 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -15,20 +15,26 @@ import ( "testing" "time" + "go.uber.org/zap" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/testutils" "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" - "go.uber.org/zap" +) + +const ( + pagesInTestData = 32 + testCrossProcessPageSize = uint64(header.HugepageSize) ) var ( // Data shared across the testing processes to check if the served and received data is the same - testCrossProcessPageSize = uint64(header.HugepageSize) - testCrossProcessData, testCrossProcessSize = prepareTestData(testCrossProcessPageSize) + testCrossProcessData, testCrossProcessSize = testutils.PrepareTestData(testCrossProcessPageSize, pagesInTestData) ) // Main process, FC in our case func TestCrossProcessDoubleRegistration(t *testing.T) { - memoryArea, memoryStart := newMock2MPageMmap(testCrossProcessSize) + memoryArea, memoryStart := testutils.New2MPageMmap(testCrossProcessSize) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { @@ -73,7 +79,7 @@ func TestCrossProcessDoubleRegistration(t *testing.T) { data := testCrossProcessData - servedContent, err := data.Slice(0, int64(testCrossProcessPageSize)) + servedContent, err := data.Slice(t.Context(), 0, int64(testCrossProcessPageSize)) if err != nil { t.Fatal("cannot read content", err) } @@ -114,7 +120,7 @@ func TestHelperProcess(t *testing.T) { ppid := os.Getppid() syscall.Kill(ppid, syscall.SIGUSR1) - mappings := newMockMappings(start, testCrossProcessSize, testCrossProcessPageSize) + m := testutils.NewContiguousMap(start, testCrossProcessSize, testCrossProcessPageSize) fdExit, err := fdexit.New() if err != nil { @@ -126,7 +132,7 @@ func TestHelperProcess(t *testing.T) { data := testCrossProcessData go func() { - err := uffd.Serve(mappings, data, fdExit, zap.L()) + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 79a265fd59..86541628a3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -6,14 +6,16 @@ import ( "syscall" "testing" + "go.uber.org/zap" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/testutils" "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" - "go.uber.org/zap" ) func TestUffdMissing(t *testing.T) { pagesize := uint64(header.PageSize) - data, size := prepareTestData(pagesize) + data, size := testutils.PrepareTestData(pagesize, pagesInTestData) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { @@ -26,14 +28,14 @@ func TestUffdMissing(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - memoryArea, memoryStart := newMock4KPageMmap(size) + memoryArea, memoryStart := testutils.New4KPageMmap(size) err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) if err != nil { t.Fatal("failed to register memory", err) } - mappings := newMockMappings(memoryStart, size, pagesize) + m := testutils.NewContiguousMap(memoryStart, size, pagesize) fdExit, err := fdexit.New() if err != nil { @@ -42,13 +44,13 @@ func TestUffdMissing(t *testing.T) { defer fdExit.Close() go func() { - err := uffd.Serve(mappings, data, fdExit, zap.L()) + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { fmt.Println("[TestUffdMissing] failed to serve uffd", err) } }() - d, err := data.Slice(0, int64(pagesize)) + d, err := data.Slice(t.Context(), 0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } @@ -60,7 +62,7 @@ func TestUffdMissing(t *testing.T) { func TestUffdWriteProtect(t *testing.T) { pagesize := uint64(header.PageSize) - data, size := prepareTestData(pagesize) + data, size := testutils.PrepareTestData(pagesize, pagesInTestData) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { @@ -73,7 +75,7 @@ func TestUffdWriteProtect(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - memoryArea, memoryStart := newMock4KPageMmap(size) + memoryArea, memoryStart := testutils.New4KPageMmap(size) err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_WP) if err != nil { @@ -91,10 +93,10 @@ func TestUffdWriteProtect(t *testing.T) { } defer fdExit.Close() - mappings := newMockMappings(memoryStart, size, pagesize) + m := testutils.NewContiguousMap(memoryStart, size, pagesize) go func() { - err := uffd.Serve(mappings, data, fdExit, zap.L()) + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) } @@ -108,7 +110,7 @@ func TestUffdWriteProtect(t *testing.T) { func TestUffdWriteProtectWithMissing(t *testing.T) { pagesize := uint64(header.PageSize) - data, size := prepareTestData(pagesize) + data, size := testutils.PrepareTestData(pagesize, pagesInTestData) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { @@ -121,7 +123,7 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - memoryArea, memoryStart := newMock4KPageMmap(size) + memoryArea, memoryStart := testutils.New4KPageMmap(size) err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { @@ -133,7 +135,7 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { t.Fatal("failed to write protect memory", err) } - mappings := newMockMappings(memoryStart, size, pagesize) + m := testutils.NewContiguousMap(memoryStart, size, pagesize) fdExit, err := fdexit.New() if err != nil { @@ -142,13 +144,13 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { defer fdExit.Close() go func() { - err := uffd.Serve(mappings, data, fdExit, zap.L()) + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) } }() - d, err := data.Slice(0, int64(pagesize)) + d, err := data.Slice(t.Context(), 0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } @@ -165,7 +167,7 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { // We are trying to simulate registering the missing handler in the FC and then registering the missing+wp handler again in the orchestrator func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { pagesize := uint64(header.PageSize) - data, size := prepareTestData(pagesize) + data, size := testutils.PrepareTestData(pagesize, pagesInTestData) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { @@ -178,7 +180,7 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - memoryArea, memoryStart := newMock4KPageMmap(size) + memoryArea, memoryStart := testutils.New4KPageMmap(size) // done in the FC err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) @@ -200,7 +202,7 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { t.Fatal("failed to write protect memory", err) } - mappings := newMockMappings(memoryStart, size, pagesize) + m := testutils.NewContiguousMap(memoryStart, size, pagesize) fdExit, err := fdexit.New() if err != nil { @@ -209,13 +211,13 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { defer fdExit.Close() go func() { - err := uffd.Serve(mappings, data, fdExit, zap.L()) + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) } }() - d, err := data.Slice(0, int64(pagesize)) + d, err := data.Slice(t.Context(), 0, int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } From e91945f22362a4a985d8c1134896d7d93395ffab Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 17:18:14 -0700 Subject: [PATCH 17/49] Simplify offset check --- .../internal/sandbox/uffd/memory/firecracker.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go b/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go index 67dbf0e8e8..b3c88da0a2 100644 --- a/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go +++ b/packages/orchestrator/internal/sandbox/uffd/memory/firecracker.go @@ -7,13 +7,9 @@ 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 || m.BaseHostVirtAddr+m.Size <= hostVirtAddr { - // Outside of this mapping - - continue + if hostVirtAddr >= m.BaseHostVirtAddr && hostVirtAddr < m.BaseHostVirtAddr+m.Size { + return m.relativeOffset(hostVirtAddr), uint64(m.PageSize), nil } - - return m.relativeOffset(hostVirtAddr), uint64(m.PageSize), nil } return 0, 0, fmt.Errorf("address %d not found in any mapping", hostVirtAddr) From d4c2ea01ab3fda3834ad181ff653fafd8269e14d Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 17:20:05 -0700 Subject: [PATCH 18/49] Clarify comment --- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 6635c46239..3f3dc86c95 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -1,7 +1,7 @@ package userfaultfd /* -This is the flow of the UFFD events: +This is the flow of the UFFD events when using write protection: ```mermaid flowchart TD From 091607e8c75af9ee86605806e83f14f3082471ea Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 18:19:04 -0700 Subject: [PATCH 19/49] Add test constant --- .../sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index ac3774ca10..985aa6d171 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -25,6 +25,7 @@ import ( const ( pagesInTestData = 32 testCrossProcessPageSize = uint64(header.HugepageSize) + inheritedFdIndex = 3 ) var ( @@ -104,7 +105,7 @@ func TestHelperProcess(t *testing.T) { start := uintptr(startRaw) - uffdFile := os.NewFile(uintptr(3), "userfaultfd") + uffdFile := os.NewFile(uintptr(inheritedFdIndex), "userfaultfd") uffd := NewUserfaultfdFromFd(uffdFile.Fd()) // done in the FC From 7ae5f702df8abd8c1ab3afb8e35230c6aa81a786 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 18:20:04 -0700 Subject: [PATCH 20/49] Fix incorrect log --- .../sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 985aa6d171..0501f6e9ad 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -135,7 +135,7 @@ func TestHelperProcess(t *testing.T) { go func() { err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) if err != nil { - fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) + fmt.Println("[TestHelperProcess] failed to serve uffd", err) } }() From f9f95c501da7d609ebba933c29d8adaa69e57cea Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 18:38:04 -0700 Subject: [PATCH 21/49] Fix format --- .../uffd/userfaultfd/userfaultfd_cross_process_test.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 0501f6e9ad..83d02dc824 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -28,10 +28,8 @@ const ( inheritedFdIndex = 3 ) -var ( - // Data shared across the testing processes to check if the served and received data is the same - testCrossProcessData, testCrossProcessSize = testutils.PrepareTestData(testCrossProcessPageSize, pagesInTestData) -) +// Data shared across the testing processes to check if the served and received data is the same +var testCrossProcessData, testCrossProcessSize = testutils.PrepareTestData(testCrossProcessPageSize, pagesInTestData) // Main process, FC in our case func TestCrossProcessDoubleRegistration(t *testing.T) { From 708d54bd534cb8fccfa91f6810df06bb198d9ac2 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:03:30 -0700 Subject: [PATCH 22/49] Remove conversion --- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 3f3dc86c95..be63bd17aa 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -36,7 +36,7 @@ func (u *userfaultfd) Serve( ) error { pollFds := []unix.PollFd{ {Fd: int32(u.fd), Events: unix.POLLIN}, - {Fd: int32(fdExit.Reader()), Events: unix.POLLIN}, + {Fd: fdExit.Reader(), Events: unix.POLLIN}, } var eg errgroup.Group From 2aa94495d20af7883a23fb3d36dbe22f5443ffe8 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:04:49 -0700 Subject: [PATCH 23/49] Make method public --- .../internal/sandbox/uffd/userfaultfd/userfaultfd.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index de6dcb4294..2d907bacd0 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -67,7 +67,7 @@ func (u *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error return nil } -func (u *userfaultfd) removeWriteProtection(addr uintptr, size uint64) error { +func (u *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { return u.writeProtect(addr, size, 0) } From dc7dc60576083fd003541bdf1c9ee06276d39700 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:06:12 -0700 Subject: [PATCH 24/49] Add command context --- .../sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 83d02dc824..9387a5dc63 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -54,7 +54,7 @@ func TestCrossProcessDoubleRegistration(t *testing.T) { sigc := make(chan os.Signal, 1) signal.Notify(sigc, syscall.SIGUSR1) - cmd := exec.Command(os.Args[0], "-test.run=TestHelperProcess") + cmd := exec.CommandContext(t.Context(), os.Args[0], "-test.run=TestHelperProcess") cmd.Env = append(os.Environ(), "GO_TEST_HELPER_PROCESS=1") cmd.Env = append(cmd.Env, fmt.Sprintf("GO_MMAP_START=%d", memoryStart)) From 6331bd2802660b15a271979f4d0e29107f8a5047 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:06:41 -0700 Subject: [PATCH 25/49] Fix error formatting --- .../internal/sandbox/uffd/userfaultfd/userfaultfd.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 2d907bacd0..061580bed6 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -50,7 +50,7 @@ func (u *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_REGISTER, uintptr(unsafe.Pointer(®ister))) if errno != 0 { - return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %v (ret=%d)", errno, ret) + return fmt.Errorf("UFFDIO_REGISTER ioctl failed: %w (ret=%d)", errno, ret) } return nil From 97ee756707da2bc711e5ec2f4f90300df6e5d1e1 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:07:15 -0700 Subject: [PATCH 26/49] Fix error formatting --- .../internal/sandbox/uffd/userfaultfd/userfaultfd.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 061580bed6..c90561a724 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -17,7 +17,7 @@ type userfaultfd struct { func newUserfaultfd(flags uintptr) (*userfaultfd, error) { uffd, _, errno := syscall.Syscall(NR_userfaultfd, flags, 0, 0) if errno != 0 { - return nil, fmt.Errorf("userfaultfd syscall failed: %v", errno) + return nil, fmt.Errorf("userfaultfd syscall failed: %w", errno) } return NewUserfaultfdFromFd(uffd), nil @@ -36,7 +36,7 @@ func (u *userfaultfd) configureApi(features CULong) error { api := NewUffdioAPI(UFFD_API, features) ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_API, uintptr(unsafe.Pointer(&api))) if errno != 0 { - return fmt.Errorf("UFFDIO_API ioctl failed: %v (ret=%d)", errno, ret) + return fmt.Errorf("UFFDIO_API ioctl failed: %w (ret=%d)", errno, ret) } return nil From 03b0c4677b68d02759ed795ecdff28902f0592b5 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:08:31 -0700 Subject: [PATCH 27/49] Improve error comparison --- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index be63bd17aa..6153f676a8 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -170,7 +170,7 @@ outerLoop: } copyErr := u.copy(addr, b, pagesize, copyMode) - if copyErr == unix.EEXIST { + if errors.Is(copyErr, unix.EEXIST) { logger.Debug("UFFD serve page already mapped", zap.Any("offset", addr), zap.Any("pagesize", pagesize)) // Page is already mapped From a331214d1dd1a24c27e337cb412b48e4b22303ab Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:22:12 -0700 Subject: [PATCH 28/49] Fix error formatting --- .../internal/sandbox/uffd/userfaultfd/userfaultfd.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index c90561a724..46d2f25529 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -61,7 +61,7 @@ func (u *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_WRITEPROTECT, uintptr(unsafe.Pointer(®ister))) if errno != 0 { - return fmt.Errorf("UFFDIO_WRITEPROTECT ioctl failed: %v (ret=%d)", errno, ret) + return fmt.Errorf("UFFDIO_WRITEPROTECT ioctl failed: %w (ret=%d)", errno, ret) } return nil From 6855c56c486e2d067ce96b22cbdbd9f39e6cd2ea Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 21:24:49 -0700 Subject: [PATCH 29/49] Change process exit --- .../userfaultfd/userfaultfd_cross_process_test.go | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 9387a5dc63..932a465c11 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -98,7 +98,8 @@ func TestHelperProcess(t *testing.T) { startRaw, err := strconv.Atoi(mmapStart) if err != nil { fmt.Print("exit parsing mmap start", err) - os.Exit(1) + + t.Fatal("exit parsing mmap start", err) } start := uintptr(startRaw) @@ -111,7 +112,8 @@ func TestHelperProcess(t *testing.T) { err = uffd.Register(start, testCrossProcessSize, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) if err != nil { fmt.Print("exit registering uffd", err) - os.Exit(1) + + t.Fatal("exit registering uffd", err) } fmt.Println("after register") @@ -124,7 +126,8 @@ func TestHelperProcess(t *testing.T) { fdExit, err := fdexit.New() if err != nil { fmt.Print("exit creating fd exit", err) - os.Exit(1) + + t.Fatal("exit creating fd exit", err) } defer fdExit.Close() @@ -138,6 +141,4 @@ func TestHelperProcess(t *testing.T) { }() time.Sleep(10 * time.Second) - - os.Exit(0) } From 359c17d54938e96ab2b41c604dcdb2ac3e03fd47 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 23:00:52 -0700 Subject: [PATCH 30/49] Trigger build From 835a32da9dbe438bf87cd5087f0075385a2e63d0 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sat, 4 Oct 2025 05:19:26 +0000 Subject: [PATCH 31/49] Enable unpriviledged uffd mode in GH PR tests --- .github/workflows/pr-tests.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/pr-tests.yml b/.github/workflows/pr-tests.yml index b72c17302a..28e27c5c22 100644 --- a/.github/workflows/pr-tests.yml +++ b/.github/workflows/pr-tests.yml @@ -33,6 +33,11 @@ jobs: with: job_key: unit-tests-${{ matrix.package }} + - name: Enable unprivileged uffd mode + cmd: | + sudo modprobe userfaultfd + sudo sysctl -w kernel.unprivileged_userns_clone=1 + - name: Run tests working-directory: ${{ matrix.package }} run: go test -v ${{ matrix.test_path }} From 76c3bf3c2e1f4cf64f2bf840ca08ee084a81cc31 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 23:16:19 -0700 Subject: [PATCH 32/49] Fix uffd unpriviledged enable --- .github/workflows/pr-tests.yml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pr-tests.yml b/.github/workflows/pr-tests.yml index 28e27c5c22..c8b4495a3f 100644 --- a/.github/workflows/pr-tests.yml +++ b/.github/workflows/pr-tests.yml @@ -34,9 +34,8 @@ jobs: job_key: unit-tests-${{ matrix.package }} - name: Enable unprivileged uffd mode - cmd: | - sudo modprobe userfaultfd - sudo sysctl -w kernel.unprivileged_userns_clone=1 + run: | + echo 1 | sudo tee /proc/sys/vm/unprivileged_userfaultfd - name: Run tests working-directory: ${{ matrix.package }} From 0ba2a14626ffb39217ca496684dd10268f40d54f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Fri, 3 Oct 2025 23:29:22 -0700 Subject: [PATCH 33/49] Use 4k pages in uffd cross process test --- .../uffd/userfaultfd/userfaultfd_cross_process_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go index 932a465c11..9f3430ec44 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go @@ -24,7 +24,7 @@ import ( const ( pagesInTestData = 32 - testCrossProcessPageSize = uint64(header.HugepageSize) + testCrossProcessPageSize = uint64(header.PageSize) inheritedFdIndex = 3 ) @@ -33,7 +33,7 @@ var testCrossProcessData, testCrossProcessSize = testutils.PrepareTestData(testC // Main process, FC in our case func TestCrossProcessDoubleRegistration(t *testing.T) { - memoryArea, memoryStart := testutils.New2MPageMmap(testCrossProcessSize) + memoryArea, memoryStart := testutils.New4KPageMmap(testCrossProcessSize) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { From 52d2fc70943ebede269bc516e8456b5a78c30b56 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 05:13:52 +0000 Subject: [PATCH 34/49] Add access checks to tests --- .../sandbox/uffd/testutils/contiguous_map.go | 19 ++++++++ .../uffd/userfaultfd/userfaultfd_test.go | 43 +++++++++++++------ 2 files changed, 49 insertions(+), 13 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go b/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go index 54377a2090..589f64ac09 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go @@ -1,11 +1,16 @@ package testutils +import "sync" + // 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 []uintptr + mu sync.RWMutex } func NewContiguousMap(start uintptr, size, pagesize uint64) *ContiguousMap { @@ -20,5 +25,19 @@ func (m *ContiguousMap) GetOffset(addr uintptr) (int64, uint64, error) { offset := addr - m.start pagesize := m.pagesize + m.mu.Lock() + m.accessedOffsets = append(m.accessedOffsets, offset) + m.mu.Unlock() + return int64(offset), pagesize, nil } + +func (m *ContiguousMap) Accessed() []uintptr { + m.mu.RLock() + defer m.mu.RUnlock() + + accessedOffsets := make([]uintptr, len(m.accessedOffsets)) + copy(accessedOffsets, m.accessedOffsets) + + return accessedOffsets +} diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 86541628a3..244876f064 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -3,6 +3,7 @@ package userfaultfd import ( "bytes" "fmt" + "reflect" "syscall" "testing" @@ -50,14 +51,20 @@ func TestUffdMissing(t *testing.T) { } }() - d, err := data.Slice(t.Context(), 0, int64(pagesize)) + accessOffset := uint64(0) + + d, err := data.Slice(t.Context(), int64(accessOffset), int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } - if !bytes.Equal(memoryArea[0:pagesize], d) { + if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) } + + if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { + t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) + } } func TestUffdWriteProtect(t *testing.T) { @@ -102,9 +109,13 @@ func TestUffdWriteProtect(t *testing.T) { } }() - memoryArea[0] = 'A' + accessOffset := 0 + + memoryArea[accessOffset] = 'A' - // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd + if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { + t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) + } } func TestUffdWriteProtectWithMissing(t *testing.T) { @@ -155,13 +166,17 @@ func TestUffdWriteProtectWithMissing(t *testing.T) { t.Fatal("cannot read content", err) } - if !bytes.Equal(memoryArea[0:pagesize], d) { + accessOffset := uint64(0) + + if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) } - memoryArea[0] = 'A' + memoryArea[accessOffset] = 'A' - // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd + if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { + t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) + } } // We are trying to simulate registering the missing handler in the FC and then registering the missing+wp handler again in the orchestrator @@ -188,8 +203,6 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { t.Fatal("failed to register memory", err) } - // TODO: Can we reregister after triggering missing and still properly handle such a page later? - // done little later in the orchestrator // both flags needs to be present err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) @@ -217,16 +230,20 @@ func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { } }() - d, err := data.Slice(t.Context(), 0, int64(pagesize)) + accessOffset := uint64(0) + + d, err := data.Slice(t.Context(), int64(accessOffset), int64(pagesize)) if err != nil { t.Fatal("cannot read content", err) } - if !bytes.Equal(memoryArea[0:pagesize], d) { + if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) } - memoryArea[0] = 'A' + memoryArea[accessOffset] = 'A' - // TODO: the write should be unblocked here, ideally we should also wait to check it was blocked then unblocked from the uffd + if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { + t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) + } } From f092687d9ee4476eb2f77237b14b31fc916f70ac Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 18:48:55 +0000 Subject: [PATCH 35/49] Cleanup --- .github/workflows/pr-tests.yml | 6 + .../sandbox/uffd/testutils/contiguous_map.go | 29 +- .../internal/sandbox/uffd/testutils/mmap.go | 21 +- .../sandbox/uffd/testutils/test_data.go | 51 +-- .../sandbox/uffd/userfaultfd/serve.go | 1 + .../sandbox/uffd/userfaultfd/userfaultfd.go | 30 +- .../userfaultfd_cross_process_test.go | 144 -------- .../uffd/userfaultfd/userfaultfd_test.go | 347 ++++++------------ packages/shared/pkg/utils/map_keys.go | 10 + 9 files changed, 204 insertions(+), 435 deletions(-) delete mode 100644 packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go create mode 100644 packages/shared/pkg/utils/map_keys.go diff --git a/.github/workflows/pr-tests.yml b/.github/workflows/pr-tests.yml index c8b4495a3f..3bf04ed6d7 100644 --- a/.github/workflows/pr-tests.yml +++ b/.github/workflows/pr-tests.yml @@ -37,6 +37,12 @@ jobs: 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 }} diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go b/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go index 589f64ac09..6eff07b789 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/contiguous_map.go @@ -1,6 +1,10 @@ package testutils -import "sync" +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. @@ -9,15 +13,16 @@ type ContiguousMap struct { size uint64 pagesize uint64 - accessedOffsets []uintptr + accessedOffsets map[uint64]struct{} mu sync.RWMutex } func NewContiguousMap(start uintptr, size, pagesize uint64) *ContiguousMap { return &ContiguousMap{ - start: start, - size: size, - pagesize: pagesize, + start: start, + size: size, + pagesize: pagesize, + accessedOffsets: make(map[uint64]struct{}), } } @@ -26,18 +31,16 @@ func (m *ContiguousMap) GetOffset(addr uintptr) (int64, uint64, error) { pagesize := m.pagesize m.mu.Lock() - m.accessedOffsets = append(m.accessedOffsets, offset) + m.accessedOffsets[uint64(offset)] = struct{}{} m.mu.Unlock() return int64(offset), pagesize, nil } -func (m *ContiguousMap) Accessed() []uintptr { - m.mu.RLock() - defer m.mu.RUnlock() - - accessedOffsets := make([]uintptr, len(m.accessedOffsets)) - copy(accessedOffsets, m.accessedOffsets) +func (m *ContiguousMap) Map() map[uint64]struct{} { + return m.accessedOffsets +} - return accessedOffsets +func (m *ContiguousMap) Keys() []uint64 { + return utils.MapKeys(m.accessedOffsets) } diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go b/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go index fc87df6bc9..20c8863de6 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go @@ -2,6 +2,7 @@ package testutils import ( "context" + "fmt" "math" "syscall" "unsafe" @@ -23,15 +24,19 @@ 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 NewPageMmap(size, pagesize uint64) ([]byte, uintptr, error) { + if pagesize == header.PageSize { + return newMockMmap(size, header.PageSize, 0) + } + + if pagesize == header.HugepageSize { + return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) + } -func New2MPageMmap(size uint64) ([]byte, uintptr) { - return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) + panic(fmt.Sprintf("unsupported page size: %d", pagesize)) } -func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr) { +func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr, error) { l := int(math.Ceil(float64(size)/float64(pagesize)) * float64(pagesize)) b, err := syscall.Mmap( -1, @@ -41,8 +46,8 @@ func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr) { syscall.MAP_PRIVATE|syscall.MAP_ANONYMOUS|flags, ) if err != nil { - return nil, 0 + return nil, 0, err } - return b, uintptr(unsafe.Pointer(&b[0])) + return b, uintptr(unsafe.Pointer(&b[0])), nil } diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go b/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go index 57991751b1..27f201470c 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go @@ -1,31 +1,40 @@ 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]) - } +import ( + "crypto/rand" +) - return dst -} - -func PrepareTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size uint64) { +func GenerateTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size uint64) { size = pagesize * pagesInTestData + n := int(size) + buf := make([]byte, n) + if _, err := rand.Read(buf); err != nil { + panic(err) + } + data = newMockSlicer( - repeatToSize( - []byte("Hello from userfaultfd! This is our test content that should be readable after the page fault."), - size, - ), + 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 := 0; i < minLen; i++ { + if a[i] != b[i] { + return i, b[i], a[i] + } + } + if len(a) != len(b) { + return minLen, 0, 0 + } + return -1, 0, 0 +} diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 6153f676a8..61836de8f3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -158,6 +158,7 @@ outerLoop: var copyMode CULong + // fmt.Fprintf(os.Stdout, "Serving page: %d %d\n", offset, pagesize) b, sliceErr := src.Slice(ctx, offset, int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go index 46d2f25529..0f345b12f2 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go @@ -5,6 +5,8 @@ import ( "fmt" "syscall" "unsafe" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) var ErrUnexpectedEventType = errors.New("unexpected event type") @@ -32,7 +34,14 @@ func NewUserfaultfdFromFd(fd uintptr) *userfaultfd { // features: UFFD_FEATURE_MISSING_HUGETLBFS // This is already called by the FC -func (u *userfaultfd) configureApi(features CULong) error { +func (u *userfaultfd) configureApi(pagesize uint64) error { + var features CULong + + // Only set the hugepage feature if we're using hugepages + if pagesize == header.HugepageSize { + features |= UFFD_FEATURE_MISSING_HUGETLBFS + } + api := NewUffdioAPI(UFFD_API, features) ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_API, uintptr(unsafe.Pointer(&api))) if errno != 0 { @@ -56,25 +65,6 @@ func (u *userfaultfd) Register(addr uintptr, size uint64, mode CULong) error { return nil } -func (u *userfaultfd) writeProtect(addr uintptr, size uint64, mode CULong) error { - register := NewUffdioWriteProtect(CULong(addr), CULong(size), mode) - - ret, _, errno := syscall.Syscall(syscall.SYS_IOCTL, u.fd, UFFDIO_WRITEPROTECT, uintptr(unsafe.Pointer(®ister))) - if errno != 0 { - return fmt.Errorf("UFFDIO_WRITEPROTECT ioctl failed: %w (ret=%d)", errno, ret) - } - - return nil -} - -func (u *userfaultfd) RemoveWriteProtection(addr uintptr, size uint64) error { - return u.writeProtect(addr, size, 0) -} - -func (u *userfaultfd) AddWriteProtection(addr uintptr, size uint64) error { - return u.writeProtect(addr, size, UFFDIO_WRITEPROTECT_MODE_WP) -} - // mode: UFFDIO_COPY_MODE_WP // When we use both missing and wp, we need to use UFFDIO_COPY_MODE_WP, otherwise copying would unprotect the page func (u *userfaultfd) copy(addr uintptr, data []byte, pagesize uint64, mode CULong) error { diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go deleted file mode 100644 index 9f3430ec44..0000000000 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_cross_process_test.go +++ /dev/null @@ -1,144 +0,0 @@ -package userfaultfd - -// This tests is creating uffd in a process and handling the page faults in another process. -// It also tests reregistering the uffd with the additional wp flag in the another process (in "orchestrator") after registering the missing handler already (in "FC"), -// simulating the case we have with the write protection being set up after FC already registered the uffd. - -import ( - "bytes" - "fmt" - "os" - "os/exec" - "os/signal" - "strconv" - "syscall" - "testing" - "time" - - "go.uber.org/zap" - - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/testutils" - "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" -) - -const ( - pagesInTestData = 32 - testCrossProcessPageSize = uint64(header.PageSize) - inheritedFdIndex = 3 -) - -// Data shared across the testing processes to check if the served and received data is the same -var testCrossProcessData, testCrossProcessSize = testutils.PrepareTestData(testCrossProcessPageSize, pagesInTestData) - -// Main process, FC in our case -func TestCrossProcessDoubleRegistration(t *testing.T) { - memoryArea, memoryStart := testutils.New4KPageMmap(testCrossProcessSize) - - uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } - defer uffd.Close() - - err = uffd.configureApi(0) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } - - err = uffd.Register(memoryStart, testCrossProcessSize, UFFDIO_REGISTER_MODE_MISSING) - if err != nil { - t.Fatal("failed to register memory", err) - } - - sigc := make(chan os.Signal, 1) - signal.Notify(sigc, syscall.SIGUSR1) - - cmd := exec.CommandContext(t.Context(), os.Args[0], "-test.run=TestHelperProcess") - cmd.Env = append(os.Environ(), "GO_TEST_HELPER_PROCESS=1") - cmd.Env = append(cmd.Env, fmt.Sprintf("GO_MMAP_START=%d", memoryStart)) - - // Passing the fd to the child process - uffdFile := os.NewFile(uffd.fd, "userfaultfd") - - cmd.ExtraFiles = []*os.File{uffdFile} - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - - if err := cmd.Start(); err != nil { - t.Fatalf("failed to start helper: %v", err) - } - defer func() { - cmd.Process.Kill() - cmd.Wait() - }() - - <-sigc - fmt.Println("child signaled ready") - - data := testCrossProcessData - - servedContent, err := data.Slice(t.Context(), 0, int64(testCrossProcessPageSize)) - if err != nil { - t.Fatal("cannot read content", err) - } - - if !bytes.Equal(memoryArea[0:testCrossProcessPageSize], servedContent) { - t.Fatal("content mismatch", string(servedContent)) - } -} - -// Secondary process, orchestrator in our case -func TestHelperProcess(t *testing.T) { - if os.Getenv("GO_TEST_HELPER_PROCESS") != "1" { - return - } - - mmapStart := os.Getenv("GO_MMAP_START") - startRaw, err := strconv.Atoi(mmapStart) - if err != nil { - fmt.Print("exit parsing mmap start", err) - - t.Fatal("exit parsing mmap start", err) - } - - start := uintptr(startRaw) - - uffdFile := os.NewFile(uintptr(inheritedFdIndex), "userfaultfd") - uffd := NewUserfaultfdFromFd(uffdFile.Fd()) - - // done in the FC - // Check: The reregistration works - err = uffd.Register(start, testCrossProcessSize, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) - if err != nil { - fmt.Print("exit registering uffd", err) - - t.Fatal("exit registering uffd", err) - } - - fmt.Println("after register") - - ppid := os.Getppid() - syscall.Kill(ppid, syscall.SIGUSR1) - - m := testutils.NewContiguousMap(start, testCrossProcessSize, testCrossProcessPageSize) - - fdExit, err := fdexit.New() - if err != nil { - fmt.Print("exit creating fd exit", err) - - t.Fatal("exit creating fd exit", err) - } - defer fdExit.Close() - - data := testCrossProcessData - - go func() { - err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestHelperProcess] failed to serve uffd", err) - } - }() - - time.Sleep(10 * time.Second) -} diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 244876f064..8101e0a6e1 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -14,236 +14,125 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) -func TestUffdMissing(t *testing.T) { - pagesize := uint64(header.PageSize) - data, size := testutils.PrepareTestData(pagesize, pagesInTestData) - - uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } - defer uffd.Close() - - err = uffd.configureApi(0) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } - - memoryArea, memoryStart := testutils.New4KPageMmap(size) - - err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) - if err != nil { - t.Fatal("failed to register memory", err) - } - - m := testutils.NewContiguousMap(memoryStart, size, pagesize) - - fdExit, err := fdexit.New() - if err != nil { - t.Fatal("failed to create fd exit", err) - } - defer fdExit.Close() - - go func() { - err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestUffdMissing] failed to serve uffd", err) - } - }() - - accessOffset := uint64(0) - - d, err := data.Slice(t.Context(), int64(accessOffset), int64(pagesize)) - if err != nil { - t.Fatal("cannot read content", err) - } - - if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { - t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) - } - - if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { - t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) - } -} - -func TestUffdWriteProtect(t *testing.T) { - pagesize := uint64(header.PageSize) - data, size := testutils.PrepareTestData(pagesize, pagesInTestData) - - uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } - defer uffd.Close() - - err = uffd.configureApi(0) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } - - memoryArea, memoryStart := testutils.New4KPageMmap(size) - - err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_WP) - if err != nil { - t.Fatal("failed to register memory", err) - } - - err = uffd.AddWriteProtection(memoryStart, size) - if err != nil { - t.Fatal("failed to write protect memory", err) - } - - fdExit, err := fdexit.New() - if err != nil { - t.Fatal("failed to create fd exit", err) - } - defer fdExit.Close() - - m := testutils.NewContiguousMap(memoryStart, size, pagesize) - - go func() { - err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) - } - }() - - accessOffset := 0 - - memoryArea[accessOffset] = 'A' - - if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { - t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) - } +type pageTest struct { + pagesize uint64 + pagesInData uint64 + operationOffset uint64 } -func TestUffdWriteProtectWithMissing(t *testing.T) { - pagesize := uint64(header.PageSize) - - data, size := testutils.PrepareTestData(pagesize, pagesInTestData) - - uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } - defer uffd.Close() - - err = uffd.configureApi(0) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } - - memoryArea, memoryStart := testutils.New4KPageMmap(size) - - err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) - if err != nil { - t.Fatal("failed to register memory", err) - } - - err = uffd.AddWriteProtection(memoryStart, size) - if err != nil { - t.Fatal("failed to write protect memory", err) - } - - m := testutils.NewContiguousMap(memoryStart, size, pagesize) - - fdExit, err := fdexit.New() - if err != nil { - t.Fatal("failed to create fd exit", err) - } - defer fdExit.Close() - - go func() { - err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) - } - }() - - d, err := data.Slice(t.Context(), 0, int64(pagesize)) - if err != nil { - t.Fatal("cannot read content", err) - } - - accessOffset := uint64(0) - - if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { - t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) - } - - memoryArea[accessOffset] = 'A' - - if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { - t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) - } -} - -// We are trying to simulate registering the missing handler in the FC and then registering the missing+wp handler again in the orchestrator -func TestUffdWriteProtectWithMissingDoubleRegistration(t *testing.T) { - pagesize := uint64(header.PageSize) - data, size := testutils.PrepareTestData(pagesize, pagesInTestData) - - uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } - defer uffd.Close() - - err = uffd.configureApi(0) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } - - memoryArea, memoryStart := testutils.New4KPageMmap(size) - - // done in the FC - err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) - if err != nil { - t.Fatal("failed to register memory", err) - } - - // done little later in the orchestrator - // both flags needs to be present - err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING|UFFDIO_REGISTER_MODE_WP) - if err != nil { - t.Fatal("failed to register memory", err) - } - - err = uffd.AddWriteProtection(memoryStart, size) - if err != nil { - t.Fatal("failed to write protect memory", err) - } - - m := testutils.NewContiguousMap(memoryStart, size, pagesize) - - fdExit, err := fdexit.New() - if err != nil { - t.Fatal("failed to create fd exit", err) - } - defer fdExit.Close() - - go func() { - err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestUffdWriteProtect] failed to serve uffd", err) - } - }() - - accessOffset := uint64(0) - - d, err := data.Slice(t.Context(), int64(accessOffset), int64(pagesize)) - if err != nil { - t.Fatal("cannot read content", err) - } - - if !bytes.Equal(memoryArea[accessOffset:accessOffset+pagesize], d) { - t.Fatalf("content mismatch: want %q, got %q", d, memoryArea[:pagesize]) - } - - memoryArea[accessOffset] = 'A' - - if !reflect.DeepEqual(m.Accessed(), []uintptr{uintptr(accessOffset)}) { - t.Fatalf("accessed mismatch: want %v, got %v", []uintptr{uintptr(accessOffset)}, m.Accessed()) +func TestUffdMissing(t *testing.T) { + tests := []pageTest{ + // Standard 4K page, operation at start + { + pagesize: header.PageSize, + pagesInData: 32, + operationOffset: 0, + }, + // Standard 4K page, operation at middle + { + pagesize: header.PageSize, + pagesInData: 32, + operationOffset: 16 * header.PageSize, + }, + // Standard 4K page, operation at last page + { + pagesize: header.PageSize, + pagesInData: 32, + operationOffset: 31 * header.PageSize, + }, + // Hugepage, operation at start + { + pagesize: header.HugepageSize, + pagesInData: 8, + operationOffset: 0, + }, + // Hugepage, operation at middle + { + pagesize: header.HugepageSize, + pagesInData: 8, + operationOffset: 4 * header.HugepageSize, + }, + // Hugepage, operation at last page + { + pagesize: header.HugepageSize, + pagesInData: 8, + operationOffset: 7 * header.HugepageSize, + }, + } + + for _, tt := range tests { + t.Run(fmt.Sprintf("pagesize-%d-offset-%d", tt.pagesize, tt.operationOffset), func(t *testing.T) { + data, size := testutils.GenerateTestData(tt.pagesize, tt.pagesInData) + + uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) + if err != nil { + t.Fatal("failed to create userfaultfd", err) + } + t.Cleanup(func() { + uffd.Close() + }) + + err = uffd.configureApi(tt.pagesize) + if err != nil { + t.Fatal("failed to configure uffd api", err) + } + + memoryArea, memoryStart, err := testutils.NewPageMmap(size, tt.pagesize) + if err != nil { + t.Fatal("failed to create page mmap", err) + } + + err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) + if err != nil { + t.Fatal("failed to register memory", err) + } + + m := testutils.NewContiguousMap(memoryStart, size, tt.pagesize) + + fdExit, err := fdexit.New() + if err != nil { + t.Fatal("failed to create fd exit", err) + } + t.Cleanup(func() { + fdExit.SignalExit() + fdExit.Close() + }) + + exitUffd := make(chan struct{}, 1) + + go func() { + err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) + if err != nil { + fmt.Println("[TestUffdMissing] failed to serve uffd", err) + } + + exitUffd <- struct{}{} + }() + + d, err := data.Slice(t.Context(), int64(tt.operationOffset), int64(tt.pagesize)) + if err != nil { + t.Fatal("cannot read content", err) + } + + if !bytes.Equal(memoryArea[tt.operationOffset:tt.operationOffset+tt.pagesize], d) { + idx, want, got := testutils.DiffByte(memoryArea[tt.operationOffset:tt.operationOffset+tt.pagesize], d) + t.Fatalf("content mismatch: want %q, got %q at index %d", want, got, idx) + } + + if !reflect.DeepEqual(m.Map(), map[uint64]struct{}{tt.operationOffset: {}}) { + t.Fatalf("accessed mismatch: should be accessed %v, actually accessed %v", []uint64{tt.operationOffset}, m.Keys()) + } + + signalExitErr := fdExit.SignalExit() + if signalExitErr != nil { + t.Fatal("failed to signal exit", err) + } + + select { + case <-exitUffd: + case <-t.Context().Done(): + t.Fatal("timeout waiting for uffd to exit") + } + }) } } diff --git a/packages/shared/pkg/utils/map_keys.go b/packages/shared/pkg/utils/map_keys.go new file mode 100644 index 0000000000..a2d4b410d2 --- /dev/null +++ b/packages/shared/pkg/utils/map_keys.go @@ -0,0 +1,10 @@ +package utils + +func MapKeys[K comparable](m map[K]struct{}) []K { + keys := make([]K, 0, len(m)) + for key := range m { + keys = append(keys, key) + } + + return keys +} From 3f7216a24ee6782aaf8c052cb6b3ee9383208492 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 18:53:20 +0000 Subject: [PATCH 36/49] Remove unused constants for now --- .../sandbox/uffd/userfaultfd/constants.go | 35 ++++--------------- 1 file changed, 7 insertions(+), 28 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go index 274fb7b77a..657fb14ec4 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/constants.go @@ -27,20 +27,10 @@ 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_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 + UFFDIO_API = C.UFFDIO_API + UFFDIO_REGISTER = C.UFFDIO_REGISTER + UFFDIO_COPY = C.UFFDIO_COPY UFFD_FEATURE_MISSING_HUGETLBFS = C.UFFD_FEATURE_MISSING_HUGETLBFS ) @@ -53,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 { @@ -87,16 +76,6 @@ func NewUffdioCopy(b []byte, address CULong, pagesize CULong, mode CULong, bytes } } -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 } From 7fb669e350352640801e8b1ff09779118fc691a5 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 18:59:29 +0000 Subject: [PATCH 37/49] Fix lint issue --- .../internal/sandbox/uffd/testutils/test_data.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go b/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go index 27f201470c..7a01d2b8cb 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go @@ -13,9 +13,7 @@ func GenerateTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size panic(err) } - data = newMockSlicer( - buf, - ) + data = newMockSlicer(buf) return data, size } @@ -28,13 +26,16 @@ func DiffByte(a, b []byte) (idx int, want, got byte) { if len(b) < minLen { minLen = len(b) } - for i := 0; i < minLen; i++ { + + 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 } From 052b09f78ca33d16b14cf3ff4566aa12e852fc4f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 19:05:34 +0000 Subject: [PATCH 38/49] Clarify naming --- .../sandbox/uffd/testutils/content_slicer.go | 17 +++++++++++++++++ .../uffd/testutils/{mmap.go => page_mmap.go} | 19 +++---------------- .../{test_data.go => random_data.go} | 6 +++--- .../uffd/userfaultfd/userfaultfd_test.go | 16 ++++++++-------- 4 files changed, 31 insertions(+), 27 deletions(-) create mode 100644 packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go rename packages/orchestrator/internal/sandbox/uffd/testutils/{mmap.go => page_mmap.go} (57%) rename packages/orchestrator/internal/sandbox/uffd/testutils/{test_data.go => random_data.go} (79%) diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go b/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go new file mode 100644 index 0000000000..f678401446 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go @@ -0,0 +1,17 @@ +package testutils + +import "context" + +// contentSlicer is a slicer that returns content via exposing 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 +} diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go similarity index 57% rename from packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go rename to packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go index 20c8863de6..ca1801280a 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/mmap.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go @@ -1,7 +1,6 @@ package testutils import ( - "context" "fmt" "math" "syscall" @@ -12,31 +11,19 @@ import ( "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 NewPageMmap(size, pagesize uint64) ([]byte, uintptr, error) { if pagesize == header.PageSize { - return newMockMmap(size, header.PageSize, 0) + return newMmap(size, header.PageSize, 0) } if pagesize == header.HugepageSize { - return newMockMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) + return newMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) } panic(fmt.Sprintf("unsupported page size: %d", pagesize)) } -func newMockMmap(size, pagesize uint64, flags int) ([]byte, uintptr, error) { +func newMmap(size, pagesize uint64, flags int) ([]byte, uintptr, error) { l := int(math.Ceil(float64(size)/float64(pagesize)) * float64(pagesize)) b, err := syscall.Mmap( -1, diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go b/packages/orchestrator/internal/sandbox/uffd/testutils/random_data.go similarity index 79% rename from packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go rename to packages/orchestrator/internal/sandbox/uffd/testutils/random_data.go index 7a01d2b8cb..780b080b4c 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/test_data.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/random_data.go @@ -4,8 +4,8 @@ import ( "crypto/rand" ) -func GenerateTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size uint64) { - size = pagesize * pagesInTestData +func RandomPages(pagesize, numberOfPages uint64) (data *contentSlicer, size uint64) { + size = pagesize * numberOfPages n := int(size) buf := make([]byte, n) @@ -13,7 +13,7 @@ func GenerateTestData(pagesize, pagesInTestData uint64) (data *mockSlicer, size panic(err) } - data = newMockSlicer(buf) + data = newContentSlicer(buf) return data, size } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 8101e0a6e1..5e38cb1ba3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -16,7 +16,7 @@ import ( type pageTest struct { pagesize uint64 - pagesInData uint64 + numberOfPages uint64 operationOffset uint64 } @@ -25,44 +25,44 @@ func TestUffdMissing(t *testing.T) { // Standard 4K page, operation at start { pagesize: header.PageSize, - pagesInData: 32, + numberOfPages: 32, operationOffset: 0, }, // Standard 4K page, operation at middle { pagesize: header.PageSize, - pagesInData: 32, + numberOfPages: 32, operationOffset: 16 * header.PageSize, }, // Standard 4K page, operation at last page { pagesize: header.PageSize, - pagesInData: 32, + numberOfPages: 32, operationOffset: 31 * header.PageSize, }, // Hugepage, operation at start { pagesize: header.HugepageSize, - pagesInData: 8, + numberOfPages: 8, operationOffset: 0, }, // Hugepage, operation at middle { pagesize: header.HugepageSize, - pagesInData: 8, + numberOfPages: 8, operationOffset: 4 * header.HugepageSize, }, // Hugepage, operation at last page { pagesize: header.HugepageSize, - pagesInData: 8, + numberOfPages: 8, operationOffset: 7 * header.HugepageSize, }, } for _, tt := range tests { t.Run(fmt.Sprintf("pagesize-%d-offset-%d", tt.pagesize, tt.operationOffset), func(t *testing.T) { - data, size := testutils.GenerateTestData(tt.pagesize, tt.pagesInData) + data, size := testutils.RandomPages(tt.pagesize, tt.numberOfPages) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) if err != nil { From 1066947c4ba032e7c04060eda660f8eb9bad230e Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 19:08:28 +0000 Subject: [PATCH 39/49] Clarify comment --- .../internal/sandbox/uffd/testutils/content_slicer.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go b/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go index f678401446..f87e6365f1 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/content_slicer.go @@ -2,7 +2,7 @@ package testutils import "context" -// contentSlicer is a slicer that returns content via exposing the Slicer interface. +// contentSlicer exposes byte slice via the Slicer interface. // This is used for testing purposes. type contentSlicer struct { content []byte From d06985c24a0c6cf46ef32192cfbce993c99ff0cc Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 19:17:45 +0000 Subject: [PATCH 40/49] Remove yet non-relevant diagram --- .../internal/sandbox/uffd/userfaultfd/serve.go | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index 61836de8f3..af40f32186 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -1,16 +1,5 @@ package userfaultfd -/* -This is the flow of the UFFD events when using write protection: - -```mermaid -flowchart TD -A[missing page] -- write (WRITE flag) --> B(COPY) --> C[mark as dirty] -A -- read (0 flag) --> D(COPY + WP protect) --> E[faulted page] -E -- write (WP|WRITE flag) --> F(remove WP) --> C -``` -*/ - import ( "context" "errors" From 179ba743124dca560a51aad5b5d3c11836ea3278 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 19:19:40 +0000 Subject: [PATCH 41/49] Remove comment --- packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index af40f32186..c2cf1b78b3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -147,7 +147,6 @@ outerLoop: var copyMode CULong - // fmt.Fprintf(os.Stdout, "Serving page: %d %d\n", offset, pagesize) b, sliceErr := src.Slice(ctx, offset, int64(pagesize)) if sliceErr != nil { signalErr := fdExit.SignalExit() From bacedce8b0271625ce2c9e5c753ffb13035f65a0 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Sun, 5 Oct 2025 19:21:54 +0000 Subject: [PATCH 42/49] Remove write protection field --- packages/orchestrator/internal/sandbox/uffd/uffd.go | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/uffd.go b/packages/orchestrator/internal/sandbox/uffd/uffd.go index 0449161da0..93a047456d 100644 --- a/packages/orchestrator/internal/sandbox/uffd/uffd.go +++ b/packages/orchestrator/internal/sandbox/uffd/uffd.go @@ -40,8 +40,6 @@ type Uffd struct { memfile *block.TrackedSliceDevice socketPath string - - writeProtection bool } var _ MemoryBackend = (*Uffd)(nil) @@ -58,12 +56,11 @@ 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, - writeProtection: true, + exit: utils.NewErrorOnce(), + readyCh: make(chan struct{}, 1), + fdExit: fdExit, + memfile: trackedMemfile, + socketPath: socketPath, }, nil } From 5d4c576e018e9675894b42f5c8537e597f28da46 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Mon, 6 Oct 2025 18:00:09 +0000 Subject: [PATCH 43/49] Add explicit mmap cleanup --- .../internal/sandbox/uffd/testutils/page_mmap.go | 14 +++++++++----- .../sandbox/uffd/userfaultfd/userfaultfd_test.go | 5 ++++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go index ca1801280a..a61339a916 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go @@ -11,7 +11,7 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) -func NewPageMmap(size, pagesize uint64) ([]byte, uintptr, error) { +func NewPageMmap(size, pagesize uint64) ([]byte, uintptr, func() error, error) { if pagesize == header.PageSize { return newMmap(size, header.PageSize, 0) } @@ -20,10 +20,10 @@ func NewPageMmap(size, pagesize uint64) ([]byte, uintptr, error) { return newMmap(size, header.HugepageSize, unix.MAP_HUGETLB|unix.MAP_HUGE_2MB) } - panic(fmt.Sprintf("unsupported page size: %d", pagesize)) + return nil, 0, nil, fmt.Errorf("unsupported page size: %d", pagesize) } -func newMmap(size, pagesize uint64, flags int) ([]byte, uintptr, error) { +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, @@ -33,8 +33,12 @@ func newMmap(size, pagesize uint64, flags int) ([]byte, uintptr, error) { syscall.MAP_PRIVATE|syscall.MAP_ANONYMOUS|flags, ) if err != nil { - return nil, 0, err + return nil, 0, nil, err } - return b, uintptr(unsafe.Pointer(&b[0])), nil + close := func() error { + return syscall.Munmap(b) + } + + return b, uintptr(unsafe.Pointer(&b[0])), close, nil } diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 5e38cb1ba3..efee604cf2 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -77,10 +77,13 @@ func TestUffdMissing(t *testing.T) { t.Fatal("failed to configure uffd api", err) } - memoryArea, memoryStart, err := testutils.NewPageMmap(size, tt.pagesize) + memoryArea, memoryStart, unmap, err := testutils.NewPageMmap(size, tt.pagesize) if err != nil { t.Fatal("failed to create page mmap", err) } + t.Cleanup(func() { + unmap() + }) err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) if err != nil { From fe1f562089c34b9e01cc1d702f6d729102bd60c8 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Mon, 6 Oct 2025 18:52:27 +0000 Subject: [PATCH 44/49] Improve test names --- .../sandbox/uffd/userfaultfd/userfaultfd_test.go | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index efee604cf2..adbedeb5a7 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -15,6 +15,7 @@ import ( ) type pageTest struct { + name string pagesize uint64 numberOfPages uint64 operationOffset uint64 @@ -22,38 +23,38 @@ type pageTest struct { func TestUffdMissing(t *testing.T) { tests := []pageTest{ - // Standard 4K page, operation at start { + name: "standard 4k page, operation at start", pagesize: header.PageSize, numberOfPages: 32, operationOffset: 0, }, - // Standard 4K page, operation at middle { + name: "standard 4k page, operation at middle", pagesize: header.PageSize, numberOfPages: 32, operationOffset: 16 * header.PageSize, }, - // Standard 4K page, operation at last page { + name: "standard 4k page, operation at last page", pagesize: header.PageSize, numberOfPages: 32, operationOffset: 31 * header.PageSize, }, - // Hugepage, operation at start { + name: "hugepage, operation at start", pagesize: header.HugepageSize, numberOfPages: 8, operationOffset: 0, }, - // Hugepage, operation at middle { + name: "hugepage, operation at middle", pagesize: header.HugepageSize, numberOfPages: 8, operationOffset: 4 * header.HugepageSize, }, - // Hugepage, operation at last page { + name: "hugepage, operation at last page", pagesize: header.HugepageSize, numberOfPages: 8, operationOffset: 7 * header.HugepageSize, @@ -61,7 +62,7 @@ func TestUffdMissing(t *testing.T) { } for _, tt := range tests { - t.Run(fmt.Sprintf("pagesize-%d-offset-%d", tt.pagesize, tt.operationOffset), func(t *testing.T) { + t.Run(tt.name, func(t *testing.T) { data, size := testutils.RandomPages(tt.pagesize, tt.numberOfPages) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) From 70015bd73987b4143369cda05ca0b58992ce89bd Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Mon, 6 Oct 2025 19:27:55 +0000 Subject: [PATCH 45/49] Fix test error message --- .../internal/sandbox/uffd/userfaultfd/userfaultfd_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index adbedeb5a7..16437644b1 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -135,7 +135,7 @@ func TestUffdMissing(t *testing.T) { select { case <-exitUffd: case <-t.Context().Done(): - t.Fatal("timeout waiting for uffd to exit") + t.Fatal("context done", t.Context().Err()) } }) } From 69208763eec9bdd192be0c0e08c55f7cae4760ff Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Mon, 6 Oct 2025 20:40:26 +0000 Subject: [PATCH 46/49] Cleanup tests --- .../uffd/userfaultfd/userfaultfd_test.go | 45 +++++++------------ 1 file changed, 15 insertions(+), 30 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 16437644b1..720fb52743 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -2,11 +2,11 @@ package userfaultfd import ( "bytes" - "fmt" - "reflect" "syscall" "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "go.uber.org/zap" "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit" @@ -66,37 +66,30 @@ func TestUffdMissing(t *testing.T) { data, size := testutils.RandomPages(tt.pagesize, tt.numberOfPages) uffd, err := newUserfaultfd(syscall.O_CLOEXEC | syscall.O_NONBLOCK) - if err != nil { - t.Fatal("failed to create userfaultfd", err) - } + require.NoError(t, err) + t.Cleanup(func() { uffd.Close() }) err = uffd.configureApi(tt.pagesize) - if err != nil { - t.Fatal("failed to configure uffd api", err) - } + require.NoError(t, err) memoryArea, memoryStart, unmap, err := testutils.NewPageMmap(size, tt.pagesize) - if err != nil { - t.Fatal("failed to create page mmap", err) - } + require.NoError(t, err) + t.Cleanup(func() { unmap() }) err = uffd.Register(memoryStart, size, UFFDIO_REGISTER_MODE_MISSING) - if err != nil { - t.Fatal("failed to register memory", err) - } + require.NoError(t, err) m := testutils.NewContiguousMap(memoryStart, size, tt.pagesize) fdExit, err := fdexit.New() - if err != nil { - t.Fatal("failed to create fd exit", err) - } + require.NoError(t, err) + t.Cleanup(func() { fdExit.SignalExit() fdExit.Close() @@ -106,36 +99,28 @@ func TestUffdMissing(t *testing.T) { go func() { err := uffd.Serve(t.Context(), m, data, fdExit, zap.L()) - if err != nil { - fmt.Println("[TestUffdMissing] failed to serve uffd", err) - } + assert.NoError(t, err) exitUffd <- struct{}{} }() d, err := data.Slice(t.Context(), int64(tt.operationOffset), int64(tt.pagesize)) - if err != nil { - t.Fatal("cannot read content", err) - } + require.NoError(t, err) if !bytes.Equal(memoryArea[tt.operationOffset:tt.operationOffset+tt.pagesize], d) { idx, want, got := testutils.DiffByte(memoryArea[tt.operationOffset:tt.operationOffset+tt.pagesize], d) t.Fatalf("content mismatch: want %q, got %q at index %d", want, got, idx) } - if !reflect.DeepEqual(m.Map(), map[uint64]struct{}{tt.operationOffset: {}}) { - t.Fatalf("accessed mismatch: should be accessed %v, actually accessed %v", []uint64{tt.operationOffset}, m.Keys()) - } + assert.Equal(t, m.Map(), map[uint64]struct{}{tt.operationOffset: {}}) signalExitErr := fdExit.SignalExit() - if signalExitErr != nil { - t.Fatal("failed to signal exit", err) - } + require.NoError(t, signalExitErr) select { case <-exitUffd: case <-t.Context().Done(): - t.Fatal("context done", t.Context().Err()) + t.Fatal("context done before exit", t.Context().Err()) } }) } From cb5cfdbbc52750140e0ede5dac199f60d6a7a76f Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Mon, 6 Oct 2025 20:55:49 +0000 Subject: [PATCH 47/49] Fix lint error --- .../internal/sandbox/uffd/userfaultfd/userfaultfd_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go index 720fb52743..9e9af112b3 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd_test.go @@ -112,7 +112,7 @@ func TestUffdMissing(t *testing.T) { t.Fatalf("content mismatch: want %q, got %q at index %d", want, got, idx) } - assert.Equal(t, m.Map(), map[uint64]struct{}{tt.operationOffset: {}}) + assert.Equal(t, map[uint64]struct{}{tt.operationOffset: {}}, m.Map()) signalExitErr := fdExit.SignalExit() require.NoError(t, signalExitErr) From ac77fe558da675a9d9355093cb23cda45b77432c Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Tue, 7 Oct 2025 22:12:00 +0000 Subject: [PATCH 48/49] Put back offset log on uffd panic --- .../orchestrator/internal/sandbox/uffd/userfaultfd/serve.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go index c2cf1b78b3..ef161b6815 100644 --- a/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go +++ b/packages/orchestrator/internal/sandbox/uffd/userfaultfd/serve.go @@ -141,7 +141,7 @@ outerLoop: eg.Go(func() error { defer func() { if r := recover(); r != nil { - logger.Error("UFFD serve panic", zap.Any("pagesize", pagesize), zap.Any("panic", r)) + logger.Error("UFFD serve panic", zap.Any("offset", offset), zap.Any("pagesize", pagesize), zap.Any("panic", r)) } }() From c11e5ef7df61136c54fdda1cf065f323f546e5e5 Mon Sep 17 00:00:00 2001 From: Tomas Valenta Date: Wed, 8 Oct 2025 01:40:58 +0000 Subject: [PATCH 49/49] Fix lint errors --- .../orchestrator/internal/sandbox/template/mask_template.go | 2 +- .../orchestrator/internal/sandbox/uffd/testutils/page_mmap.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/orchestrator/internal/sandbox/template/mask_template.go b/packages/orchestrator/internal/sandbox/template/mask_template.go index d1df5e9751..1e25ce6e36 100644 --- a/packages/orchestrator/internal/sandbox/template/mask_template.go +++ b/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() } diff --git a/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go index a61339a916..073a21c5c0 100644 --- a/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go +++ b/packages/orchestrator/internal/sandbox/uffd/testutils/page_mmap.go @@ -36,9 +36,9 @@ func newMmap(size, pagesize uint64, flags int) ([]byte, uintptr, func() error, e return nil, 0, nil, err } - close := func() error { + closeMmap := func() error { return syscall.Munmap(b) } - return b, uintptr(unsafe.Pointer(&b[0])), close, nil + return b, uintptr(unsafe.Pointer(&b[0])), closeMmap, nil }