Add LoRA path sanitizer (leejet#1156)

Author: roj234

build.bat

@@ -3,3 +3,4 @@ pushd %~dp0
 echo Compiling wait patiently
 clang -Iggml -Ithirdparty entry.cpp -w --shared -DGGML_MAX_NAME=256 -O2 -flto=thin -ffunction-sections -fdata-sections -Lincludes -lstdc++ -lggml -lggml-base -lggml-cpu-x64 -lggml-cuda -DSD_USE_CUDA -DSD_BUILD_DLL -o stable-diffusion.dll -Wl,-s
 clang -Iggml -Ithirdparty -I. -Iexamples -w examples\cli\main.cpp -lstdc++ -lstable-diffusion -Lincludes -o sd-cli.exe -Wl,-s
+clang -Iggml -Ithirdparty -I. -Iexamples -w examples\server\main.cpp -lstdc++ -lstable-diffusion -Lincludes -DSD_SAFE_LORA -o sd-server.exe -Wl,-s -D_WIN32_WINNT=0xA00 -lws2_32

examples/common/common.hpp

@@ -37,6 +37,19 @@ namespace fs = std::filesystem;
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
+static bool is_path_safe(const fs::path& base_dir, const fs::path& target_path) {
+    try {
+        // 获取目标路径的绝对路径（weakly_canonical 允许文件尚不存在，但能处理 .. 和 .）
+        fs::path abs_target = fs::weakly_canonical(target_path);
+
+        auto base_str = base_dir.string();
+        auto target_str = abs_target.string();
+        return target_str.rfind(base_str, 0) == 0;
+    } catch (...) {
+        return false;
+    }
+}
+
 const char* modes_str[] = {
     "img_gen",
     "vid_gen",
@@ -1632,15 +1645,16 @@ struct SDGenerationParams {
     }
 
     void extract_and_remove_lora(const std::string& lora_model_dir) {
-        if (lora_model_dir.empty()) {
-            return;
-        }
+        if (lora_model_dir.empty()) return;
+
         static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
         static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"};
         std::smatch m;
 
         std::string tmp = prompt;
 
+        fs::path base_path = fs::absolute(lora_model_dir);
+
         while (std::regex_search(tmp, m, re)) {
             std::string raw_path      = m[1].str();
             const std::string raw_mul = m[2].str();
@@ -1663,10 +1677,22 @@ struct SDGenerationParams {
 
             fs::path final_path;
             if (is_absolute_path(raw_path)) {
+    #ifdef SD_SAFE_LORA
+                LOG_WARN("This LoRA path is forbidden [SD_SAFE_LORA]: %s", raw_path.c_str());
+                goto next_match;
+    #else
                 final_path = raw_path;
+    #endif
             } else {
-                final_path = fs::path(lora_model_dir) / raw_path;
+                final_path = base_path / raw_path;
+    #ifdef SD_SAFE_LORA
+                if (!is_path_safe(base_path, final_path)) {
+                    LOG_WARN("This LoRA path is forbidden [SD_SAFE_LORA]: %s", raw_path.c_str());
+                    goto next_match;
+                }
+    #endif
             }
+
             if (!fs::exists(final_path)) {
                 bool found = false;
                 for (const auto& ext : valid_ext) {
@@ -1686,13 +1712,17 @@ struct SDGenerationParams {
                 }
             }
 
+{
             const std::string key = final_path.lexically_normal().string();
 
             if (is_high_noise)
                 high_noise_lora_map[key] += mul;
             else
                 lora_map[key] += mul;
+}
+
 
+    next_match:
             prompt = std::regex_replace(prompt, re, "", std::regex_constants::format_first_only);
 
             tmp = m.suffix().str();
@@ -1999,7 +2029,7 @@ uint8_t* load_image_common(bool from_memory,
     int c = 0;
     const char* image_path;
     uint8_t* image_buffer = nullptr;
-    
+
     bool is_qoi = false;
     if (from_memory) {
         // magic bytes
@@ -2039,7 +2069,7 @@ uint8_t* load_image_common(bool from_memory,
                 image_path_or_bytes, &width, &height, &c, expected_channel);
         }
     }
-    
+
     if (image_buffer == nullptr) {
         LOG_ERROR("load image from '%s' failed", image_path);
         return nullptr;
@@ -2180,6 +2210,6 @@ bool save_image_as_qoi(const char* filename, int width, int height, int channels
     desc.height = height;
     desc.channels = channels;
     desc.colorspace = QOI_SRGB;
-    
+
     return qoi_write(filename, data, &desc) > 0;
 }

examples/server/main.cpp

@@ -254,7 +254,7 @@ std::vector<uint8_t> write_image_to_vector(
             desc.height = height;
             desc.channels = channels;
             desc.colorspace = QOI_SRGB;
-            
+
             int out_len = 0;
             void* qoi_data = qoi_encode(image, &desc, &out_len);
             if (qoi_data) {
@@ -289,6 +289,7 @@ int main(int argc, const char** argv) {
     SDContextParams ctx_params;
     SDGenerationParams default_gen_params;
     parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
+    ctx_params.lora_apply_mode = LORA_APPLY_AT_RUNTIME;
 
     sd_set_log_callback(sd_log_cb, (void*)&svr_params);
     log_verbose = svr_params.verbose;
@@ -425,7 +426,7 @@ int main(int argc, const char** argv) {
             if (gen_params.sample_params.sample_steps > 100)
                 gen_params.sample_params.sample_steps = 100;
 
-            if (!gen_params.process_and_check(IMG_GEN, "")) {
+            if (!gen_params.process_and_check(IMG_GEN, ctx_params.lora_model_dir)) {
                 res.status = 400;
                 res.set_content(R"({"error":"invalid params"})", "application/json");
                 return;
@@ -482,8 +483,8 @@ int main(int argc, const char** argv) {
                     continue;
                 }
                 auto image_bytes = write_image_to_vector(
-                    output_format == "jpeg" ? ImageFormat::JPEG : 
-                    output_format == "qoi" ? ImageFormat::QOI : 
+                    output_format == "jpeg" ? ImageFormat::JPEG :
+                    output_format == "qoi" ? ImageFormat::QOI :
                     ImageFormat::PNG,
                                                          results[i].data,
                                                          results[i].width,
@@ -609,7 +610,7 @@ int main(int argc, const char** argv) {
             if (gen_params.sample_params.sample_steps > 100)
                 gen_params.sample_params.sample_steps = 100;
 
-            if (!gen_params.process_and_check(IMG_GEN, "")) {
+            if (!gen_params.process_and_check(IMG_GEN, ctx_params.lora_model_dir)) {
                 res.status = 400;
                 res.set_content(R"({"error":"invalid params"})", "application/json");
                 return;
@@ -705,8 +706,8 @@ int main(int argc, const char** argv) {
                 if (results[i].data == nullptr)
                     continue;
                 auto image_bytes = write_image_to_vector(
-                    output_format == "jpeg" ? ImageFormat::JPEG : 
-                    output_format == "qoi" ? ImageFormat::QOI : 
+                    output_format == "jpeg" ? ImageFormat::JPEG :
+                    output_format == "qoi" ? ImageFormat::QOI :
                     ImageFormat::PNG,
                                                          results[i].data,
                                                          results[i].width,