diff --git a/src/model_io/pickle_io.cpp b/src/model_io/pickle_io.cpp index 1ea7a902c..a597e66b3 100644 --- a/src/model_io/pickle_io.cpp +++ b/src/model_io/pickle_io.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -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(element_count); + if (count > std::numeric_limits::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(dimension); + if (element_count > static_cast(std::numeric_limits::max()) / size) { + return false; + } + element_count *= size; + } + + if (!has_zero_dimension && + element_count > static_cast(std::numeric_limits::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)); diff --git a/src/model_io/torch_legacy_io.cpp b/src/model_io/torch_legacy_io.cpp index 0b0c78f8f..3170fe8d4 100644 --- a/src/model_io/torch_legacy_io.cpp +++ b/src/model_io/torch_legacy_io.cpp @@ -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(tensor_nbytes) > storage_nbytes - tensor_storage.offset) { return false; } diff --git a/src/model_io/torch_zip_io.cpp b/src/model_io/torch_zip_io.cpp index 9eaf6c53a..a5db29a11 100644 --- a/src/model_io/torch_zip_io.cpp +++ b/src/model_io/torch_zip_io.cpp @@ -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(tensor_nbytes) > entry_size - tensor_storage.offset) { set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'"); return false; } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index f1436b8a7..69f3d8915 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -1077,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, std::vector read_buffer; std::vector convert_buffer; + std::vector 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(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(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(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); diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 2f44b0003..8db303dda 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -54,9 +54,9 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) { } static bool device_supports_param_op(ggml_backend_dev_t device, - ggml_tensor* weight, - enum ggml_op op, - ggml_backend_buffer_type_t buft) { + ggml_tensor* weight, + enum ggml_op op, + ggml_backend_buffer_type_t buft) { if (op == GGML_OP_NONE) { return true; } diff --git a/src/model_manager.h b/src/model_manager.h index 881a17714..3dc633864 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -131,9 +131,9 @@ class ModelManager : public RunnerWeightManager { ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size = nullptr, - bool allow_split_buffer = false, - bool params_follow_compute_backend = false, + size_t* registered_tensor_size = nullptr, + bool allow_split_buffer = false, + bool params_follow_compute_backend = false, const std::map* tensor_ops = nullptr); bool unregister_param_tensors(const std::string& desc,