DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-10 06:07:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: validate safetensors data offsets (#1754)

Browse Source
This commit is contained in:
leejet 2026-07-06 22:21:08 +08:00 committed by GitHub
parent 4fcc6fef10
commit e22272ee63
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/model_io/safetensors_io.cpp
View File

@ -43,7 +43,7 @@ bool is_safetensors_file(const std::string& file_path) {
} }
size_t header_size_ = model_io::read_u64(header_size_buf); size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_ || header_size_ <= 2) { if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN || header_size_ <= 2) {
return false; return false;
} }
@ -114,10 +114,11 @@ bool read_safetensors_file(const std::string& file_path,
} }
size_t header_size_ = model_io::read_u64(header_size_buf); size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_) { if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN) {
set_error(error, "invalid safetensor file '" + file_path + "'"); set_error(error, "invalid safetensor file '" + file_path + "'");
return false; return false;
} }
const size_t data_start = ST_HEADER_SIZE_LEN + header_size_;
// read header // read header
std::vector<char> header_buf; std::vector<char> header_buf;
@ -156,6 +157,10 @@ bool read_safetensors_file(const std::string& file_path,
size_t begin = tensor_info["data_offsets"][0].get<size_t>(); size_t begin = tensor_info["data_offsets"][0].get<size_t>();
size_t end = tensor_info["data_offsets"][1].get<size_t>(); size_t end = tensor_info["data_offsets"][1].get<size_t>();
if (begin > end || end > file_size_ - data_start) {
set_error(error, "data offsets out of bounds for tensor '" + name + "'");
return false;
}
ggml_type type = safetensors_dtype_to_ggml_type(dtype); ggml_type type = safetensors_dtype_to_ggml_type(dtype);
if (type == GGML_TYPE_COUNT) { if (type == GGML_TYPE_COUNT) {
@ -187,7 +192,7 @@ bool read_safetensors_file(const std::string& file_path,
n_dims = 1; n_dims = 1;
} }
TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin); TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
tensor_storage.reverse_ne(); tensor_storage.reverse_ne();
size_t tensor_data_size = end - begin; size_t tensor_data_size = end - begin;