fix(snapshot): prevent lost OCI index updates during concurrent saves

src/hyperlight_host/src/sandbox/snapshot/file/mod.rs

@@ -11,6 +11,7 @@ mod media_types;
 pub(crate) mod reference;
 mod transport;
 
+use std::fs::OpenOptions;
 use std::path::{Path, PathBuf};
 
 use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
@@ -146,6 +147,62 @@ fn read_layout_marker(path: &Path) -> crate::Result<()> {
     Ok(())
 }
 
+/// Validate an existing OCI layout marker, if present.
+/// Returns whether the marker already exists.
+fn validate_existing_layout(path: &Path) -> crate::Result<bool> {
+    let marker = path.join("oci-layout");
+    let exists = marker
+        .try_exists()
+        .map_err(|e| crate::new_error!("save: failed to stat {:?}: {}", marker, e))?;
+
+    if exists {
+        read_layout_marker(path)?;
+    }
+
+    Ok(exists)
+}
+
+/// Validate and read an existing OCI index without modifying the layout.
+/// Called before writing blobs and again under the index lock.
+fn read_existing_index_manifests(index_path: &Path) -> crate::Result<Vec<Descriptor>> {
+    if !index_path
+        .try_exists()
+        .map_err(|e| crate::new_error!("save: failed to stat {:?}: {}", index_path, e))?
+    {
+        return Ok(Vec::new());
+    }
+
+    let bytes = read_bounded(index_path, MAX_JSON_BLOB_SIZE)
+        .map_err(|e| crate::new_error!("save: failed to read existing index.json: {}", e))?;
+
+    let existing: ImageIndex = serde_json::from_slice(&bytes).map_err(|e| {
+        crate::new_error!(
+            "save: existing index.json is not a valid OCI image index: {}",
+            e
+        )
+    })?;
+
+    if existing.schema_version() != SCHEMA_VERSION {
+        return Err(crate::new_error!(
+            "save: unsupported OCI index schemaVersion {} (expected {})",
+            existing.schema_version(),
+            SCHEMA_VERSION
+        ));
+    }
+
+    if let Some(media_type) = existing.media_type()
+        && !matches!(media_type, MediaType::ImageIndex)
+    {
+        return Err(crate::new_error!(
+            "save: OCI index has unexpected media type {} (expected {})",
+            media_type,
+            MediaType::ImageIndex
+        ));
+    }
+
+    Ok(existing.manifests().to_vec())
+}
+
 fn load_manifest(
     path: &Path,
     blobs_dir: &Path,
@@ -397,60 +454,38 @@ impl Snapshot {
             }
         }
 
-        // Validate any pre-existing `oci-layout` marker before
-        // touching anything else, so a foreign layout (future
-        // version, hand-edited file) is reported without altering
-        // the directory.
-        let layout_marker = path.join("oci-layout");
-        let marker_existed = layout_marker
-            .try_exists()
-            .map_err(|e| crate::new_error!("save: failed to stat {:?}: {}", layout_marker, e))?;
-        if marker_existed {
-            let bytes = read_bounded(&layout_marker, MAX_JSON_BLOB_SIZE).map_err(|e| {
-                crate::new_error!("save: failed to read existing oci-layout: {}", e)
-            })?;
-            let v: serde_json::Value = serde_json::from_slice(&bytes).map_err(|e| {
-                crate::new_error!("save: existing oci-layout is not valid JSON: {}", e)
-            })?;
-            match v.get("imageLayoutVersion").and_then(|s| s.as_str()) {
-                Some(s) if s == OCI_LAYOUT_VERSION => {}
-                Some(other) => {
-                    return Err(crate::new_error!(
-                        "save: existing imageLayoutVersion {:?} is unsupported (expected {:?})",
-                        other,
-                        OCI_LAYOUT_VERSION
-                    ));
-                }
-                None => {
-                    return Err(crate::new_error!(
-                        "save: existing oci-layout is missing imageLayoutVersion"
-                    ));
-                }
-            }
-        }
-
+        // Validate layout and index before writing snapshot blobs.
+        validate_existing_layout(path)?;
         let index_path = path.join("index.json");
-        let index_existed = index_path
-            .try_exists()
-            .map_err(|e| crate::new_error!("save: failed to stat {:?}: {}", index_path, e))?;
-        let mut manifests: Vec<Descriptor> = if index_existed {
-            let bytes = read_bounded(&index_path, MAX_JSON_BLOB_SIZE).map_err(|e| {
-                crate::new_error!("save: failed to read existing index.json: {}", e)
-            })?;
-            let existing: ImageIndex = serde_json::from_slice(&bytes).map_err(|e| {
-                crate::new_error!(
-                    "save: existing index.json is not a valid OCI image index: {}",
-                    e
-                )
-            })?;
-            existing.manifests().to_vec()
-        } else {
-            Vec::new()
-        };
+        read_existing_index_manifests(&index_path)?;
+        let layout_marker = path.join("oci-layout");
 
+        // Generate content-addressed snapshot blobs outside the index lock.
         let new_desc = self.write_blobs_and_build_descriptor(path, tag, &cfg, &cfg_bytes)?;
         let written_digest = OciDigest::from_oci_spec_digest(new_desc.digest());
 
+        // Serialize layout validation and the index read-modify-write
+        // transaction across cooperating processes.
+        // Use a stable lock file, not index.json, which is atomically replaced.
+        let lock_path = path.join(".hyperlight-index.lock");
+        let lock_file = OpenOptions::new()
+            .create(true)
+            .read(true)
+            .write(true)
+            .open(&lock_path)
+            .map_err(|e| {
+                crate::new_error!("save: failed to open index lock {:?}: {}", lock_path, e)
+            })?;
+        lock_file.lock().map_err(|e| {
+            crate::new_error!("save: failed to acquire index lock {:?}: {}", lock_path, e)
+        })?;
+        // File lock is released when lock_file is dropped on all exit paths.
+
+        // Revalidate under the lock in case the layout changed.
+        let marker_existed = validate_existing_layout(path)?;
+
+        let mut manifests = read_existing_index_manifests(&index_path)?;
+
         // Replacement is by tag, not by digest: a new snapshot may
         // hash to a different value but still claim the same logical
         // ref. Blobs from the replaced manifest become orphans.

src/hyperlight_host/src/sandbox/snapshot/file_tests.rs

@@ -189,6 +189,97 @@ fn snapshot_metadata_namespaces_are_independent() {
 
 // Round-trip via OCI layout on disk.
 
+#[test]
+fn save_rejects_malformed_existing_index_without_writing_blobs() {
+    let snap = create_snapshot();
+    let dir = tempfile::tempdir().unwrap();
+    let path = dir.path().join("malformed-index");
+
+    std::fs::create_dir(&path).unwrap();
+
+    std::fs::write(
+        path.join("oci-layout"),
+        br#"{"imageLayoutVersion":"1.0.0"}"#,
+    )
+    .unwrap();
+
+    std::fs::write(path.join("index.json"), b"not-json").unwrap();
+
+    let err = snap
+        .save(&path, &OciTag::new("latest").unwrap())
+        .unwrap_err();
+
+    assert!(err.to_string().contains("index.json"));
+
+    assert_eq!(std::fs::read(path.join("index.json")).unwrap(), b"not-json");
+
+    assert!(
+        !path.join("blobs").exists(),
+        "save wrote snapshot blobs before validating the index"
+    );
+
+    assert!(
+        !path.join(".hyperlight-index.lock").exists(),
+        "save created a lock file before validating the index"
+    );
+}
+
+#[test]
+fn concurrent_snapshot_saves_preserve_all_tags() {
+    use std::collections::HashSet;
+    use std::sync::Barrier;
+    use std::thread;
+
+    const WRITERS: usize = 8;
+
+    let snapshot = create_snapshot();
+    let dir = tempfile::tempdir().unwrap();
+    let path = dir.path().join("concurrent");
+    let barrier = Arc::new(Barrier::new(WRITERS));
+
+    let handles: Vec<_> = (0..WRITERS)
+        .map(|i| {
+            let snapshot = Arc::clone(&snapshot);
+            let barrier = Arc::clone(&barrier);
+            let path = path.clone();
+
+            thread::spawn(move || {
+                let tag = format!("writer-{i}");
+                barrier.wait();
+                snapshot.save(&path, &OciTag::new(&tag).unwrap()).unwrap();
+            })
+        })
+        .collect();
+
+    for handle in handles {
+        handle.join().unwrap();
+    }
+
+    let index: Value =
+        serde_json::from_slice(&std::fs::read(path.join("index.json")).unwrap()).unwrap();
+
+    let tags: HashSet<String> = index["manifests"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|manifest| {
+            manifest["annotations"]["org.opencontainers.image.ref.name"]
+                .as_str()
+                .unwrap()
+                .to_owned()
+        })
+        .collect();
+
+    assert_eq!(tags.len(), WRITERS);
+
+    for i in 0..WRITERS {
+        assert!(
+            tags.contains(&format!("writer-{i}")),
+            "missing snapshot tag writer-{i}"
+        );
+    }
+}
+
 #[test]
 fn round_trip_save_load_call() {
     let snapshot = create_snapshot();
@@ -1541,6 +1632,15 @@ fn save_rejects_unsupported_existing_layout_version() {
         br#"{"imageLayoutVersion":"99.0.0"}"#,
     )
     .unwrap();
+    // Capture every existing entry and its bytes before the failed save.
+    let before: Vec<_> = std::fs::read_dir(&path)
+        .unwrap()
+        .map(|entry| {
+            let entry = entry.unwrap();
+            (entry.file_name(), std::fs::read(entry.path()).unwrap())
+        })
+        .collect();
+
     let err = snap
         .save(&path, &OciTag::new("latest").unwrap())
         .unwrap_err();
@@ -1549,9 +1649,17 @@ fn save_rejects_unsupported_existing_layout_version() {
         msg.contains("imageLayoutVersion") || msg.contains("unsupported"),
         "expected unsupported-version error, got: {msg}"
     );
-    assert!(
-        !path.join("index.json").exists(),
-        "save must not have written index.json"
+    let after: Vec<_> = std::fs::read_dir(&path)
+        .unwrap()
+        .map(|entry| {
+            let entry = entry.unwrap();
+            (entry.file_name(), std::fs::read(entry.path()).unwrap())
+        })
+        .collect();
+
+    assert_eq!(
+        before, after,
+        "unsupported layout must leave directory contents unchanged"
     );
 }