fix(kms): clear measurement cache entries by their .json file name

dstack/kms/src/main_service.rs

@@ -3,7 +3,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::{
-    path::Path,
     sync::{
         atomic::{AtomicU64, Ordering},
         Arc,
@@ -120,42 +119,12 @@ fn short_lived_cache<V: Clone + Send + Sync + 'static>() -> Cache<(), V> {
         .build()
 }
 
-/// remove a single cache entry (a hex-named subdir/file) under `parent_dir`, or
-/// everything when `sub_dir == "all"`. A non-hex key is rejected to keep the
-/// deletion confined to the cache.
-fn remove_cache(parent_dir: &Path, sub_dir: &str) -> Result<()> {
-    if sub_dir.is_empty() {
-        return Ok(());
-    }
-    if sub_dir == "all" {
-        if parent_dir.exists() {
-            fs::remove_dir_all(parent_dir)?;
-        }
-        return Ok(());
-    }
-    if !sub_dir.chars().all(|c| c.is_ascii_hexdigit()) {
-        bail!("invalid cache key");
-    }
-    let path = parent_dir.join(sub_dir);
-    if path.is_dir() {
-        fs::remove_dir_all(path)?;
-    } else if path.is_file() {
-        fs::remove_file(path)?;
-    }
-    Ok(())
-}
-
 impl KmsState {
     /// clear cached image and measurement material for the given hashes. Used by
     /// the admin `ClearImageCache` RPC; authorization is enforced by the admin
     /// listener's HTTP authenticator, not here.
     pub(crate) fn clear_image_cache(&self, image_hash: &str, config_hash: &str) -> Result<()> {
-        let images_dir = self.config.image.cache_dir.join("images");
-        remove_cache(&images_dir, image_hash).context("failed to clear image cache")?;
-        // measurement cache is kept by the verifier under measurements/.
-        let mr_cache_dir = self.config.image.cache_dir.join("measurements");
-        remove_cache(&mr_cache_dir, config_hash).context("failed to clear measurement cache")?;
-        Ok(())
+        self.verifier.clear_cache(image_hash, config_hash)
     }
 
     pub fn new(config: KmsConfig) -> Result<Self> {
@@ -649,7 +618,7 @@ mod tests {
     };
     use cc_eventlog::RuntimeEvent;
     use sha2::{Digest, Sha256, Sha384};
-    use std::time::Duration;
+    use std::{path::Path, time::Duration};
 
     const SNP_ENABLED: KeyReleasePolicy = KeyReleasePolicy {
         snp: true,
@@ -662,28 +631,6 @@ mod tests {
         nitro_enclave: true,
     };
 
-    #[test]
-    fn remove_cache_only_deletes_the_named_hex_entry() {
-        let dir = tempfile::tempdir().unwrap();
-        let root = dir.path();
-        let keep = root.join("cafe");
-        let drop = root.join("beef");
-        fs::create_dir_all(&keep).unwrap();
-        fs::create_dir_all(&drop).unwrap();
-
-        remove_cache(root, "beef").unwrap();
-        assert!(!drop.exists(), "named entry must be removed");
-        assert!(keep.exists(), "other entries must be kept");
-
-        // empty key is a no-op; a non-hex key is rejected before touching disk.
-        remove_cache(root, "").unwrap();
-        assert!(remove_cache(root, "../escape").is_err());
-        assert!(keep.exists());
-
-        // "all" clears the whole cache dir.
-        remove_cache(root, "all").unwrap();
-        assert!(!root.exists());
-    }
     use std::collections::BTreeMap;
 
     fn hex_of(byte: u8, len: usize) -> String {

dstack/verifier/src/verification.rs

@@ -132,6 +132,29 @@ pub struct CvmVerifier {
     download_locks: Mutex<HashMap<String, Weak<AsyncMutex<()>>>>,
 }
 
+/// Remove the entry for `key` (rejecting non-hex keys so deletion stays inside
+/// the cache), everything under `dir` for `"all"`, or nothing for an empty key.
+fn remove_cache(dir: &Path, key: &str, entry: impl FnOnce(&str) -> PathBuf) -> Result<()> {
+    match key {
+        "" => {}
+        "all" => {
+            if dir.exists() {
+                fs_err::remove_dir_all(dir)?;
+            }
+        }
+        _ if key.chars().all(|c| c.is_ascii_hexdigit()) => {
+            let path = entry(key);
+            if path.is_dir() {
+                fs_err::remove_dir_all(path)?;
+            } else if path.is_file() {
+                fs_err::remove_file(path)?;
+            }
+        }
+        _ => bail!("invalid cache key"),
+    }
+    Ok(())
+}
+
 impl CvmVerifier {
     pub fn new(
         image_cache_dir: String,
@@ -363,6 +386,19 @@ impl CvmVerifier {
         Path::new(&self.image_cache_dir).join("images")
     }
 
+    /// Clear the cached image and measurements. Each key is a hex cache key,
+    /// `"all"` to clear the whole cache, or empty to leave it untouched.
+    pub fn clear_cache(&self, image_hash: &str, config_hash: &str) -> Result<()> {
+        remove_cache(&self.images_dir(), image_hash, |key| {
+            self.images_dir().join(key)
+        })
+        .context("failed to clear image cache")?;
+        remove_cache(&self.measurement_cache_dir(), config_hash, |key| {
+            self.measurement_cache_path(key)
+        })
+        .context("failed to clear measurement cache")
+    }
+
     fn image_content_digest(image_dir: &Path) -> Result<Option<Vec<u8>>> {
         let sha256sum_path = image_dir.join("sha256sum.txt");
         if !sha256sum_path.exists() {
@@ -1699,6 +1735,34 @@ mod tests {
         }
     }
 
+    #[test]
+    fn clear_cache_only_deletes_the_named_entries() {
+        let directory = tempfile::tempdir().unwrap();
+        let mut verifier = test_verifier();
+        verifier.image_cache_dir = directory.path().display().to_string();
+        for key in ["cafe", "beef"] {
+            fs_err::create_dir_all(verifier.images_dir().join(key)).unwrap();
+            let path = verifier.measurement_cache_path(key);
+            fs_err::create_dir_all(path.parent().unwrap()).unwrap();
+            fs_err::write(path, "{}").unwrap();
+        }
+
+        verifier.clear_cache("beef", "beef").unwrap();
+        assert!(!verifier.images_dir().join("beef").exists());
+        assert!(!verifier.measurement_cache_path("beef").exists());
+        assert!(verifier.images_dir().join("cafe").exists());
+        assert!(verifier.measurement_cache_path("cafe").exists());
+
+        // empty keys are a no-op; a non-hex key is rejected before touching disk.
+        verifier.clear_cache("", "").unwrap();
+        assert!(verifier.clear_cache("../escape", "").is_err());
+        assert!(verifier.clear_cache("", "../escape").is_err());
+
+        verifier.clear_cache("all", "all").unwrap();
+        assert!(!verifier.images_dir().exists());
+        assert!(!verifier.measurement_cache_dir().exists());
+    }
+
     #[test]
     fn measurement_cache_version_mismatch_is_ignored_and_replaced() {
         let directory = tempfile::tempdir().unwrap();