From 12c110da5c612ecc5b09256ad12ae8c64473b2fb Mon Sep 17 00:00:00 2001 From: robertomeroni Date: Mon, 10 Aug 2026 20:50:33 +0200 Subject: [PATCH] Bound GGUF tensor data offsets against the file mapping gguflib computes a tensor's data pointer as ctx->data + ctx->data_off + the tensor's offset field, in unsigned arithmetic and without comparing the result against the mapping. A crafted offset therefore produces a pointer outside the mapped file, and mlx's own memcpy in extract_tensor_data then reads it. If the addition wraps, the pointer lands back inside the mapping and the wrong bytes are read silently. Validate offset and size against the file in load_arrays, which is the single point both the plain and the quantized paths pass through. --- mlx/io/gguf.cpp | 23 ++++++++++++ tests/load_tests.cpp | 89 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+) 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;