Fix: Out-of-bounds read in GGUF metadata loader

mlx/io/gguf.cpp

@@ -3,6 +3,7 @@
 #include <cstdint>
 #include <cstring>
 #include <fstream>
+#include <limits>
 #include <numeric>
 
 #include "mlx/io/gguf.h"
@@ -87,6 +88,139 @@ std::tuple<allocator::Buffer, Dtype> extract_tensor_data(gguf_tensor* tensor) {
   return {buffer, float16};
 }
 
+// Mirror check_tensor_in_file(): gguf_get_key() leaves key.val / ctx->off
+// pointing at the value but performs no bounds checking, and the value lengths
+// are read straight from the file. Bound the whole value (and, for arrays,
+// every element) against the mmap'd file once per key in load_metadata(),
+// before set_mx_value_from_gguf consumes it. Lengths that would not fit in the
+// int the downstream array() / std::string constructors take are rejected here.
+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);
+
+  size_t fixed = 0;
+  switch (type) {
+    case GGUF_VALUE_TYPE_BOOL:
+    case GGUF_VALUE_TYPE_UINT8:
+    case GGUF_VALUE_TYPE_INT8:
+      fixed = 1;
+      break;
+    case GGUF_VALUE_TYPE_UINT16:
+    case GGUF_VALUE_TYPE_INT16:
+      fixed = 2;
+      break;
+    case GGUF_VALUE_TYPE_UINT32:
+    case GGUF_VALUE_TYPE_INT32:
+    case GGUF_VALUE_TYPE_FLOAT32:
+      fixed = 4;
+      break;
+    case GGUF_VALUE_TYPE_UINT64:
+    case GGUF_VALUE_TYPE_INT64:
+    case GGUF_VALUE_TYPE_FLOAT64:
+      fixed = 8;
+      break;
+    default:
+      break;
+  }
+  if (fixed) {
+    if (fixed > avail(base)) {
+      throw std::runtime_error(
+          "[load_gguf] Metadata value extends past the end of the file. "
+          "Perhaps an incomplete download or corrupt file?");
+    }
+    return;
+  }
+
+  // gguf_string = { uint64_t len; char string[] }.
+  if (type == GGUF_VALUE_TYPE_STRING) {
+    if (sizeof(uint64_t) > avail(base) ||
+        val->string.len > static_cast<uint64_t>(std::numeric_limits<int>::max()) ||
+        sizeof(uint64_t) + val->string.len > avail(base)) {
+      throw std::runtime_error(
+          "[load_gguf] String metadata value extends past the end of the file.");
+    }
+    return;
+  }
+
+  // Array header = { uint32_t type; uint64_t len; } (gguf_array_header_size),
+  // followed by the elements.
+  if (type == GGUF_VALUE_TYPE_ARRAY) {
+    if (gguf_array_header_size > avail(base)) {
+      throw std::runtime_error(
+          "[load_gguf] Metadata value extends past the end of the file. "
+          "Perhaps an incomplete download or corrupt file?");
+    }
+    if (val->array.len > static_cast<uint64_t>(std::numeric_limits<int>::max())) {
+      throw std::runtime_error(
+          "[load_gguf] Array metadata value length is too large.");
+    }
+
+    size_t elt_size = 0;
+    switch (val->array.type) {
+      case GGUF_VALUE_TYPE_BOOL:
+      case GGUF_VALUE_TYPE_UINT8:
+      case GGUF_VALUE_TYPE_INT8:
+        elt_size = 1;
+        break;
+      case GGUF_VALUE_TYPE_UINT16:
+      case GGUF_VALUE_TYPE_INT16:
+        elt_size = 2;
+        break;
+      case GGUF_VALUE_TYPE_UINT32:
+      case GGUF_VALUE_TYPE_INT32:
+      case GGUF_VALUE_TYPE_FLOAT32:
+        elt_size = 4;
+        break;
+      case GGUF_VALUE_TYPE_UINT64:
+      case GGUF_VALUE_TYPE_INT64:
+      case GGUF_VALUE_TYPE_FLOAT64:
+        elt_size = 8;
+        break;
+      default:
+        break;
+    }
+    const uint8_t* elt = base + gguf_array_header_size;
+    if (elt_size) {
+      if (val->array.len > avail(elt) / elt_size) {
+        throw std::runtime_error(
+            "[load_gguf] Array metadata value extends past the end of the file.");
+      }
+      return;
+    }
+    // String array: each element is a length-prefixed string, so walk them.
+    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)) {
+          throw std::runtime_error(
+              "[load_gguf] Array metadata value extends past the end of the file.");
+        }
+        uint64_t slen = reinterpret_cast<const gguf_string*>(p)->len;
+        if (slen > static_cast<uint64_t>(std::numeric_limits<int>::max()) ||
+            sizeof(uint64_t) + slen > avail(p)) {
+          throw std::runtime_error(
+              "[load_gguf] Array metadata value extends past the end of the file.");
+        }
+        p += sizeof(uint64_t) + slen;
+      }
+      return;
+    }
+    // Unsupported element type (e.g. nested array): header is bounded; the
+    // consumer rejects the format without reading the payload.
+    return;
+  }
+
+  throw std::runtime_error("[load_gguf] Received unexpected type.");
+}
+
 void set_mx_value_from_gguf(
     gguf_ctx* ctx,
     uint32_t type,
@@ -206,6 +340,7 @@ std::unordered_map<std::string, GGUFMetaData> load_metadata(gguf_ctx* ctx) {
   while (gguf_get_key(ctx, &key)) {
     std::string key_name = std::string(key.name, key.namelen);
     auto& val = metadata.insert({key_name, GGUFMetaData{}}).first->second;
+    check_metadata_value_in_file(ctx, key.type, key.val);
     set_mx_value_from_gguf(ctx, key.type, key.val, val);
   }
   return metadata;

tests/load_tests.cpp

@@ -257,6 +257,114 @@ 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>;