Minimize change to main branch.

Author: Feli

mlx/io/gguf.cpp

@@ -132,8 +132,7 @@ void set_mx_value_from_gguf(
       value = array(val->boolval, bool_);
       break;
     case GGUF_VALUE_TYPE_STRING:
-      value =
-          std::string(val->string.string, static_cast<int>(val->string.len));
+      value = std::string(val->string.string, val->string.len);
       break;
     case GGUF_VALUE_TYPE_FLOAT64:
       value = array(val->float64, float32);
@@ -182,7 +181,7 @@ void set_mx_value_from_gguf(
           for (auto& str : strs) {
             auto str_val = reinterpret_cast<gguf_string*>(data);
             data += (str_val->len + sizeof(gguf_string));
-            str = std::string(str_val->string, static_cast<int>(str_val->len));
+            str = std::string(str_val->string, str_val->len);
             ctx->off += (str_val->len + sizeof(gguf_string));
           }
           value = std::move(strs);
@@ -208,21 +207,109 @@ void set_mx_value_from_gguf(
   }
 }
 
+inline size_t gguf_value_type_size(uint32_t type) {
+  switch (type) {
+    case GGUF_VALUE_TYPE_BOOL:
+    case GGUF_VALUE_TYPE_UINT8:
+    case GGUF_VALUE_TYPE_INT8:
+      return 1;
+    case GGUF_VALUE_TYPE_UINT16:
+    case GGUF_VALUE_TYPE_INT16:
+      return 2;
+    case GGUF_VALUE_TYPE_UINT32:
+    case GGUF_VALUE_TYPE_INT32:
+    case GGUF_VALUE_TYPE_FLOAT32:
+      return 4;
+    case GGUF_VALUE_TYPE_UINT64:
+    case GGUF_VALUE_TYPE_INT64:
+    case GGUF_VALUE_TYPE_FLOAT64:
+      return 8;
+    default:
+      return 0;
+  }
+}
+
+void check_metadata_value_in_file(
+    const gguf_ctx* ctx,
+    uint32_t type,
+    const gguf_value* val) {
+  auto end = ctx->data + ctx->size;
+  // Bytes available from a pointer up to the end of the mapping; 0 if the
+  // pointer lies outside [ctx->data, end].
+  auto avail = [&](const uint8_t* p) -> size_t {
+    return (p < ctx->data || p > end) ? 0 : static_cast<size_t>(end - p);
+  };
+  auto base = reinterpret_cast<const uint8_t*>(val);
+  auto fail = [](const char* what) {
+    std::ostringstream msg;
+    msg << "[load_gguf] " << what
+        << " Perhaps an incomplete download or corrupt file?";
+    throw std::runtime_error(msg.str());
+  };
+
+  size_t fixed = gguf_value_type_size(type);
+  if (fixed) {
+    if (fixed > avail(base)) {
+      fail("Metadata value extends past the end of the file.");
+    }
+    return;
+  }
+
+  auto check_string = [&](const uint8_t* p) -> const uint8_t* {
+    uint64_t len = reinterpret_cast<const gguf_string*>(p)->len;
+    if (sizeof(uint64_t) + len > avail(p)) {
+      fail("String metadata value extends past the end of the file.");
+    }
+    return p + sizeof(uint64_t) + len;
+  };
+
+  if (type == GGUF_VALUE_TYPE_STRING) {
+    if (sizeof(uint64_t) > avail(base)) {
+      fail("String metadata value extends past the end of the file.");
+    }
+    check_string(base);
+    return;
+  }
+
+  if (type == GGUF_VALUE_TYPE_ARRAY) {
+    if (gguf_array_header_size > avail(base)) {
+      fail("Metadata value extends past the end of the file.");
+    }
+    const uint8_t* elt = base + gguf_array_header_size;
+    size_t elt_size = gguf_value_type_size(val->array.type);
+    if (elt_size) {
+      if (val->array.len > avail(elt) / elt_size) {
+        fail("Array metadata value extends past the end of the file.");
+      }
+      return;
+    }
+    if (val->array.type == GGUF_VALUE_TYPE_STRING) {
+      const uint8_t* p = elt;
+      for (uint64_t i = 0; i < val->array.len; i++) {
+        if (sizeof(uint64_t) > avail(p)) {
+          fail("Array metadata value extends past the end of the file.");
+        }
+        p = check_string(p);
+      }
+    }
+    return;
+  }
+
+  throw std::runtime_error("[load_gguf] Received unexpected type.");
+}
+
 std::unordered_map<std::string, GGUFMetaData> load_metadata(gguf_ctx* ctx) {
   std::unordered_map<std::string, GGUFMetaData> metadata;
   gguf_key key;
   while (gguf_get_key(ctx, &key)) {
+    check_metadata_value_in_file(ctx, key.type, key.val);
     std::string key_name = std::string(key.name, key.namelen);
     auto& val = metadata.insert({key_name, GGUFMetaData{}}).first->second;
     set_mx_value_from_gguf(ctx, key.type, key.val, val);
   }
   return metadata;
 }
 
-// gguflib computes weights_data as ctx->data + ctx->data_off + the tensor's
-// offset field in unsigned arithmetic, without comparing the result against the
-// mapping, so a crafted offset can point outside the file or -- if the addition
-// wraps -- back inside it at the wrong bytes.
 void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) {
   auto fail = [&tensor](const std::string& what) {
     std::ostringstream msg;

mlx/io/gguf_quants.cpp

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <cstring>
+#include <numeric>
 
 #include "mlx/io/gguf.h"
 
@@ -121,27 +122,28 @@ void gguf_load_quantized(
   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());
   }
 
   auto weights_shape = shape;
   weights_shape.back() /= (weights_per_byte * 4);
 
-  auto checked_product = [&](const Shape& dims) {
-    size_t product = 1;
-    for (auto dim : dims) {
-      if (__builtin_mul_overflow(product, static_cast<size_t>(dim), &product)) {
-        std::ostringstream msg;
-        msg << "[load_gguf] tensor " << name << " shape size overflow";
-        throw std::runtime_error(msg.str());
-      }
+  auto checked_product = [&](size_t product, ShapeElem dim) {
+    if (__builtin_mul_overflow(product, static_cast<size_t>(dim), &product)) {
+      std::ostringstream msg;
+      msg << "[load_gguf] tensor " << name << " shape size overflow";
+      throw std::runtime_error(msg.str());
     }
     return product;
   };
 
+  const size_t weights_count = std::accumulate(
+      weights_shape.begin(), weights_shape.end(), size_t{1}, checked_product);
+
   shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block;
-  const size_t block_count = checked_product(shape);
+  const size_t block_count =
+      std::accumulate(shape.begin(), shape.end(), size_t{1}, checked_product);
 
   const uint64_t bytes_per_block = tensor.type == GGUF_TYPE_Q4_0
       ? 18
@@ -158,45 +160,35 @@ void gguf_load_quantized(
     throw std::runtime_error(msg.str());
   }
 
-  const size_t weights_bytes_per_block =
-      tensor.type == GGUF_TYPE_Q8_0 ? 32 : 16;
   size_t w_nbytes;
-  if (__builtin_mul_overflow(block_count, weights_bytes_per_block, &w_nbytes)) {
+  if (__builtin_mul_overflow(uint32.size(), weights_count, &w_nbytes)) {
     std::ostringstream msg;
     msg << "[load_gguf] tensor " << name << " weights size overflow";
     throw std::runtime_error(msg.str());
   }
 
-  auto weights_buffer = allocator::malloc(w_nbytes);
-  if (!weights_buffer.raw_ptr()) {
-    std::ostringstream msg;
-    msg << "[load_gguf] tensor " << name << " allocation failed";
-    throw std::runtime_error(msg.str());
-  }
-  array weights(weights_buffer, std::move(weights_shape), uint32);
+  auto checked_malloc = [&](size_t size, const char* what) {
+    auto buffer = allocator::malloc(size);
+    if (!buffer.raw_ptr()) {
+      std::ostringstream msg;
+      msg << "[load_gguf] tensor " << name << " " << what
+          << " allocation failed";
+      throw std::runtime_error(msg.str());
+    }
+    return buffer;
+  };
+
+  array weights(
+      checked_malloc(w_nbytes, "weights"), std::move(weights_shape), uint32);
 
   size_t sb_nbytes;
   if (__builtin_mul_overflow(float16.size(), block_count, &sb_nbytes)) {
     std::ostringstream msg;
     msg << "[load_gguf] tensor " << name << " scales/biases size overflow";
     throw std::runtime_error(msg.str());
   }
-
-  auto scales_buffer = allocator::malloc(sb_nbytes);
-  if (!scales_buffer.raw_ptr()) {
-    std::ostringstream msg;
-    msg << "[load_gguf] tensor " << name << " scales/biases allocation failed";
-    throw std::runtime_error(msg.str());
-  }
-  array scales(scales_buffer, shape, float16);
-
-  auto biases_buffer = allocator::malloc(sb_nbytes);
-  if (!biases_buffer.raw_ptr()) {
-    std::ostringstream msg;
-    msg << "[load_gguf] tensor " << name << " biases allocation failed";
-    throw std::runtime_error(msg.str());
-  }
-  array biases(biases_buffer, std::move(shape), float16);
+  array scales(checked_malloc(sb_nbytes, "scales"), shape, float16);
+  array biases(checked_malloc(sb_nbytes, "biases"), std::move(shape), float16);
 
   if (tensor.type == GGUF_TYPE_Q4_0) {
     extract_q4_0_data(tensor, weights, scales, biases);

tests/load_tests.cpp

@@ -1,6 +1,5 @@
 // Copyright © 2023 Apple Inc.
 
-#include <cstdint>
 #include <filesystem>
 #include <fstream>
 #include <stdexcept>
@@ -258,6 +257,120 @@ TEST_CASE("test gguf tensor data offset validation") {
   }
 }
 
+// Writes a metadata-only GGUF (no tensors) whose metadata KV section is
+// `kv_section` verbatim, so a caller can encode values whose lengths exceed the
+// file to exercise check_metadata_value_in_file(). `kv_count` must match the
+// number of KV pairs encoded in `kv_section`.
+void write_raw_gguf_metadata(
+    const std::string& path,
+    uint64_t kv_count,
+    const std::vector<char>& kv_section) {
+  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(0); // tensor_count
+  u64(kv_count); // metadata_kv_count
+  out.write(kv_section.data(), kv_section.size());
+}
+
+TEST_CASE("test gguf metadata value validation") {
+  // A STRING/ARRAY metadata value claiming a length larger than the file must
+  // be rejected rather than read past the end of the mapping. See PR #4212.
+
+  auto append_string_kv = [](std::vector<char>& b,
+                             const std::string& key,
+                             uint64_t claimed_len,
+                             bool write_payload) {
+    auto put = [&](const void* p, size_t n) {
+      b.insert(
+          b.end(),
+          static_cast<const char*>(p),
+          static_cast<const char*>(p) + n);
+    };
+    uint64_t klen = key.size();
+    put(&klen, 8);
+    put(key.data(), key.size());
+    uint32_t vt = 8; // GGUF_VALUE_TYPE_STRING
+    put(&vt, 4);
+    put(&claimed_len, 8);
+    if (write_payload) {
+      b.insert(b.end(), claimed_len, '\0');
+    }
+  };
+
+  auto append_array_kv = [](std::vector<char>& b,
+                            const std::string& key,
+                            uint32_t elt_type,
+                            uint64_t claimed_len) {
+    auto put = [&](const void* p, size_t n) {
+      b.insert(
+          b.end(),
+          static_cast<const char*>(p),
+          static_cast<const char*>(p) + n);
+    };
+    uint64_t klen = key.size();
+    put(&klen, 8);
+    put(key.data(), key.size());
+    uint32_t vt = 9; // GGUF_VALUE_TYPE_ARRAY
+    put(&vt, 4);
+    put(&elt_type, 4);
+    put(&claimed_len, 8);
+  };
+
+  SUBCASE("valid empty and small strings load") {
+    std::vector<char> kv;
+    append_string_kv(kv, "empty", 0, false);
+    append_string_kv(kv, "small", 5, true);
+    std::string file_path = get_temp_file("test_gguf_meta_ok.gguf");
+    write_raw_gguf_metadata(file_path, 2, kv);
+    auto [weights, metadata] = load_gguf(file_path);
+    CHECK(weights.empty());
+    CHECK(std::get<std::string>(metadata.at("empty")) == "");
+    CHECK(std::get<std::string>(metadata.at("small")) == std::string(5, '\0'));
+  }
+
+  SUBCASE("string length extends past the end of the file") {
+    // Claims 100 bytes of payload, none of which are present.
+    std::vector<char> kv;
+    append_string_kv(kv, "s", 100, false);
+    std::string file_path = get_temp_file("test_gguf_meta_str_past.gguf");
+    write_raw_gguf_metadata(file_path, 1, kv);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("string length far past the end of the file") {
+    std::vector<char> kv;
+    append_string_kv(kv, "s", 1ull << 40, false);
+    std::string file_path = get_temp_file("test_gguf_meta_str_far.gguf");
+    write_raw_gguf_metadata(file_path, 1, kv);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("fixed-size array length extends past the end of the file") {
+    // GGUF_VALUE_TYPE_UINT8 = 0; claims 2^40 elements, none present.
+    std::vector<char> kv;
+    append_array_kv(kv, "a", 0, 1ull << 40);
+    std::string file_path = get_temp_file("test_gguf_meta_arr_past.gguf");
+    write_raw_gguf_metadata(file_path, 1, kv);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+
+  SUBCASE("string array element length extends past the end of the file") {
+    // GGUF_VALUE_TYPE_STRING = 8; two elements, neither present.
+    std::vector<char> kv;
+    append_array_kv(kv, "a", 8, 2);
+    std::string file_path = get_temp_file("test_gguf_meta_strarr_past.gguf");
+    write_raw_gguf_metadata(file_path, 1, kv);
+    CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
+  }
+}
+
 TEST_CASE("test gguf metadata") {
   std::string file_path = get_temp_file("test_arr.gguf");
   using dict = std::unordered_map<std::string, array>;
@@ -510,7 +623,7 @@ TEST_CASE("test gguf quantized tensor security") {
 
   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;
+    uint64_t big_dim = 1ull << 31;
     write_raw_quantized_gguf(file_path, {big_dim}, 8 /* Q8_0 */, big_dim);
     CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error);
   }