Files
llama.cpp/gguf-py/tests/test_gguf_reader_validation.py
hclandGitHub 5788b510a1 gguf-py: validate n_dims and guard against uint64 overflow in reader (#25401)
The Python GGUF reader lacked two guards the C++ loader has:
- n_dims read as uint32 with no GGML_MAX_DIMS bound -> crafted file with
  huge n_dims triggers oversized memmap read / OOM.
- np.prod(dims) on uint64 wraps silently -> a crafted dims triple can
  overflow to a tiny element count, passing an undersized read through.

Add a GGML_MAX_DIMS check and compute the element count with Python ints.

Fixes #25378
2026-08-04 12:12:48 +03:00

38 lines
1.2 KiB
Python

import struct
import numpy as np
import pytest
from gguf.gguf_reader import GGUFReader
def _write_gguf(path, n_dims_field, dims):
buf = b'GGUF' + struct.pack('<IQQ', 3, 1, 0) # version 3, 1 tensor, 0 kv
name = b'bad_tensor'
buf += struct.pack('<Q', len(name)) + name
buf += struct.pack('<I', n_dims_field)
for d in dims:
buf += struct.pack('<Q', d)
buf += struct.pack('<I', 0) # dtype F32
buf += struct.pack('<Q', 0) # tensor offset
buf += b'\x00' * 64
path.write_bytes(buf)
def test_n_dims_upper_bound(tmp_path):
# crafted file claims 1_000_000 dims; must be rejected, not read past EOF
p = tmp_path / 'evil_ndims.gguf'
_write_gguf(p, 1_000_000, [1] * 8)
with pytest.raises(ValueError, match='exceeds GGML_MAX_DIMS'):
GGUFReader(p)
def test_dims_product_no_uint64_wraparound(tmp_path):
# dims whose true product overflows uint64; np.prod would wrap to 4 and
# silently pass an undersized read. The reader must not accept it.
dims = [4194305, 4194305, 211106198978564]
assert int(np.prod(np.array(dims, dtype=np.uint64))) == 4 # the wrap bug
p = tmp_path / 'evil_overflow.gguf'
_write_gguf(p, len(dims), dims)
with pytest.raises(ValueError):
GGUFReader(p)