diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index a56919f99..baac0147d 100644 --- a/dstack/kms/src/main_service.rs +++ b/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() -> 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 { @@ -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 { diff --git a/dstack/verifier/src/verification.rs b/dstack/verifier/src/verification.rs index d11bc4090..949f36d0b 100644 --- a/dstack/verifier/src/verification.rs +++ b/dstack/verifier/src/verification.rs @@ -132,6 +132,29 @@ pub struct CvmVerifier { download_locks: Mutex>>>, } +/// 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>> { 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();