cpiovfs snapshot support: initrd re-map on restore

Signed-off-by: danbugs <danilochiarlone@gmail.com>

Author: danbugs

examples/python-agent-driver/kraft.yaml

@@ -11,17 +11,20 @@ unikraft:
     CONFIG_LIBUKINTCTLR_HYPERLIGHT: 'y'
     CONFIG_HYPERLIGHT_MAX_GUEST_LOG_LEVEL: 4
 
-    CONFIG_LIBUKPRINT_KLVL_CRIT: 'y'
+    CONFIG_LIBUKPRINT_KLVL_INFO: 'y'
     CONFIG_LIBUKPRINT_PRINT_TIME: 'n'
-    CONFIG_LIBUKPRINT_PRINT_SRCNAME: 'n'
+    CONFIG_LIBUKPRINT_PRINT_SRCNAME: 'y'
     CONFIG_LIBUKBOOT_BANNER_NONE: 'y'
     CONFIG_OPTIMIZE_SIZE: 'y'
     CONFIG_OPTIMIZE_PIE: 'y'
 
     CONFIG_LIBVFSCORE: 'y'
     CONFIG_LIBVFSCORE_AUTOMOUNT_CI: 'y'
-    CONFIG_LIBVFSCORE_AUTOMOUNT_CI_INITRD: 'y'
-    CONFIG_LIBRAMFS: 'y'
+    CONFIG_LIBVFSCORE_AUTOMOUNT_CI_CUSTOM: 'y'
+    CONFIG_LIBVFSCORE_AUTOMOUNT_CI0_MP: '/'
+    CONFIG_LIBVFSCORE_AUTOMOUNT_CI0_DRIVER: 'cpiovfs'
+    CONFIG_LIBVFSCORE_AUTOMOUNT_CI0_UKOPTS_IFINITRD0: 'y'
+    CONFIG_LIBCPIOVFS: 'y'
     CONFIG_LIBUKCPIO: 'y'
 
     CONFIG_LIBDEVFS: 'y'

host/examples/test_cpiovfs.rs

@@ -0,0 +1,41 @@
+use hyperlight_unikraft::Sandbox;
+
+fn main() -> anyhow::Result<()> {
+    let kernel = std::env::args()
+        .nth(1)
+        .expect("usage: test_cpiovfs <kernel> <initrd>");
+    let initrd = std::env::args()
+        .nth(2)
+        .expect("usage: test_cpiovfs <kernel> <initrd>");
+
+    eprintln!("=== Test: build + init + snapshot + restore + call ===");
+    {
+        let mut sbox = Sandbox::builder(&kernel)
+            .initrd_file(&initrd)
+            .heap_size(3 * 512 * 1024 * 1024)
+            .build()?;
+        eprintln!("  build OK");
+        sbox.restore()?;
+        let _: () = sbox.call_named("init", ())?;
+        eprintln!("  init OK");
+        sbox.snapshot_now()?;
+        eprintln!("  snapshot OK");
+
+        let snap_path = "/tmp/cpiovfs_snapshot.hls";
+        sbox.save_snapshot(snap_path)?;
+        let snap_size = std::fs::metadata(snap_path)?.len();
+        eprintln!(
+            "  snapshot size: {} MiB ({} bytes)",
+            snap_size / 1024 / 1024,
+            snap_size
+        );
+
+        sbox.restore()?;
+        eprintln!("  restore OK");
+        let _: () = sbox.call_named("run", "print('test ok')".to_string())?;
+        eprintln!("  call OK");
+    }
+
+    eprintln!("=== All tests passed ===");
+    Ok(())
+}

host/src/bin/pyhl.rs

@@ -464,8 +464,12 @@ fn cmd_run(args: RunArgs) -> Result<()> {
         .map(|m| parse_mount(m))
         .collect::<Result<_>>()?;
 
+    let initrd = home.join(INITRD_FILE);
+
     let t_load = Instant::now();
-    let mut sandbox = if run_preopens.is_empty() {
+    let mut sandbox = if initrd.is_file() {
+        Sandbox::from_snapshot_file_with_initrd(&snapshot, &run_preopens, &initrd)?
+    } else if run_preopens.is_empty() {
         Sandbox::from_snapshot_file(&snapshot)?
     } else {
         Sandbox::from_snapshot_file_with(&snapshot, &run_preopens)?

host/src/lib.rs

@@ -1387,7 +1387,7 @@ impl Sandbox {
     /// a 2.5 GB snapshot — enough to double the whole `pyhl run` wall
     /// time on simple scripts.
     pub fn from_snapshot_file<P: AsRef<Path>>(path: P) -> Result<Self> {
-        Self::from_snapshot_file_with(path, &[])
+        Self::from_snapshot_file_full(path, &[], None)
     }
 
     /// Load a previously-persisted snapshot and register a
@@ -1402,6 +1402,26 @@ impl Sandbox {
     /// fixed at setup time because it lives in the snapshot's memory
     /// image.
     pub fn from_snapshot_file_with<P: AsRef<Path>>(path: P, preopens: &[Preopen]) -> Result<Self> {
+        Self::from_snapshot_file_full(path, preopens, None)
+    }
+
+    /// Load a snapshot with an initrd file re-mapped at the standard
+    /// guest VA (0xC000_0000). Required when the snapshot was taken
+    /// from a cpiovfs-backed guest whose VFS nodes point into the
+    /// initrd region.
+    pub fn from_snapshot_file_with_initrd<P: AsRef<Path>, I: AsRef<Path>>(
+        path: P,
+        preopens: &[Preopen],
+        initrd: I,
+    ) -> Result<Self> {
+        Self::from_snapshot_file_full(path, preopens, Some(initrd.as_ref().to_path_buf()))
+    }
+
+    fn from_snapshot_file_full<P: AsRef<Path>>(
+        path: P,
+        preopens: &[Preopen],
+        initrd: Option<std::path::PathBuf>,
+    ) -> Result<Self> {
         let loaded = Snapshot::from_file_unchecked(path.as_ref())?;
         let arc = Arc::new(loaded);
 
@@ -1421,13 +1441,18 @@ impl Sandbox {
             })?;
         }
 
-        let inner = MultiUseSandbox::from_snapshot(arc.clone(), host_funcs, None)?;
+        let mut inner = MultiUseSandbox::from_snapshot(arc.clone(), host_funcs, None)?;
+
+        const INITRD_MAP_BASE: u64 = 0xC000_0000;
+        if let Some(ref initrd_path) = initrd {
+            inner.map_file_cow(initrd_path, INITRD_MAP_BASE, Some("initrd"))?;
+        }
 
         Ok(Self {
             inner,
             snapshot: Some(arc),
-            file_mapping_path: None,
-            file_mapping_base: 0,
+            file_mapping_path: initrd,
+            file_mapping_base: INITRD_MAP_BASE,
         })
     }
 }

host/src/pyhl.rs

@@ -222,7 +222,10 @@ impl Runtime {
                 snap.display()
             );
         }
-        let sandbox = if mounts.is_empty() {
+        let initrd = home.join(INITRD_FILE);
+        let sandbox = if initrd.is_file() {
+            Sandbox::from_snapshot_file_with_initrd(&snap, mounts, &initrd)?
+        } else if mounts.is_empty() {
             Sandbox::from_snapshot_file(&snap)?
         } else {
             Sandbox::from_snapshot_file_with(&snap, mounts)?