Bound GGUF metadata string/array values against the file mapping

.github/workflows/release.yml

@@ -16,7 +16,7 @@ on:
         required: false
         type: boolean
   schedule:
-    - cron: 33 6 * * 1-5
+    - cron: 33 6 * * *
 
 # In jobs we must use |*publish| instead of |inputs.publish| because we can not
 # set default value for workflow_dispatch inputs reliably.

.github/workflows/update_bypass_list.yml

@@ -5,8 +5,9 @@ on:
   workflow_dispatch:
   pull_request_target:
     types:
-      - opened
       - closed
+  schedule:
+    - cron: 33 6 * * *
 
 permissions:
   contents: write

mlx/io/gguf.cpp

@@ -124,8 +124,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);
@@ -174,7 +173,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);
@@ -200,21 +199,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;

tests/load_tests.cpp

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