Merge pull request #187 from Dstack-TEE/kms-clear-cache

kms: Add RPC to clear image cache

Author: kvinwang

kms/dstack-app/compose-dev.yaml

@@ -66,6 +66,7 @@ services:
       - /var/run/dstack.sock:/var/run/dstack.sock
     environment:
       - IMAGE_DOWNLOAD_URL=${IMAGE_DOWNLOAD_URL}
+      - ADMIN_TOKEN_HASH=${ADMIN_TOKEN_HASH}
     ports:
       - 8000:8000
     depends_on:

kms/dstack-app/deploy-to-vmm.sh

@@ -40,6 +40,9 @@ GIT_REV=HEAD
 
 # The DStack OS image name to use for the KMS app
 OS_IMAGE=dstack-0.5.0
+
+# The admin token for the KMS app
+ADMIN_TOKEN=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 32 | head -n 1)
 EOF
   echo "Please edit the .env file and set the required variables, then run this script again."
   exit 1
@@ -69,6 +72,8 @@ COMPOSE_TMP=$(mktemp)
 
 GIT_REV=$(git rev-parse $GIT_REV)
 
+ADMIN_TOKEN_HASH=$(echo -n $ADMIN_TOKEN | sha256sum | cut -d' ' -f1)
+
 cp compose-dev.yaml "$COMPOSE_TMP"
 
 subvar() {
@@ -79,6 +84,7 @@ subvar ETH_RPC_URL
 subvar KMS_CONTRACT_ADDR
 subvar GIT_REV
 subvar IMAGE_DOWNLOAD_URL
+subvar ADMIN_TOKEN_HASH
 
 echo "Docker compose file:"
 cat "$COMPOSE_TMP"

kms/dstack-app/entrypoint.sh

@@ -2,6 +2,9 @@
 set -e
 
 cat <<EOF > ./kms.toml
+[core]
+admin_token_hash = "${ADMIN_TOKEN_HASH}"
+
 [core.image]
 verify = true
 cache_dir = "./images"

kms/kms.toml

@@ -21,6 +21,7 @@ mandatory = false
 [core]
 cert_dir = "/etc/kms/certs"
 subject_postfix = ".dstack"
+admin_token_hash = ""
 
 [core.image]
 verify = true

kms/rpc/proto/kms_rpc.proto

@@ -89,6 +89,14 @@ service KMS {
   rpc GetTempCaCert(google.protobuf.Empty) returns (GetTempCaCertResponse);
   // Sign a certificate
   rpc SignCert(SignCertRequest) returns (SignCertResponse);
+  // Clear the image cache
+  rpc ClearImageCache(ClearImageCacheRequest) returns (google.protobuf.Empty);
+}
+
+message ClearImageCacheRequest {
+  string token = 1;
+  string image_hash = 2;
+  string config_hash = 3;
 }
 
 message BootstrapRequest {

kms/src/config.rs

@@ -34,6 +34,8 @@ pub(crate) struct KmsConfig {
     pub auth_api: AuthApi,
     pub onboard: OnboardConfig,
     pub image: ImageConfig,
+    #[serde(with = "serde_human_bytes")]
+    pub admin_token_hash: Vec<u8>,
 }
 
 impl KmsConfig {

kms/src/main_service.rs

@@ -1,11 +1,15 @@
-use std::{ffi::OsStr, path::Path, sync::Arc};
+use std::{
+    ffi::OsStr,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use anyhow::{bail, Context, Result};
 use dstack_kms_rpc::{
     kms_server::{KmsRpc, KmsServer},
-    AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse,
-    GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest,
-    SignCertResponse,
+    AppId, AppKeyResponse, ClearImageCacheRequest, GetAppKeyRequest, GetKmsKeyRequest,
+    GetMetaResponse, GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse,
+    SignCertRequest, SignCertResponse,
 };
 use dstack_types::VmConfig;
 use fs_err as fs;
@@ -149,8 +153,41 @@ impl RpcHandler {
             .await
     }
 
+    fn image_cache_dir(&self) -> PathBuf {
+        self.state.config.image.cache_dir.join("images")
+    }
+
+    fn mr_cache_dir(&self) -> PathBuf {
+        self.state.config.image.cache_dir.join("computed")
+    }
+
+    fn remove_cache(&self, parent_dir: &PathBuf, sub_dir: &str) -> Result<()> {
+        if sub_dir.is_empty() {
+            return Ok(());
+        }
+        if sub_dir == "all" {
+            fs::remove_dir_all(parent_dir)?;
+        } else {
+            let path = parent_dir.join(sub_dir);
+            if path.is_dir() {
+                fs::remove_dir_all(path)?;
+            } else {
+                fs::remove_file(path)?;
+            }
+        }
+        Ok(())
+    }
+
+    fn ensure_admin(&self, token: &str) -> Result<()> {
+        let token_hash = sha2::Sha256::new_with_prefix(token).finalize();
+        if token_hash.as_slice() != self.state.config.admin_token_hash.as_slice() {
+            bail!("Invalid token");
+        }
+        Ok(())
+    }
+
     fn get_cached_mrs(&self, key: &str) -> Result<Mrs> {
-        let path = self.state.config.image.cache_dir.join("computed").join(key);
+        let path = self.mr_cache_dir().join(key);
         if !path.exists() {
             bail!("Cached MRs not found");
         }
@@ -161,7 +198,7 @@ impl RpcHandler {
     }
 
     fn cache_mrs(&self, key: &str, mrs: &Mrs) -> Result<()> {
-        let path = self.state.config.image.cache_dir.join("computed").join(key);
+        let path = self.mr_cache_dir().join(key);
         fs::create_dir_all(path.parent().unwrap()).context("Failed to create cache directory")?;
         safe_write::safe_write(
             &path,
@@ -194,7 +231,7 @@ impl RpcHandler {
         }
 
         // Create a directory for the image if it doesn't exist
-        let image_dir = self.state.config.image.cache_dir.join(&hex_os_image_hash);
+        let image_dir = self.image_cache_dir().join(&hex_os_image_hash);
         // Check if metadata.json exists, if not download the image
         let metadata_path = image_dir.join("metadata.json");
         if !metadata_path.exists() {
@@ -253,7 +290,7 @@ impl RpcHandler {
             .replace("{OS_IMAGE_HASH}", hex_os_image_hash);
 
         // Create a temporary directory for extraction within the cache directory
-        let cache_dir = self.state.config.image.cache_dir.join("tmp");
+        let cache_dir = self.image_cache_dir().join("tmp");
         fs::create_dir_all(&cache_dir).context("Failed to create cache directory")?;
         let auto_delete_temp_dir = tempfile::Builder::new()
             .prefix("tmp-download-")
@@ -578,6 +615,15 @@ impl KmsRpc for RpcHandler {
             ],
         })
     }
+
+    async fn clear_image_cache(self, request: ClearImageCacheRequest) -> Result<()> {
+        self.ensure_admin(&request.token)?;
+        self.remove_cache(&self.image_cache_dir(), &request.image_hash)
+            .context("Failed to clear image cache")?;
+        self.remove_cache(&self.mr_cache_dir(), &request.config_hash)
+            .context("Failed to clear MR cache")?;
+        Ok(())
+    }
 }
 
 impl RpcCall<KmsState> for RpcHandler {