diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d10493efff..b53f003ff4 100644 --- a/.github/workflows/release.yml +++ b/.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. diff --git a/.github/workflows/update_bypass_list.yml b/.github/workflows/update_bypass_list.yml index a129c51ce2..b9b71eee8b 100644 --- a/.github/workflows/update_bypass_list.yml +++ b/.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 diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 40cca573e5..6c27c46987 100644 --- a/mlx/io/gguf.cpp +++ b/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(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(data); data += (str_val->len + sizeof(gguf_string)); - str = std::string(str_val->string, static_cast(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,10 +199,102 @@ 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(end - p); + }; + auto base = reinterpret_cast(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(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 load_metadata(gguf_ctx* ctx) { std::unordered_map 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); @@ -211,10 +302,6 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { 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; diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 8974919476..6ef7bc276e 100644 --- a/tests/load_tests.cpp +++ b/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& 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;