From 25655b5681aa5f1f8d7fc3c286b67666fdee740b Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Fri, 23 Oct 2026 07:19:56 +0800 Subject: [PATCH 1/2] fix(kms): clear measurement cache entries by their .json file name Signed-off-by: Leechael Yim --- dstack/kms/src/main_service.rs | 41 ++++++++++++++++++++--------- dstack/verifier/src/lib.rs | 2 +- dstack/verifier/src/verification.rs | 5 +++- 3 files changed, 33 insertions(+), 15 deletions(-) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 6a5af3941..a42ddcd3c 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -18,7 +18,7 @@ use dstack_kms_rpc::{ GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, SignCertResponse, }; -use dstack_verifier::{CvmVerifier, VerificationDetails}; +use dstack_verifier::{CvmVerifier, VerificationDetails, MEASUREMENT_CACHE_SUFFIX}; use fs_err as fs; use k256::ecdsa::SigningKey; use moka::future::Cache; @@ -120,10 +120,10 @@ 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<()> { +/// remove a single cache entry (a hex-named subdir/file, plus `suffix`) 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, suffix: &str) -> Result<()> { if sub_dir.is_empty() { return Ok(()); } @@ -136,7 +136,7 @@ fn remove_cache(parent_dir: &Path, sub_dir: &str) -> Result<()> { if !sub_dir.chars().all(|c| c.is_ascii_hexdigit()) { bail!("invalid cache key"); } - let path = parent_dir.join(sub_dir); + let path = parent_dir.join(format!("{sub_dir}{suffix}")); if path.is_dir() { fs::remove_dir_all(path)?; } else if path.is_file() { @@ -151,10 +151,11 @@ impl KmsState { /// 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/. + remove_cache(&images_dir, image_hash, "").context("failed to clear image cache")?; + // measurement cache is kept by the verifier as measurements/.json. let mr_cache_dir = self.config.image.cache_dir.join("measurements"); - remove_cache(&mr_cache_dir, config_hash).context("failed to clear measurement cache")?; + remove_cache(&mr_cache_dir, config_hash, MEASUREMENT_CACHE_SUFFIX) + .context("failed to clear measurement cache")?; Ok(()) } @@ -688,19 +689,33 @@ mod tests { fs::create_dir_all(&keep).unwrap(); fs::create_dir_all(&drop).unwrap(); - remove_cache(root, "beef").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()); + 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(); + remove_cache(root, "all", "").unwrap(); assert!(!root.exists()); } + + #[test] + fn remove_cache_deletes_the_measurement_cache_file() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + let keep = root.join(format!("cafe{MEASUREMENT_CACHE_SUFFIX}")); + let drop = root.join(format!("beef{MEASUREMENT_CACHE_SUFFIX}")); + fs::write(&keep, "{}").unwrap(); + fs::write(&drop, "{}").unwrap(); + + remove_cache(root, "beef", MEASUREMENT_CACHE_SUFFIX).unwrap(); + assert!(!drop.exists(), "named measurement entry must be removed"); + assert!(keep.exists(), "other measurement entries must be kept"); + } use std::collections::BTreeMap; fn hex_of(byte: u8, len: usize) -> String { diff --git a/dstack/verifier/src/lib.rs b/dstack/verifier/src/lib.rs index d5b70a2a6..6e9ed5645 100644 --- a/dstack/verifier/src/lib.rs +++ b/dstack/verifier/src/lib.rs @@ -16,4 +16,4 @@ mod verification; pub use dstack_mr::TdxMeasurements; pub use types::{PolicyBootInfo, VerificationDetails, VerificationRequest, VerificationResponse}; -pub use verification::{policy_tcb_fields, CvmVerifier}; +pub use verification::{policy_tcb_fields, CvmVerifier, MEASUREMENT_CACHE_SUFFIX}; diff --git a/dstack/verifier/src/verification.rs b/dstack/verifier/src/verification.rs index f4f862db2..407671e2c 100644 --- a/dstack/verifier/src/verification.rs +++ b/dstack/verifier/src/verification.rs @@ -96,6 +96,9 @@ fn decode_key_provider_info(bytes: &[u8]) -> Option.json`. +pub const MEASUREMENT_CACHE_SUFFIX: &str = ".json"; + /// Upper bound on both a downloaded image archive and its extracted contents. /// Measurement images carry no rootfs; released ones are under 25 MB extracted. const MAX_IMAGE_SIZE: u64 = 100 << 20; @@ -170,7 +173,7 @@ impl CvmVerifier { fn measurement_cache_path(&self, cache_key: &str) -> PathBuf { self.measurement_cache_dir() - .join(format!("{cache_key}.json")) + .join(format!("{cache_key}{MEASUREMENT_CACHE_SUFFIX}")) } /// Hash only what the measurement reads: `image` and the measurement documents From fd803dda674718d88d82ac4201902c3ed2b82591 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 27 Oct 2026 21:54:28 -0700 Subject: [PATCH 2/2] refactor(kms): let the verifier clear its own cache The image and measurement cache layout is owned by CvmVerifier, so move cache clearing there and drop the exported MEASUREMENT_CACHE_SUFFIX. KMS no longer hard-codes the images/ and measurements/ paths, which is how the .json suffix got out of sync in the first place. Signed-off-by: Kevin Wang --- dstack/kms/src/main_service.rs | 74 ++--------------------------- dstack/verifier/src/lib.rs | 2 +- dstack/verifier/src/verification.rs | 69 +++++++++++++++++++++++++-- 3 files changed, 69 insertions(+), 76 deletions(-) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index a9fd8f90f..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, @@ -18,7 +17,7 @@ use dstack_kms_rpc::{ GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, SignCertResponse, }; -use dstack_verifier::{CvmVerifier, VerificationDetails, MEASUREMENT_CACHE_SUFFIX}; +use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; use k256::ecdsa::SigningKey; use moka::future::Cache; @@ -120,43 +119,12 @@ fn short_lived_cache() -> Cache<(), V> { .build() } -/// remove a single cache entry (a hex-named subdir/file, plus `suffix`) 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, suffix: &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(format!("{sub_dir}{suffix}")); - 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 as measurements/.json. - let mr_cache_dir = self.config.image.cache_dir.join("measurements"); - remove_cache(&mr_cache_dir, config_hash, MEASUREMENT_CACHE_SUFFIX) - .context("failed to clear measurement cache")?; - Ok(()) + self.verifier.clear_cache(image_hash, config_hash) } pub fn new(config: KmsConfig) -> Result { @@ -650,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, @@ -663,42 +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()); - } - - #[test] - fn remove_cache_deletes_the_measurement_cache_file() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path(); - let keep = root.join(format!("cafe{MEASUREMENT_CACHE_SUFFIX}")); - let drop = root.join(format!("beef{MEASUREMENT_CACHE_SUFFIX}")); - fs::write(&keep, "{}").unwrap(); - fs::write(&drop, "{}").unwrap(); - - remove_cache(root, "beef", MEASUREMENT_CACHE_SUFFIX).unwrap(); - assert!(!drop.exists(), "named measurement entry must be removed"); - assert!(keep.exists(), "other measurement entries must be kept"); - } use std::collections::BTreeMap; fn hex_of(byte: u8, len: usize) -> String { diff --git a/dstack/verifier/src/lib.rs b/dstack/verifier/src/lib.rs index 6e9ed5645..d5b70a2a6 100644 --- a/dstack/verifier/src/lib.rs +++ b/dstack/verifier/src/lib.rs @@ -16,4 +16,4 @@ mod verification; pub use dstack_mr::TdxMeasurements; pub use types::{PolicyBootInfo, VerificationDetails, VerificationRequest, VerificationResponse}; -pub use verification::{policy_tcb_fields, CvmVerifier, MEASUREMENT_CACHE_SUFFIX}; +pub use verification::{policy_tcb_fields, CvmVerifier}; diff --git a/dstack/verifier/src/verification.rs b/dstack/verifier/src/verification.rs index fe367786e..949f36d0b 100644 --- a/dstack/verifier/src/verification.rs +++ b/dstack/verifier/src/verification.rs @@ -96,9 +96,6 @@ fn decode_key_provider_info(bytes: &[u8]) -> Option.json`. -pub const MEASUREMENT_CACHE_SUFFIX: &str = ".json"; - /// Upper bound on both a downloaded image archive and its extracted contents. /// Measurement images carry no rootfs; released ones are under 25 MB extracted. const MAX_IMAGE_SIZE: u64 = 100 << 20; @@ -135,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, @@ -173,7 +193,7 @@ impl CvmVerifier { fn measurement_cache_path(&self, cache_key: &str) -> PathBuf { self.measurement_cache_dir() - .join(format!("{cache_key}{MEASUREMENT_CACHE_SUFFIX}")) + .join(format!("{cache_key}.json")) } /// Hash only what the measurement reads: `image` and the measurement documents @@ -366,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() { @@ -1702,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();