fix: prevent torch checkpoint offset overflow (#1832)

src/model_io/pickle_io.cpp

@@ -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) {
+static bool checked_pickle_byte_count(int64_t element_count,
-    if (tensor.tensor_storage.nelements() == 0) {
+                                      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));

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;
+            uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
-            current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
+            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();
+            int64_t tensor_nbytes   = tensor_storage.nbytes_to_read();
-            if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
+            if (tensor_nbytes < 0 ||
+                tensor_storage.offset > storage_nbytes ||
+                static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
                 return false;
             }
 

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();
+        int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
-        if (tensor_storage.offset + tensor_nbytes > entry_size) {
+        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;
         }

src/model_loader.cpp

@@ -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_buffer.resize(entry_size);
-                                zip_entry_noallocread(zip, (void*)read_buffer.data(), 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);