diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 758488e4a8..40cca573e5 100644 --- a/mlx/io/gguf.cpp +++ b/mlx/io/gguf.cpp @@ -211,6 +211,28 @@ 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; + msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) + << "' " << what << ". Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + }; + if (tensor.offset < ctx->data_off) { + fail("has a data offset that overflows the data section"); + } + if (tensor.offset > ctx->size) { + fail("has a data offset past the end of the file"); + } + if (tensor.bsize > ctx->size - tensor.offset) { + fail("extends past the end of the file"); + } +} + std::unordered_map load_arrays(gguf_ctx* ctx) { std::unordered_map array_map; gguf_tensor tensor; @@ -225,6 +247,7 @@ std::unordered_map load_arrays(gguf_ctx* ctx) { }; while (gguf_get_tensor(ctx, &tensor)) { + check_tensor_in_file(ctx, tensor); if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1 || tensor.type == GGUF_TYPE_Q8_0) { gguf_load_quantized(array_map, tensor); diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 33ad34bc15..8974919476 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -168,6 +168,95 @@ TEST_CASE("test gguf") { } } +// Writes a one-tensor GGUF (name "t", ndim 1, dim 4, type F32) whose tensor +// data offset field is set verbatim to `tensor_data_offset`. Writes +// `data_bytes` bytes of tensor data, defaulting to the full four floats. +void write_raw_gguf( + const std::string& path, + uint64_t tensor_data_offset, + size_t data_bytes = 4 * sizeof(float)) { + 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(1); // tensor_count + u64(0); // metadata_kv_count + u64(1); // tensor name length + out.write("t", 1); + u32(1); // ndim + u64(4); // dim[0] + u32(0); // GGUF_TYPE_F32 + u64(tensor_data_offset); + while (out.tellp() % 32 != 0) { // default GGUF alignment + out.put(0); + } + std::vector data(data_bytes, 0); + out.write(data.data(), data.size()); +} + +TEST_CASE("test gguf tensor data offset validation") { + // A crafted tensor data offset must be rejected rather than turned into a + // pointer outside the mapping. See ml-explore/mlx#4136. + SUBCASE("valid offset loads") { + std::string file_path = get_temp_file("test_gguf_offset_ok.gguf"); + write_raw_gguf(file_path, 0); + auto [weights, metadata] = load_gguf(file_path); + CHECK_EQ(weights.size(), 1); + CHECK(array_equal(weights.at("t"), zeros({4}, float32)).item()); + } + + SUBCASE("offset past the end of the file") { + std::string file_path = get_temp_file("test_gguf_offset_past_end.gguf"); + write_raw_gguf(file_path, 1ull << 20); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset far past the end of the file") { + std::string file_path = get_temp_file("test_gguf_offset_far_past.gguf"); + write_raw_gguf(file_path, 1ull << 40); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset that overflows the data section base") { + // Wraps back to an in-mapping address, so an end-pointer-only check would + // silently read the wrong bytes instead of reading out of bounds. + std::string file_path = get_temp_file("test_gguf_offset_wrap.gguf"); + write_raw_gguf(file_path, ~0ull - 8); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor extends past the end of the file") { + // In-range offset, but the data is truncated: only the extent check + // catches this. + std::string file_path = get_temp_file("test_gguf_truncated.gguf"); + write_raw_gguf(file_path, 0, 4 * sizeof(float) - 1); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor starts inside the file but ends past it") { + // A small offset, so the start is in range and only the extent decides. + // Pins the extent check to the tensor's own start rather than to the + // start of the data section. + std::string file_path = get_temp_file("test_gguf_partial_overrun.gguf"); + write_raw_gguf(file_path, 8); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("offset just past the end of the file") { + // Only a few bytes past the end rather than far outside it, so the + // resulting pointer is still in the mapped page and reads succeed + // silently. Pins the offset bound to the file size exactly. + std::string file_path = get_temp_file("test_gguf_offset_just_past.gguf"); + write_raw_gguf(file_path, 20); + 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;