From 27ff8298fbe57fc8ce0dcde4f8bd829a6ea5dc80 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 26 Mar 2026 08:15:01 +0000
Subject: [PATCH 1/2] fix: prevent path traversal in KMS remove_cache

---
 kms/src/main_service.rs | 56 +++++++++++++++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 8 deletions(-)

diff --git a/kms/src/main_service.rs b/kms/src/main_service.rs
index 9965fc8d..35a935cf 100644
--- a/kms/src/main_service.rs
+++ b/kms/src/main_service.rs
@@ -2,7 +2,10 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{path::PathBuf, sync::Arc};
+use std::{
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use anyhow::{bail, Context, Result};
 use dstack_kms_rpc::{
@@ -142,20 +145,57 @@ impl RpcHandler {
         self.state.config.image.cache_dir.join("images")
     }
 
-    fn remove_cache(&self, parent_dir: &PathBuf, sub_dir: &str) -> Result<()> {
+    fn remove_cache(&self, parent_dir: &Path, 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)?;
+
+            return Ok(());
+        }
+
+        let sub_path = Path::new(sub_dir);
+        if sub_path.is_absolute() {
+            bail!("Invalid cache path");
+        }
+
+        let mut cleaned = PathBuf::new();
+        for component in sub_path.components() {
+            use std::path::Component;
+
+            match component {
+                Component::Normal(part) => cleaned.push(part),
+                Component::CurDir => {}
+                Component::ParentDir | Component::RootDir | Component::Prefix(_) => {
+                    bail!("Invalid cache path");
+                }
             }
         }
+
+        if cleaned.as_os_str().is_empty() {
+            // Only separators or current-dir components – nothing to do.
+            return Ok(());
+        }
+
+        let path = parent_dir.join(&cleaned);
+
+        let canonical_parent = parent_dir
+            .canonicalize()
+            .unwrap_or_else(|_| parent_dir.to_path_buf());
+        let canonical = path
+            .canonicalize()
+            .unwrap_or_else(|_| canonical_parent.join(&cleaned));
+
+        if !canonical.starts_with(&canonical_parent) {
+            bail!("Invalid cache path");
+        }
+
+        if canonical.is_dir() {
+            fs::remove_dir_all(canonical)?;
+        } else {
+            fs::remove_file(canonical)?;
+        }
         Ok(())
     }
 

From 01a89cb262896d93d3e2c97b3e2e463b2f89403b Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Fri, 27 Mar 2026 11:01:01 +0000
Subject: [PATCH 2/2] kms: restrict clear_image_cache key to hex

---
 kms/src/main_service.rs | 45 ++++++++---------------------------------
 1 file changed, 8 insertions(+), 37 deletions(-)

diff --git a/kms/src/main_service.rs b/kms/src/main_service.rs
index 35a935cf..cb682c7b 100644
--- a/kms/src/main_service.rs
+++ b/kms/src/main_service.rs
@@ -149,53 +149,24 @@ impl RpcHandler {
         if sub_dir.is_empty() {
             return Ok(());
         }
+
         if sub_dir == "all" {
             fs::remove_dir_all(parent_dir)?;
-
             return Ok(());
         }
 
-        let sub_path = Path::new(sub_dir);
-        if sub_path.is_absolute() {
-            bail!("Invalid cache path");
+        if !sub_dir.chars().all(|c| c.is_ascii_hexdigit()) {
+            bail!("Invalid cache key");
         }
 
-        let mut cleaned = PathBuf::new();
-        for component in sub_path.components() {
-            use std::path::Component;
-
-            match component {
-                Component::Normal(part) => cleaned.push(part),
-                Component::CurDir => {}
-                Component::ParentDir | Component::RootDir | Component::Prefix(_) => {
-                    bail!("Invalid cache path");
-                }
-            }
-        }
+        let path = parent_dir.join(sub_dir);
 
-        if cleaned.as_os_str().is_empty() {
-            // Only separators or current-dir components – nothing to do.
-            return Ok(());
+        if path.is_dir() {
+            fs::remove_dir_all(path)?;
+        } else if path.is_file() {
+            fs::remove_file(path)?;
         }
 
-        let path = parent_dir.join(&cleaned);
-
-        let canonical_parent = parent_dir
-            .canonicalize()
-            .unwrap_or_else(|_| parent_dir.to_path_buf());
-        let canonical = path
-            .canonicalize()
-            .unwrap_or_else(|_| canonical_parent.join(&cleaned));
-
-        if !canonical.starts_with(&canonical_parent) {
-            bail!("Invalid cache path");
-        }
-
-        if canonical.is_dir() {
-            fs::remove_dir_all(canonical)?;
-        } else {
-            fs::remove_file(canonical)?;
-        }
         Ok(())
     }