Validate model variable sizes before allocation

src/models/model.cc

@@ -2,6 +2,8 @@
 
 #include <spdlog/spdlog.h>
 
+#include <limits>
+
 #include "ctranslate2/models/model_factory.h"
 #include "ctranslate2/ops/ops.h"
 #include "ctranslate2/utils.h"
@@ -619,7 +621,6 @@ namespace ctranslate2 {
 
       // Load the variables.
       const auto num_variables = consume<uint32_t>(model_file);
-      model->_variable_index.reserve(num_variables);
 
       // check config for tensor parallel
       bool multi_query_attention = false;
@@ -655,9 +656,31 @@ namespace ctranslate2 {
           num_bytes = consume<uint32_t>(model_file) * item_size;
         }
 
-        StorageView variable(std::move(shape), dtype);
-        if (num_bytes != variable.size() * variable.item_size())
+        const dim_t item_size = StorageView(dtype).item_size();
+        // The payload is the raw tensor data that follows this variable header.
+        // Check that it fits in the file before allocating the StorageView.
+        const auto payload_position = model_file.tellg();
+        if (payload_position == std::streampos(-1))
           throw std::runtime_error("Variable '" + name + "' has an invalid payload size");
+        model_file.seekg(0, std::ios::end);
+        const auto model_end = model_file.tellg();
+        model_file.clear();
+        model_file.seekg(payload_position);
+        if (model_end == std::streampos(-1) || model_end < payload_position)
+          throw std::runtime_error("Variable '" + name + "' has an invalid payload size");
+        dim_t variable_size = 1;
+        for (const dim_t dim : shape) {
+          if (dim == 0 || variable_size > std::numeric_limits<dim_t>::max() / dim)
+            throw std::runtime_error("Variable '" + name + "' has an invalid shape");
+          variable_size *= dim;
+        }
+        if (item_size == 0
+            || variable_size > std::numeric_limits<dim_t>::max() / item_size
+            || num_bytes != variable_size * item_size
+            || static_cast<size_t>(num_bytes) > static_cast<size_t>(model_end - payload_position))
+          throw std::runtime_error("Variable '" + name + "' has an invalid payload size");
+
+        StorageView variable(std::move(shape), dtype);
         consume<char>(model_file, num_bytes, static_cast<char*>(variable.buffer()));
         if (tensor_parallel) {
           int outer_dim = 0;

src/models/model_reader.cc

@@ -40,6 +40,36 @@ namespace ctranslate2 {
         char* p = const_cast<char*>(base);
         setg(p, p, p + size);
       }
+
+      std::streampos seekoff(std::streamoff off,
+                             std::ios_base::seekdir dir,
+                             std::ios_base::openmode which) override {
+        if (!(which & std::ios_base::in))
+          return std::streampos(-1);
+
+        char* begin = eback();
+        char* end = egptr();
+        std::streamoff base = 0;
+        if (dir == std::ios_base::beg)
+          base = 0;
+        else if (dir == std::ios_base::cur)
+          base = gptr() - begin;
+        else if (dir == std::ios_base::end)
+          base = end - begin;
+        else
+          return std::streampos(-1);
+
+        const std::streamoff next = base + off;
+        if (next < 0 || next > end - begin)
+          return std::streampos(-1);
+
+        setg(begin, begin + next, end);
+        return next;
+      }
+
+      std::streampos seekpos(std::streampos pos, std::ios_base::openmode which) override {
+        return seekoff(static_cast<std::streamoff>(pos), std::ios_base::beg, which);
+      }
     };
 
     struct imemstream : virtual membuf, std::istream {

tests/model_test.cc

@@ -1,4 +1,5 @@
 #include <ctranslate2/models/sequence_to_sequence.h>
+#include <ctranslate2/models/model_reader.h>
 
 #include <ctranslate2/decoding.h>
 
@@ -8,6 +9,22 @@ TEST(ModelTest, ContainsModel) {
   ASSERT_TRUE(models::contains_model(default_model_dir()));
 }
 
+TEST(ModelTest, RejectsSerializedVariableZeroDimension) {
+  const std::string model_bin(
+    "\x06\x00\x00\x00\x0c\x00WhisperSpec\0\x03\x00\x00\x00"
+    "\x01\x00\x00\x00\x07\x00weight\0\x02\x01\x00\x00\x00"
+    "\x00\x00\x00\x00\x01\x00\x00\x00\x00", 49);
+
+  models::ModelMemoryReader reader("test_model");
+  reader.register_file("model.bin", model_bin);
+  try {
+    models::Model::load(reader);
+    FAIL() << "Expected runtime_error";
+  } catch (const std::runtime_error& e) {
+    EXPECT_NE(std::string(e.what()).find("invalid shape"), std::string::npos);
+  }
+}
+
 TEST(ModelTest, UpdateDecoderOutputLayer) {
   auto model = models::Model::load(default_model_dir())->as_sequence_to_sequence();
   auto& decoder = dynamic_cast<models::EncoderDecoderReplica&>(*model).decoder();