Fix safetensor overflow.

mlx/io/gguf.cpp

@@ -51,7 +51,15 @@ Shape get_shape(const gguf_tensor& tensor) {
   Shape shape;
   // The dimension order in GGML is the reverse of the order used in MLX.
   for (int i = tensor.ndim - 1; i >= 0; i--) {
-    shape.push_back(tensor.dim[i]);
+    uint64_t dim = tensor.dim[i];
+    // Reject dimensions that exceed int32 range (ShapeElem is int32_t)
+    if (dim > static_cast<uint64_t>(INT32_MAX)) {
+      std::ostringstream msg;
+      msg << "[load_gguf] tensor dimension " << i << " value " << dim
+          << " exceeds int32 range";
+      throw std::runtime_error(msg.str());
+    }
+    shape.push_back(static_cast<ShapeElem>(dim));
   }
   return shape;
 }

mlx/io/gguf_quants.cpp

@@ -111,32 +111,143 @@ void gguf_load_quantized(
 
   std::string name(tensor.name, tensor.namelen);
 
+  // Issue #4244 fix (2): reject ndim == 0. gguflib permits ndim == 0 but the
+  // code below indexes shape[shape.size()-1].
+  if (tensor.ndim == 0) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " has zero dimensions";
+    throw std::runtime_error(msg.str());
+  }
+
   auto shape = get_shape(tensor);
   const uint64_t weights_per_block = 32;
   if (shape[shape.size() - 1] % weights_per_block != 0) {
     std::ostringstream msg;
     msg << "[load_gguf] tensor " << name
-        << "has incompatible last dim shape: " << shape[shape.size() - 1];
+        << " has incompatible last dim shape: " << shape[shape.size() - 1];
+    throw std::runtime_error(msg.str());
+  }
+
+  // Issue #4244 fix (4): cross-check shape-derived element count against the
+  // GGUF-supplied num_weights using checked 64-bit multiplication.  Without
+  // this the two widths can diverge when the 64-bit product wraps but the
+  // 32-bit Shape still describes billions of elements.
+  uint64_t shape_elem_count = 1;
+  for (auto dim : shape) {
+    if (__builtin_mul_overflow(
+            shape_elem_count, static_cast<uint64_t>(dim), &shape_elem_count)) {
+      std::ostringstream msg;
+      msg << "[load_gguf] tensor " << name << " element count overflow";
+      throw std::runtime_error(msg.str());
+    }
+  }
+  if (shape_elem_count != tensor.num_weights) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name
+        << " shape element count (" << shape_elem_count
+        << ") does not match GGUF num_weights (" << tensor.num_weights << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Issue #4244 PoC drives num_weights to exactly 0 (via wrapping dims), which
+  // yields a zero-byte allocation against a multi-billion-iteration loop.
+  // Reject explicitly so we never malloc(0) and enter the extractors.
+  if (tensor.num_weights == 0) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " has zero num_weights";
     throw std::runtime_error(msg.str());
   }
 
   auto weights_shape = shape;
   weights_shape.back() /= (weights_per_byte * 4);
-  auto w_nbytes = uint32.size() *
-      std::accumulate(weights_shape.begin(),
-                      weights_shape.end(),
-                      1,
-                      std::multiplies<size_t>());
 
-  array weights(allocator::malloc(w_nbytes), std::move(weights_shape), uint32);
+  // Issue #4244 fix (1): std::accumulate init must be size_t{1} (not int
+  // literal 1) otherwise the 64-bit product from multiplies<size_t> is
+  // truncated back to int on each intermediate assignment.  We also guard the
+  // final element-size multiplication with __builtin_mul_overflow.
+  size_t w_nbytes;
+  if (__builtin_mul_overflow(
+          uint32.size(),
+          static_cast<size_t>(std::accumulate(
+              weights_shape.begin(),
+              weights_shape.end(),
+              size_t{1},
+              std::multiplies<size_t>())),
+          &w_nbytes)) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " weights size overflow";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Issue #4244 fix (5): check allocator results before use.  Without this,
+  // applying only the accumulate fix above turns the heap overflow into a
+  // null-pointer dereference SEGV when malloc returns nullptr for huge sizes.
+  void* w_ptr = allocator::malloc(w_nbytes).raw_ptr();
+  if (!w_ptr) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " allocation failed";
+    throw std::runtime_error(msg.str());
+  }
+  array weights(allocator::Buffer(w_ptr), std::move(weights_shape), uint32);
 
   // For scales and bias
   shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block;
-  auto sb_nbytes = float16.size() *
-      std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<size_t>());
 
-  array scales(allocator::malloc(sb_nbytes), shape, float16);
-  array biases(allocator::malloc(sb_nbytes), std::move(shape), float16);
+  size_t sb_nbytes;
+  if (__builtin_mul_overflow(
+          float16.size(),
+          static_cast<size_t>(std::accumulate(
+              shape.begin(), shape.end(), size_t{1}, std::multiplies<size_t>())),
+          &sb_nbytes)) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " scales/biases size overflow";
+    throw std::runtime_error(msg.str());
+  }
+
+  void* sb_ptr = allocator::malloc(sb_nbytes).raw_ptr();
+  if (!sb_ptr) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " scales/biases allocation failed";
+    throw std::runtime_error(msg.str());
+  }
+  array scales(allocator::Buffer(sb_ptr), shape, float16);
+
+  void* b_ptr = allocator::malloc(sb_nbytes).raw_ptr();
+  if (!b_ptr) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " biases allocation failed";
+    throw std::runtime_error(msg.str());
+  }
+  array biases(allocator::Buffer(b_ptr), std::move(shape), float16);
+
+  // Issue #4244 fix (4): the extractors loop on scales_arr.size() (the true
+  // shape-derived count) and read `bytes_per_block` per iteration.  Verify
+  // that the total read span fits within the header-declared bsize so we
+  // can't have a small bsize that passed check_tensor_in_file() paired with
+  // a large Shape that would otherwise drive reads past the mapping.
+  uint64_t bytes_per_block;
+  if (tensor.type == GGUF_TYPE_Q4_0) {
+    bytes_per_block = 18;
+  } else if (tensor.type == GGUF_TYPE_Q4_1) {
+    bytes_per_block = 20;
+  } else {
+    bytes_per_block = 34;
+  }
+  uint64_t total_data_needed;
+  if (__builtin_mul_overflow(
+          scales.size(), bytes_per_block, &total_data_needed)) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name << " data size overflow";
+    throw std::runtime_error(msg.str());
+  }
+  if (total_data_needed > tensor.bsize) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name
+        << " requires " << total_data_needed
+        << " bytes but GGUF header declares " << tensor.bsize;
+    throw std::runtime_error(msg.str());
+  }
+
   if (tensor.type == GGUF_TYPE_Q4_0) {
     extract_q4_0_data(tensor, weights, scales, biases);
   } else if (tensor.type == GGUF_TYPE_Q4_1) {

mlx/io/safetensors.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023 Apple Inc.
 
+#include <climits>
 #include <json.hpp>
 #include <memory>
 #include <sstream>
@@ -181,7 +182,21 @@ SafetensorsLoad load_safetensors(
     {
       size_t expected_nbytes = type.size();
       for (auto dim : shape) {
-        expected_nbytes *= static_cast<size_t>(dim);
+        if (dim < 0) {
+          std::ostringstream msg;
+          msg << "[load_safetensors] Tensor '" << item.key()
+              << "' has negative dimension " << dim;
+          throw std::runtime_error(msg.str());
+        }
+        if (__builtin_mul_overflow(
+                expected_nbytes,
+                static_cast<size_t>(dim),
+                &expected_nbytes)) {
+          std::ostringstream msg;
+          msg << "[load_safetensors] Tensor '" << item.key()
+              << "' element count overflow";
+          throw std::runtime_error(msg.str());
+        }
       }
       if (data_offsets[1] < data_offsets[0] ||
           data_offsets[1] - data_offsets[0] != expected_nbytes) {
@@ -193,7 +208,8 @@ SafetensorsLoad load_safetensors(
         throw std::runtime_error(msg.str());
       }
     }
-    if (offset + data_offsets[1] > file_size) {
+    // Use subtraction instead of addition to avoid overflow wrap
+    if (data_offsets[1] > file_size - offset) {
       std::ostringstream msg;
       msg << "[load_safetensors] Tensor '" << item.key()
           << "' invalid data offsets (" << data_offsets[0] << ", "

tests/load_tests.cpp

@@ -3,6 +3,7 @@
 #include <filesystem>
 #include <fstream>
 #include <stdexcept>
+#include <cstdint>
 #include <vector>
 
 #include "doctest/doctest.h"
@@ -442,3 +443,101 @@ TEST_CASE("test single array serialization") {
     CHECK(array_equal(a, b).item<bool>());
   }
 }
+
+// Helper to write a raw GGUF file with quantized tensor type.
+// dims: tensor dimensions in GGML order (last dimension first in file)
+// ttype: GGUF quantized type (Q4_0=1, Q4_1=2, Q8_0=8)
+// num_weights: the num_weights field from the GGUF header
+// data_len: bytes of filler data after the tensor header
+void write_raw_quantized_gguf(
+    const std::string& path,
+    const std::vector<uint64_t>& dims,
+    uint32_t ttype,
+    uint64_t num_weights,
+    size_t data_len = 0) {
+  std::ofstream out(path, std::ios::binary);
+  auto u32 = [&out](uint32_t v) {
+    out.write(reinterpret_cast<const char*>(&v), 4);
+  };
+  auto u64 = [&out](uint64_t v) {
+    out.write(reinterpret_cast<const char*>(&v), 8);
+  };
+  out.write("GGUF", 4);
+  u32(3); // version
+  u64(1); // tensor_count
+  u64(0); // metadata_kv_count
+  // Tensor name: "w"
+  u64(1); // tensor name length
+  out.write("w", 1);
+  u32(static_cast<uint32_t>(dims.size())); // ndim
+  for (auto d : dims) {
+    u64(d); // dim (GGML order)
+  }
+  u32(ttype); // type
+  u64(num_weights); // num_weights
+  u64(0); // offset (relative to data section)
+  // Align to 32 bytes
+  while (out.tellp() % 32 != 0) {
+    out.put(0);
+  }
+  // Write filler data - always write at least 1KB
+  size_t min_data = std::max(data_len, (size_t)1024);
+  std::vector<char> buf(min_data, 0x41);
+  out.write(buf.data(), min_data);
+}
+
+TEST_CASE("test gguf quantized tensor security") {
+  // Test that crafted quantized GGUF files are rejected rather than causing
+  // heap buffer overflows. See ml-explore/mlx#4244.
+
+  SUBCASE("zero num_weights rejected") {
+    std::string file_path = get_temp_file("test_gguf_qzero.gguf");
+    write_raw_quantized_gguf(file_path, {32, 32}, 8 /* Q8_0 */, 0);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("shape element count mismatch rejected (Q8_0)") {
+    // num_weights doesn't match the shape-derived element count.
+    // dims in GGML order: [32, 64] -> MLX shape [64, 32], element count = 2048
+    // But we pass num_weights = 1, which doesn't match.
+    std::string file_path = get_temp_file("test_gguf_qmismatch.gguf");
+    write_raw_quantized_gguf(file_path, {32, 64}, 8 /* Q8_0 */, 1);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("shape element count mismatch rejected (Q4_0)") {
+    // dims in GGML order: [32, 32] -> MLX shape [32, 32], element count = 1024
+    // But we pass num_weights = 100, which doesn't match.
+    std::string file_path = get_temp_file("test_gguf_qmismatch2.gguf");
+    write_raw_quantized_gguf(file_path, {32, 32}, 1 /* Q4_0 */, 100);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("dimension exceeding int32 rejected") {
+    std::string file_path = get_temp_file("test_gguf_qbigdim.gguf");
+    uint64_t big_dim = static_cast<uint64_t>(INT32_MAX) + 1;
+    write_raw_quantized_gguf(file_path, {big_dim}, 8 /* Q8_0 */, big_dim);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+}
+
+TEST_CASE("test safetensors security") {
+  SUBCASE("negative dimension rejected") {
+    std::string file_path = get_temp_file("test_neg_dim.safetensors");
+    std::string json_header =
+        R"({"tensor":{"dtype":"F32","shape":[-1,10],"data_offsets":[0,40]}})";
+    std::vector<char> data(40, 0);
+    write_raw_safetensors(file_path, json_header, data);
+    CHECK_THROWS_AS(load_safetensors(file_path), std::runtime_error);
+  }
+
+  SUBCASE("negative dimension with offset wrap rejected") {
+    std::string file_path = get_temp_file("test_neg_dim_wrap.safetensors");
+    // shape:[-82] with large data_offsets would cause offset+data_offsets[1] to wrap
+    std::string json_header =
+        R"({"tensor":{"dtype":"U8","shape":[-82],"data_offsets":[0,18446744073709551534]}})";
+    std::vector<char> data(100, 0);
+    write_raw_safetensors(file_path, json_header, data);
+    CHECK_THROWS_AS(load_safetensors(file_path), std::runtime_error);
+  }
+}