feat: enable write-protection on guest memory

UFFD provides an API to enable write-protection for memory ranges
tracked by a userfault file descriptor. Detailed information can be
found here: https://docs.kernel.org/admin-guide/mm/userfaultfd.html.

To use the feature, users need to register the memory region with
UFFDIO_REGISTER_MODE_WP. Then, users need to enable explicitly
write-protection for sub-ranges of the registered region.

Writes in pages within write-protected memory ranges can be handled in
one of two ways. In synchronous mode, writes in a protected page will
cause kernel to send a write protection event over the userfaultfd.
In asynchronous mode, the kernel will automatically handle writes to
protected pages by clearing the write-protection bit. Userspace can
later observe the write protection bit by looking into the corresponding
entry of /proc/<pid>/pagemap.

This commit, uncoditionally, enables write protection for guest memory
using the asynchronous mode.

!NOTE!: asynchronous write protection requires (host) kernel version 6.7
or later).

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

Author: bchalios

Cargo.lock

@@ -42,7 +42,11 @@ serde_json = "1.0.132"
 slab = "0.4.7"
 thiserror = "1.0.67"
 timerfd = "1.5.0"
-userfaultfd = "0.8.1"
+userfaultfd = { git = "https://github.com/e2b-dev/userfaultfd-rs", branch = "feat_write_protection", features = [
+  "linux5_7",
+  "linux5_13",
+  "linux6_7"
+] }
 utils = { path = "../utils" }
 vhost = { version = "0.13.0", features = ["vhost-user-frontend"] }
 vm-allocator = "0.1.0"

src/vmm/Cargo.toml

@@ -513,6 +513,10 @@ impl Vmm {
                 )
             };
 
+            // TODO: if we don't support UFFD/async WP, we can completely skip this bit, as the
+            // UFFD handler already tracks dirty pages through the WriteProtected events. For the
+            // time being, we always do.
+            //
             // Build dirty bitmap: check pagemap only for pages that mincore reports resident.
             let mut slot_bitmap = vec![0u64; nr_pages.div_ceil(64)];
             for page_idx in 0..nr_pages {

src/vmm/src/lib.rs

@@ -14,7 +14,7 @@ use std::sync::{Arc, Mutex};
 use seccompiler::BpfThreadMap;
 use semver::Version;
 use serde::{Deserialize, Serialize};
-use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
+use userfaultfd::{FeatureFlags, RegisterMode, Uffd, UffdBuilder};
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
 #[cfg(target_arch = "aarch64")]
@@ -536,6 +536,8 @@ pub enum GuestMemoryFromUffdError {
     Create(userfaultfd::Error),
     /// Failed to register memory address range with the userfaultfd object: {0}
     Register(userfaultfd::Error),
+    /// Failed to enable write protection on memory address range with the userfaultfd object: {0}
+    WriteProtect(userfaultfd::Error),
     /// Failed to connect to UDS Unix stream: {0}
     Connect(#[from] std::io::Error),
     /// Failed to sends file descriptor: {0}
@@ -560,6 +562,10 @@ fn guest_memory_from_uffd(
         uffd_builder.require_features(FeatureFlags::EVENT_REMOVE);
     }
 
+    uffd_builder.require_features(
+        FeatureFlags::MISSING_HUGETLBFS | FeatureFlags::WP_ASYNC,
+    );
+
     let uffd = uffd_builder
         .close_on_exec(true)
         .non_blocking(true)
@@ -568,8 +574,22 @@ fn guest_memory_from_uffd(
         .map_err(GuestMemoryFromUffdError::Create)?;
 
     for mem_region in guest_memory.iter() {
-        uffd.register(mem_region.as_ptr().cast(), mem_region.size() as _)
-            .map_err(GuestMemoryFromUffdError::Register)?;
+        uffd.register_with_mode(
+            mem_region.as_ptr().cast(),
+            mem_region.size() as _,
+            RegisterMode::MISSING | RegisterMode::WRITE_PROTECT,
+        )
+        .map_err(GuestMemoryFromUffdError::Register)?;
+
+        // If memory is backed by huge pages, we can immediately write protect it.
+        // Otherwise (memory is backed by anonymous memory), write protecting here
+        // won't have any effect, as the write-protection bit for a page will be
+        // wiped when the first page fault occurs. These cases need to be handled
+        // directly from the UFFD handler.
+        if huge_pages.is_hugetlbfs() {
+            uffd.write_protect(mem_region.as_ptr().cast(), mem_region.size() as _)
+                .map_err(GuestMemoryFromUffdError::WriteProtect)?;
+        }
     }
 
     send_uffd_handshake(mem_uds_path, &backend_mappings, &uffd)?;