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: prevent torch checkpoint offset overflow (#1832)

Browse Source
This commit is contained in:
leejet 2026-07-29 23:16:29 +08:00 committed by GitHub
parent 9cfe2af8f9
commit e92e86fb11
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 121 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
src/model_io/pickle_io.cpp
View File

@ -2,6 +2,7 @@
#include <cstdlib>
#include <cstring>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>
@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
return false;
}
static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
if (tensor.tensor_storage.nelements() == 0) {
static bool checked_pickle_byte_count(int64_t element_count,
uint64_t element_nbytes,
uint64_t* byte_count) {
if (element_count < 0 || element_nbytes == 0) {
return false;
}
uint64_t count = static_cast<uint64_t>(element_count);
if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
return false;
}
*byte_count = count * element_nbytes;
return true;
}
static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
if (raw_element_nbytes == 0) {
return false;
}
bool has_zero_dimension = false;
uint64_t element_count = 1;
for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
int64_t dimension = tensor.tensor_storage.ne[i];
if (dimension < 0) {
return false;
}
if (dimension == 0) {
has_zero_dimension = true;
continue;
}
uint64_t size = static_cast<uint64_t>(dimension);
if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
return false;
}
element_count *= size;
}
if (!has_zero_dimension &&
element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
return false;
}
if (has_zero_dimension) {
return true;
}
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
return false;
}
storage.nbytes = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
if (!checked_pickle_byte_count(pid.items[4].int_value,
storage.raw_element_nbytes,
&storage.nbytes)) {
set_error(error, "invalid storage size in torch pickle");
return false;
}
storage_nbytes[storage.key] = storage.nbytes;
stack.push_back(make_storage_value(storage));
} break;
@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.tensor_storage.is_f64 = args.items[0].storage.is_f64;
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
tensor.tensor_storage.storage_key = args.items[0].storage.key;
tensor.tensor_storage.offset = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
if (!checked_pickle_byte_count(args.items[1].int_value,
args.items[0].storage.raw_element_nbytes,
&tensor.tensor_storage.offset)) {
set_error(error, "invalid tensor storage offset in torch pickle");
return false;
}
for (const auto& item : args.items[2].items) {
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.stride[tensor.stride_n_dims++] = item.int_value;
}
if (!tensor_is_contiguous(tensor)) {
if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
set_error(error, "invalid tensor shape or stride in torch pickle");
return false;
}
stack.push_back(make_tensor_value(tensor));

17
src/model_io/torch_legacy_io.cpp
View File

@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
if (it == legacy_storage_map.end()) {
return false;
}
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
if (current_offset > file_size ||
LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
return false;
}
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
if (it->second > file_size - storage_offset) {
return false;
}
storage_offsets[storage_key] = storage_offset;
current_offset = storage_offset + it->second;
}
for (auto& tensor_storage : tensor_storages) {
@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
uint64_t base_offset = it_offset->second;
uint64_t storage_nbytes = it_size->second;
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > storage_nbytes ||
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
return false;
}

6
src/model_io/torch_zip_io.cpp
View File

@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
return false;
}
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > entry_size) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > entry_size ||
static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
return false;
}

48
src/model_loader.cpp
View File

@ -1077,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::vector<uint8_t> read_buffer;
std::vector<uint8_t> convert_buffer;
std::vector<uint8_t> zip_entry_buffer;
while (true) {
int64_t t0, t1;
@ -1115,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
auto read_data = [&](char* buf, size_t n) {
auto read_data = [&](char* buf, size_t n) -> bool {
if (zip != nullptr) {
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
return false;
}
size_t entry_size = zip_entry_size(zip);
if (tensor_storage.offset > entry_size) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
if (n > entry_size - tensor_offset) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
if (entry_size != n) {
int64_t t_memcpy_start;
read_buffer.resize(entry_size);
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
zip_entry_buffer.resize(entry_size);
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
t_memcpy_start = ggml_time_ms();
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
} else {
zip_entry_noallocread(zip, (void*)buf, n);
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
}
zip_entry_close(zip);
} else if (mmapped) {
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
} else {
file.seekg(tensor_storage.offset);
file.read(buf, n);
if (!file) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
}
return true;
};
char* read_buf = nullptr;
@ -1176,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}
t0 = ggml_time_ms();
read_data(read_buf, nbytes_to_read);
if (!read_data(read_buf, nbytes_to_read)) {
failed = true;
break;
}
t1 = ggml_time_ms();
read_time_ms.fetch_add(t1 - t0);