From 4c726067d5363f13c9874bd0c312212b1bda5d97 Mon Sep 17 00:00:00 2001
From: Yi-111-a <153097222+Yi-111-a@users.noreply.github.com>
Date: Fri, 25 Sep 2026 16:03:40 +0800
Subject: [PATCH] fix: reject cyclic safetensors index references

A weight_map entry that points back at its own index, or at another index
that points back, made parse_file recurse until the process was killed.
Track indexes already being loaded and return a normal error instead.
---
 src/model_loader.cpp | 24 ++++++++++++++++++++++++
 src/model_loader.h   |  2 ++
 2 files changed, 26 insertions(+)

diff --git a/src/model_loader.cpp b/src/model_loader.cpp
index 2bd9261db..308907f6a 100644
--- a/src/model_loader.cpp
+++ b/src/model_loader.cpp
@@ -311,7 +311,27 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
     return true;
 }
 
+static std::string canonical_loader_path(const std::string& path) {
+    std::error_code error;
+    const std::filesystem::path canonical = std::filesystem::weakly_canonical(std::filesystem::u8path(path), error);
+    if (!error) {
+        return canonical.generic_string();
+    }
+    return std::filesystem::u8path(path).lexically_normal().generic_string();
+}
+
 bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) {
+    const std::string index_key = canonical_loader_path(file_path);
+    if (!loading_safetensors_indexes_.insert(index_key).second) {
+        LOG_ERROR("cyclic safetensors index reference '%s'", file_path.c_str());
+        return false;
+    }
+    struct SafetensorsIndexGuard {
+        std::unordered_set<std::string>& active;
+        std::string key;
+        ~SafetensorsIndexGuard() { active.erase(key); }
+    } guard{loading_safetensors_indexes_, index_key};
+
     LOG_VERBOSE("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
 
     std::vector<std::string> shard_paths;
@@ -322,6 +342,10 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
     }
 
     for (const std::string& shard_path : shard_paths) {
+        if (loading_safetensors_indexes_.count(canonical_loader_path(shard_path)) != 0) {
+            LOG_ERROR("cyclic safetensors index reference '%s'", shard_path.c_str());
+            return false;
+        }
         if (!parse_file(shard_path, prefix)) {
             return false;
         }
diff --git a/src/model_loader.h b/src/model_loader.h
index e51b260b7..662dd7017 100644
--- a/src/model_loader.h
+++ b/src/model_loader.h
@@ -7,6 +7,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include "model.h"
@@ -62,6 +63,7 @@ class ModelLoader {
     std::string tensor_type_rules_;
     std::vector<FileStamp> parsed_dependencies_;
     std::map<std::string, std::set<std::string>> parsed_tensor_names_;
+    std::unordered_set<std::string> loading_safetensors_indexes_;
 
     static bool read_file_stamp(const std::string& path, FileStamp& stamp);
     static bool file_unchanged(const FileStamp& stamp);