Add test scripts to test UFFD handler + balloon interactions

get_rootfs_guest_kernel.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+set -e
+
+ARCH="$(uname -m)"
+
+latest=$(wget "http://spec.ccfc.min.s3.amazonaws.com/?prefix=firecracker-ci/v1.11/$ARCH/vmlinux-5.10&list-type=2" -O - 2>/dev/null | grep -oP "(?<=<Key>)(firecracker-ci/v1.11/$ARCH/vmlinux-5\.10\.[0-9]{1,3})(?=</Key>)")
+
+# Download a linux kernel binary
+wget "https://s3.amazonaws.com/spec.ccfc.min/${latest}"
+
+# Download a rootfs
+wget -O ubuntu-24.04.squashfs.upstream "https://s3.amazonaws.com/spec.ccfc.min/firecracker-ci/v1.11/${ARCH}/ubuntu-24.04.squashfs"
+
+# Create an ssh key for the rootfs
+unsquashfs ubuntu-24.04.squashfs.upstream
+ssh-keygen -f id_rsa -N ""
+cp -v id_rsa.pub squashfs-root/root/.ssh/authorized_keys
+mv -v id_rsa ./ubuntu-24.04.id_rsa
+# create ext4 filesystem image
+sudo chown -R root:root squashfs-root
+truncate -s 400M ubuntu-24.04.ext4
+sudo mkfs.ext4 -d squashfs-root -F ubuntu-24.04.ext4

run_firecracker.sh

@@ -0,0 +1,6 @@
+# Run in a separate terminal and leave open
+
+API_SOCKET="/tmp/firecracker.socket"
+sudo rm -f $API_SOCKET
+sudo firecracker --api-sock "${API_SOCKET}"
+

run_uffd_handler.sh

@@ -0,0 +1,2 @@
+sudo rm -rf /tmp/uffd.socket
+cargo run --example uffd_valid_handler  /tmp/uffd.socket /tmp/mem_file

snapshot_vm.sh

@@ -0,0 +1,129 @@
+#!/bin/bash
+set -e
+
+TAP_DEV="tap0"
+TAP_IP="172.16.0.1"
+MASK_SHORT="/30"
+
+# Setup network interface
+sudo ip link del "$TAP_DEV" 2> /dev/null || true
+sudo ip tuntap add dev "$TAP_DEV" mode tap
+sudo ip addr add "${TAP_IP}${MASK_SHORT}" dev "$TAP_DEV"
+sudo ip link set dev "$TAP_DEV" up
+
+# Enable ip forwarding
+sudo sh -c "echo 1 > /proc/sys/net/ipv4/ip_forward"
+sudo iptables -P FORWARD ACCEPT
+
+# This tries to determine the name of the host network interface to forward
+# VM's outbound network traffic through. If outbound traffic doesn't work,
+# double check this returns the correct interface!
+HOST_IFACE=$(ip -j route list default |jq -r '.[0].dev')
+
+# Set up microVM internet access
+sudo iptables -t nat -D POSTROUTING -o "$HOST_IFACE" -j MASQUERADE || true
+sudo iptables -t nat -A POSTROUTING -o "$HOST_IFACE" -j MASQUERADE
+
+API_SOCKET="/tmp/firecracker.socket"
+LOGFILE="./firecracker.log"
+
+# Create log file
+touch $LOGFILE
+
+# Set log file
+sudo curl -X PUT --unix-socket "${API_SOCKET}" \
+    --data "{
+        \"log_path\": \"${LOGFILE}\",
+        \"level\": \"Debug\",
+        \"show_level\": true,
+        \"show_log_origin\": true
+    }" \
+    "http://localhost/logger"
+
+KERNEL="./$(ls vmlinux* | tail -1)"
+KERNEL_BOOT_ARGS="console=ttyS0 reboot=k panic=1 pci=off"
+
+ARCH=$(uname -m)
+
+if [ ${ARCH} = "aarch64" ]; then
+    KERNEL_BOOT_ARGS="keep_bootcon ${KERNEL_BOOT_ARGS}"
+fi
+
+# Set boot source
+sudo curl -X PUT --unix-socket "${API_SOCKET}" \
+    --data "{
+        \"kernel_image_path\": \"${KERNEL}\",
+        \"boot_args\": \"${KERNEL_BOOT_ARGS}\"
+    }" \
+    "http://localhost/boot-source"
+
+ROOTFS="./ubuntu-24.04.ext4"
+
+# Set rootfs
+sudo curl -X PUT --unix-socket "${API_SOCKET}" \
+    --data "{
+        \"drive_id\": \"rootfs\",
+        \"path_on_host\": \"${ROOTFS}\",
+        \"is_root_device\": true,
+        \"is_read_only\": false
+    }" \
+    "http://localhost/drives/rootfs"
+
+# The IP address of a guest is derived from its MAC address with
+# `fcnet-setup.sh`, this has been pre-configured in the guest rootfs. It is
+# important that `TAP_IP` and `FC_MAC` match this.
+FC_MAC="06:00:AC:10:00:02"
+
+# Set network interface
+sudo curl -X PUT --unix-socket "${API_SOCKET}" \
+    --data "{
+        \"iface_id\": \"net1\",
+        \"guest_mac\": \"$FC_MAC\",
+        \"host_dev_name\": \"$TAP_DEV\"
+    }" \
+    "http://localhost/network-interfaces/net1"
+
+# Initialize balloon
+amount_mib=5
+deflate_on_oom=true
+polling_interval=1
+sudo curl --unix-socket "${API_SOCKET}" \
+    -X PUT 'http://localhost/balloon' \
+    -H 'Accept: application/json' \
+    -H 'Content-Type: application/json' \
+    -d "{
+        \"amount_mib\": $amount_mib, \
+        \"deflate_on_oom\": $deflate_on_oom, \
+        \"stats_polling_interval_s\": $polling_interval \
+    }"
+
+sleep 0.015s
+
+# Start the machine
+sudo curl -X PUT --unix-socket "${API_SOCKET}" \
+    --data "{
+        \"action_type\": \"InstanceStart\"
+    }" \
+    "http://localhost/actions"
+
+# Wait for machine to boot up properly
+sleep 10s
+
+# Snapshot the machine
+sudo curl -X PATCH --unix-socket "${API_SOCKET}" \
+ -d '{
+            "state": "Paused"
+    }' http://localhost/vm
+
+
+
+sudo curl --unix-socket "${API_SOCKET}" -i \
+ 	-X PUT 'http://localhost/snapshot/create' \
+	-H  'Accept: application/json' \
+	-H  'Content-Type: application/json' \
+	-d '{
+            "snapshot_type": "Full",
+            "snapshot_path": "/tmp/snapshot_file",
+            "mem_file_path": "/tmp/mem_file" 
+    }' 
+

src/firecracker/examples/uffd/uffd_utils.rs

@@ -121,6 +121,9 @@ impl UffdHandler {
                 //    memory from the host (through balloon device)
                 Some(MemPageState::Uninitialized) | Some(MemPageState::FromFile) => {
                     let (start, end) = self.populate_from_file(region, fault_page_addr, len);
+                    if start == 0 && end == 0 {
+                        return;
+                    }
                     self.update_mem_state_mappings(start, end, MemPageState::FromFile);
                     return;
                 }
@@ -144,9 +147,17 @@ impl UffdHandler {
         let src = self.backing_buffer as u64 + region.mapping.offset + offset;
 
         let ret = unsafe {
-            self.uffd
-                .copy(src as *const _, dst as *mut _, len, true)
-                .expect("Uffd copy failed")
+            match self.uffd.copy(src as *const _, dst as *mut _, len, true) {
+                Ok(value) => value,
+                // Catch EAGAIN errors, which occur when there is a simultaneous copy with a EVENT_REMOVE.
+                // Ignore in this case.
+                Err(Error::PartiallyCopied(bytes_copied)) => {
+                    return (0, 0);
+                }
+                Err(e) => {
+                    panic!("Uffd copy failed");
+                }
+            }
         };
 
         // Make sure the UFFD copied some bytes.

src/firecracker/examples/uffd/valid_handler.rs

@@ -37,11 +37,14 @@ fn main() {
             userfaultfd::Event::Pagefault { addr, .. } => {
                 uffd_handler.serve_pf(addr.cast(), uffd_handler.page_size)
             }
-            userfaultfd::Event::Remove { start, end } => uffd_handler.update_mem_state_mappings(
-                start as u64,
-                end as u64,
-                MemPageState::Removed,
-            ),
+            userfaultfd::Event::Remove { start, end } => {
+                println!("Received remove event");
+                uffd_handler.update_mem_state_mappings(
+                    start as u64,
+                    end as u64,
+                    MemPageState::Removed,
+                );
+            }
             _ => panic!("Unexpected event on userfaultfd"),
         }
     });

trigger_remove_events.sh

@@ -0,0 +1,21 @@
+API_SOCKET="/tmp/firecracker.socket"
+
+sudo curl --unix-socket "${API_SOCKET}" -i     -X PUT 'http://localhost/snapshot/load'     -H  'Accept: application/json'     -H  'Content-Type: application/json'     -d '{
+            "snapshot_path": "/tmp/snapshot_file",
+            "mem_backend": {
+                "backend_path": "/tmp/uffd.socket",
+                "backend_type": "Uffd"
+            },
+            "enable_diff_snapshots": true,
+            "resume_vm": true
+    }'
+
+# Inflate balloon to trigger remove events
+sudo curl --unix-socket "${API_SOCKET}" -i  -X PATCH 'http://localhost/balloon' \
+    -H 'Accept: application/json' \
+    -H 'Content-Type: application/json' \
+    -d "{
+        \"amount_mib\": 20 \
+    }"
+
+