Fix fast_page_fault_helper crashing instead of causing page faults

resources/overlay/usr/local/bin/fast_page_fault_helper.c

@@ -20,22 +20,27 @@
 int main(int argc, char *const argv[]) {
     sigset_t set;
     int signal;
+    void *ptr;
 
     sigemptyset(&set);
-    if(sigaddset(&set, SIGUSR1) == -1) {
+    if (sigaddset(&set, SIGUSR1) == -1) {
         perror("sigaddset");
-        return -1;
+        return 1;
+    }
+    if (sigprocmask(SIG_BLOCK, &set, NULL) == -1)  {
+        perror("sigprocmask");
+        return 1;
     }
 
-    void *ptr = mmap(NULL, MEM_SIZE_MIB, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
-
-    memset(ptr, 1, MEM_SIZE_MIB);
+    ptr = mmap(NULL, MEM_SIZE_MIB, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
 
-    if(MAP_FAILED == ptr) {
+    if (MAP_FAILED == ptr) {
         perror("mmap");
-        return -1;
+        return 1;
     }
 
+    memset(ptr, 1, MEM_SIZE_MIB);
+
     sigwait(&set, &signal);
 
     memset(ptr, 2, MEM_SIZE_MIB);

tests/conftest.py

@@ -252,16 +252,6 @@ def bin_seccomp_paths():
     yield demos
 
 
-@pytest.fixture
-def uffd_handler_paths():
-    """Build UFFD handler binaries."""
-    handlers = {
-        f"{handler}_handler": build_tools.get_example(f"uffd_{handler}_handler")
-        for handler in ["malicious", "valid", "fault_all"]
-    }
-    yield handlers
-
-
 @pytest.fixture(scope="session")
 def netns_factory(worker_id):
     """A network namespace factory

tests/framework/utils.py

@@ -10,7 +10,6 @@
 import re
 import select
 import signal
-import stat
 import subprocess
 import time
 import typing
@@ -133,69 +132,6 @@ def chroot(path):
         os.chdir(working_dir)
 
 
-class UffdHandler:
-    """Describe the UFFD page fault handler process."""
-
-    def __init__(self, name, socket_path, mem_file, chroot_path, log_file_name):
-        """Instantiate the handler process with arguments."""
-        self._proc = None
-        self._handler_name = name
-        self._socket_path = socket_path
-        self._mem_file = mem_file
-        self._chroot = chroot_path
-        self._log_file = log_file_name
-
-    def spawn(self, uid, gid):
-        """Spawn handler process using arguments provided."""
-
-        with chroot(self._chroot):
-            st = os.stat(self._handler_name)
-            os.chmod(self._handler_name, st.st_mode | stat.S_IEXEC)
-
-            chroot_log_file = Path("/") / self._log_file
-            with open(chroot_log_file, "w", encoding="utf-8") as logfile:
-                args = [f"/{self._handler_name}", self._socket_path, self._mem_file]
-                self._proc = subprocess.Popen(
-                    args, stdout=logfile, stderr=subprocess.STDOUT
-                )
-
-            # Give it time start and fail, if it really has too (bad things happen).
-            time.sleep(1)
-            if not self.is_running():
-                print(chroot_log_file.read_text(encoding="utf-8"))
-                assert False, "Could not start PF handler!"
-
-            # The page fault handler will create the socket path with root rights.
-            # Change rights to the jailer's.
-            os.chown(self._socket_path, uid, gid)
-
-    @property
-    def proc(self):
-        """Return UFFD handler process."""
-        return self._proc
-
-    def is_running(self):
-        """Check if UFFD process is running"""
-        return self.proc is not None and self.proc.poll() is None
-
-    @property
-    def log_file(self):
-        """Return the path to the UFFD handler's log file"""
-        return Path(self._chroot) / Path(self._log_file)
-
-    @property
-    def log_data(self):
-        """Return the log data of the UFFD handler"""
-        if self.log_file is None:
-            return ""
-        return self.log_file.read_text(encoding="utf-8")
-
-    def __del__(self):
-        """Tear down the UFFD handler process."""
-        if self.proc is not None:
-            self.proc.kill()
-
-
 class CpuMap:
     """Cpu map from real cpu cores to containers visible cores.
 

tests/framework/utils_uffd.py

@@ -0,0 +1,98 @@
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""UFFD related utility functions"""
+
+import os
+import stat
+import subprocess
+import time
+from pathlib import Path
+
+from framework.utils import chroot
+from host_tools import cargo_build
+
+SOCKET_PATH = "/firecracker-uffd.sock"
+
+
+class UffdHandler:
+    """Describe the UFFD page fault handler process."""
+
+    def __init__(self, name, socket_path, mem_file, chroot_path, log_file_name):
+        """Instantiate the handler process with arguments."""
+        self._proc = None
+        self._handler_name = name
+        self._socket_path = socket_path
+        self._mem_file = mem_file
+        self._chroot = chroot_path
+        self._log_file = log_file_name
+
+    def spawn(self, uid, gid):
+        """Spawn handler process using arguments provided."""
+
+        with chroot(self._chroot):
+            st = os.stat(self._handler_name)
+            os.chmod(self._handler_name, st.st_mode | stat.S_IEXEC)
+
+            chroot_log_file = Path("/") / self._log_file
+            with open(chroot_log_file, "w", encoding="utf-8") as logfile:
+                args = [f"/{self._handler_name}", self._socket_path, self._mem_file]
+                self._proc = subprocess.Popen(
+                    args, stdout=logfile, stderr=subprocess.STDOUT
+                )
+
+            # Give it time start and fail, if it really has too (bad things happen).
+            time.sleep(1)
+            if not self.is_running():
+                print(chroot_log_file.read_text(encoding="utf-8"))
+                assert False, "Could not start PF handler!"
+
+            # The page fault handler will create the socket path with root rights.
+            # Change rights to the jailer's.
+            os.chown(self._socket_path, uid, gid)
+
+    @property
+    def proc(self):
+        """Return UFFD handler process."""
+        return self._proc
+
+    def is_running(self):
+        """Check if UFFD process is running"""
+        return self.proc is not None and self.proc.poll() is None
+
+    @property
+    def log_file(self):
+        """Return the path to the UFFD handler's log file"""
+        return Path(self._chroot) / Path(self._log_file)
+
+    @property
+    def log_data(self):
+        """Return the log data of the UFFD handler"""
+        if self.log_file is None:
+            return ""
+        return self.log_file.read_text(encoding="utf-8")
+
+    def __del__(self):
+        """Tear down the UFFD handler process."""
+        if self.proc is not None:
+            self.proc.kill()
+
+
+def spawn_pf_handler(vm, handler_path, mem_path):
+    """Spawn page fault handler process."""
+    # Copy snapshot memory file into chroot of microVM.
+    jailed_mem = vm.create_jailed_resource(mem_path)
+    # Copy the valid page fault binary into chroot of microVM.
+    jailed_handler = vm.create_jailed_resource(handler_path)
+    handler_name = os.path.basename(jailed_handler)
+
+    uffd_handler = UffdHandler(
+        handler_name, SOCKET_PATH, jailed_mem, vm.chroot(), "uffd.log"
+    )
+    uffd_handler.spawn(vm.jailer.uid, vm.jailer.gid)
+
+    return uffd_handler
+
+
+def uffd_handler(handler_name):
+    """Retrieves the uffd handler with the given name"""
+    return cargo_build.get_example(f"uffd_{handler_name}_handler")

tests/integration_tests/functional/test_uffd.py

@@ -8,9 +8,8 @@
 import pytest
 import requests
 
-from framework.utils import Timeout, UffdHandler, check_output
-
-SOCKET_PATH = "/firecracker-uffd.sock"
+from framework.utils import Timeout, check_output
+from framework.utils_uffd import SOCKET_PATH, spawn_pf_handler, uffd_handler
 
 
 @pytest.fixture(scope="function", name="snapshot")
@@ -36,22 +35,6 @@ def snapshot_fxt(microvm_factory, guest_kernel_linux_5_10, rootfs):
     yield snapshot
 
 
-def spawn_pf_handler(vm, handler_path, mem_path):
-    """Spawn page fault handler process."""
-    # Copy snapshot memory file into chroot of microVM.
-    jailed_mem = vm.create_jailed_resource(mem_path)
-    # Copy the valid page fault binary into chroot of microVM.
-    jailed_handler = vm.create_jailed_resource(handler_path)
-    handler_name = os.path.basename(jailed_handler)
-
-    uffd_handler = UffdHandler(
-        handler_name, SOCKET_PATH, jailed_mem, vm.chroot(), "uffd.log"
-    )
-    uffd_handler.spawn(vm.jailer.uid, vm.jailer.gid)
-
-    return uffd_handler
-
-
 def test_bad_socket_path(uvm_plain, snapshot):
     """
     Test error scenario when socket path does not exist.
@@ -100,7 +83,7 @@ def test_unbinded_socket(uvm_plain, snapshot):
     vm.mark_killed()
 
 
-def test_valid_handler(uvm_plain, snapshot, uffd_handler_paths):
+def test_valid_handler(uvm_plain, snapshot):
     """
     Test valid uffd handler scenario.
     """
@@ -109,9 +92,7 @@ def test_valid_handler(uvm_plain, snapshot, uffd_handler_paths):
     vm.spawn()
 
     # Spawn page fault handler process.
-    _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_handler"], snapshot.mem
-    )
+    _pf_handler = spawn_pf_handler(vm, uffd_handler("valid"), snapshot.mem)
 
     vm.restore_from_snapshot(snapshot, resume=True, uffd_path=SOCKET_PATH)
 
@@ -128,7 +109,7 @@ def test_valid_handler(uvm_plain, snapshot, uffd_handler_paths):
     vm.ssh.check_output("true")
 
 
-def test_malicious_handler(uvm_plain, snapshot, uffd_handler_paths):
+def test_malicious_handler(uvm_plain, snapshot):
     """
     Test malicious uffd handler scenario.
 
@@ -144,9 +125,7 @@ def test_malicious_handler(uvm_plain, snapshot, uffd_handler_paths):
     vm.spawn()
 
     # Spawn page fault handler process.
-    _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["malicious_handler"], snapshot.mem
-    )
+    _pf_handler = spawn_pf_handler(vm, uffd_handler("malicious"), snapshot.mem)
 
     # We expect Firecracker to freeze while resuming from a snapshot
     # due to the malicious handler's unavailability.

tests/integration_tests/performance/test_huge_pages.py

@@ -10,7 +10,7 @@
 from framework.microvm import HugePagesConfig
 from framework.properties import global_props
 from framework.utils_ftrace import ftrace_events
-from integration_tests.functional.test_uffd import SOCKET_PATH, spawn_pf_handler
+from framework.utils_uffd import SOCKET_PATH, spawn_pf_handler, uffd_handler
 
 
 def check_hugetlbfs_in_use(pid: int, allocation_name: str):
@@ -69,9 +69,7 @@ def test_hugetlbfs_boot(uvm_plain):
     )
 
 
-def test_hugetlbfs_snapshot(
-    microvm_factory, guest_kernel_linux_5_10, rootfs, uffd_handler_paths
-):
+def test_hugetlbfs_snapshot(microvm_factory, guest_kernel_linux_5_10, rootfs):
     """
     Test hugetlbfs snapshot restore via uffd
     """
@@ -95,16 +93,14 @@ def test_hugetlbfs_snapshot(
     vm.spawn()
 
     # Spawn page fault handler process.
-    _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_handler"], snapshot.mem
-    )
+    _pf_handler = spawn_pf_handler(vm, uffd_handler("valid"), snapshot.mem)
 
     vm.restore_from_snapshot(snapshot, resume=True, uffd_path=SOCKET_PATH)
 
     check_hugetlbfs_in_use(vm.firecracker_pid, "/anon_hugepage")
 
 
-def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain, uffd_handler_paths):
+def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain):
     """
     Test hugetlbfs differential snapshot support.
 
@@ -139,9 +135,7 @@ def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain, uffd_handler_paths)
     vm.spawn()
 
     # Spawn page fault handler process.
-    _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_handler"], snapshot_merged.mem
-    )
+    _pf_handler = spawn_pf_handler(vm, uffd_handler("valid"), snapshot_merged.mem)
 
     vm.restore_from_snapshot(snapshot_merged, resume=True, uffd_path=SOCKET_PATH)
 
@@ -153,7 +147,6 @@ def test_ept_violation_count(
     microvm_factory,
     guest_kernel_linux_5_10,
     rootfs,
-    uffd_handler_paths,
     metrics,
     huge_pages,
 ):
@@ -200,9 +193,7 @@ def test_ept_violation_count(
     vm.spawn()
 
     # Spawn page fault handler process.
-    _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["fault_all_handler"], snapshot.mem
-    )
+    _pf_handler = spawn_pf_handler(vm, uffd_handler("fault_all"), snapshot.mem)
 
     with ftrace_events("kvm:*"):
         vm.restore_from_snapshot(snapshot, resume=True, uffd_path=SOCKET_PATH)