[BUG] mx.load(): unvalidated .npy shape (negative / overflowing) causes heap OOB write and read #4646

Filed as a public issue at the maintainer's request (the private advisory was closed as out-of-scope: model files are treated as trusted input, like loading a script). This is reported here as an input-validation / robustness bug in the .npy / .npz reader, not as a security vulnerability.

Summary

mx.load() parses the array shape from a .npy header (and, by inspection, .npz members go through the same parser) and builds an array without rejecting negative dimensions and without guarding the shape-product against size_t overflow. A malformed shape makes the logical element count (array.size()) disagree with the allocated byte count (nbytes()), because the allocation size wraps while the logical size does not. Evaluating the array then reads/writes past the buffer.

Reproduced on v0.32.3, CPU-only build with -fsanitize=address.

Details

1. Shape parsed without validation — mlx/io/load.cpp:297-298

int dim = std::stoi(shape_str, &pos);   // int, negative accepted, no bound
shape.push_back(dim);

std::stoi("-1") returns -1, which is pushed into the shape as-is.

2. Shape product overflows the descriptor — mlx/array.cpp:251

size *= shape[i];   // size_t *= int32 ShapeElem, no overflow guard

For shape=(-1,), size sign-extends to SIZE_MAX.

3. Allocation wraps — mlx/backend/common/load.cpp:31 + mlx/backend/no_gpu/allocator.cpp:102-104

out.set_data(allocator::malloc(out.nbytes()));   // nbytes() = SIZE_MAX*8 = 2^64-8
void* ptr = std::malloc(size + sizeof(size_t));  // (2^64-8)+8 = 0 -> malloc(0)
*static_cast<size_t*>(ptr) = size;               // 8-byte write into a 1-byte region

Two faults follow from the same unchecked shape:

• shape=(-1,): nbytes() and then size + sizeof(size_t) wrap to 0, so malloc(0) returns a 1-byte region and the write is 8 bytes past it (out-of-bounds write).
• shape=(2097152, 1048576, 1048576): element count is 2^61, so nbytes() wraps to 0 and the buffer is 8 bytes, while size() stays 2^61; a downstream op iterates size() elements over the 8-byte buffer (out-of-bounds read at mlx/backend/cpu/simd/base_simd.h:47).

NumPy itself rejects both of these headers at load time; MLX currently does not.

Reproduction

# gen.py
import struct

def make_npy(path, descr, shape):
    body = "(%s,)" % shape[0] if len(shape) == 1 else "(%s)" % ", ".join(map(str, shape))
    hdr = "{'descr': '%s', 'fortran_order': False, 'shape': %s, }" % (descr, body)
    hdr += " " * ((16 - (11 + len(hdr)) % 16) % 16) + "\n"
    with open(path, "wb") as f:
        f.write(b"\x93NUMPY\x01\x00" + struct.pack("<H", len(hdr)) + hdr.encode() + b"\x00" * 64)

make_npy("neg.npy",   "<f8", (-1,))
make_npy("wrap0.npy", "<f8", (2097152, 1048576, 1048576))

# poc.py — run once per file so the first crash does not mask the second
import sys, mlx.core as mx
a = mx.load(sys.argv[1])
mx.eval(a + mx.array(1, a.dtype))

ASan output (CPU-only, v0.32.3):

$ python poc.py neg.npy
heap-buffer-overflow WRITE of size 8
  CommonAllocator::malloc  mlx/backend/no_gpu/allocator.cpp:104
  Load::eval_cpu           mlx/backend/common/load.cpp:31
0 bytes after a 1-byte region     (a.size == 18446744073709551615)

$ python poc.py wrap0.npy
heap-buffer-overflow READ of size 8
  simd::load               mlx/backend/cpu/simd/base_simd.h:47
  Load::eval_cpu           mlx/backend/common/load.cpp:31
0 bytes after an 8-byte region

A normal file (shape=(4,)) loads and evaluates correctly.

Suggested fix

Validate the parsed shape in the npy reader (mlx/io/load.cpp), right after the shape is parsed and before the array is built:

1. reject negative dimensions;
2. compute the element count (and × itemsize) with checked multiplication and throw on overflow.

As optional defense-in-depth, guard ArrayDesc::init (array.cpp:251) and CommonAllocator::malloc (allocator.cpp:102) against the wrap.

I'll open a PR shortly (same day) implementing (1)+(2) with tests, and link it here.

I'm happy to adjust the approach to whatever you prefer.