diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 40cca573e5..90d7b82407 100644 --- a/mlx/io/gguf.cpp +++ b/mlx/io/gguf.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "mlx/io/gguf.h" @@ -87,6 +88,139 @@ std::tuple 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(end - p); + }; + auto base = reinterpret_cast(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(std::numeric_limits::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(std::numeric_limits::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(p)->len; + if (slen > static_cast(std::numeric_limits::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 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; diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 8974919476..f844f26359 100644 --- a/tests/load_tests.cpp +++ b/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& kv_section) { + std::ofstream out(path, std::ios::binary); + auto u32 = [&out](uint32_t v) { + out.write(reinterpret_cast(&v), 4); + }; + auto u64 = [&out](uint64_t v) { + out.write(reinterpret_cast(&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& 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(p), static_cast(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& 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(p), static_cast(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 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(metadata.at("empty")) == ""); + CHECK(std::get(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 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 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 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 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;