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

diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index c9ec92bd8ed498efc241a1de6d8d3f7ff1d243a7..6b0a26b63d89de68058fd0da676905e5d27dc029 100644 (file)
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -11,6 +11,7 @@ GGUF_MAGIC             = 0x46554747  # "GGUF"
 GGUF_VERSION           = 3
 GGUF_DEFAULT_ALIGNMENT = 32
 GGML_QUANT_VERSION     = 2  # GGML_QNT_VERSION from ggml.h
+GGML_MAX_DIMS          = 4  # GGML_MAX_DIMS from ggml.h
 
 #
 # metadata keys

diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py
index 0a1b85f50641b1abf7f597cdad133100f7884a3b..ea241ada285c679cdb4d36abfa764d126c51ac8e 100644 (file)
--- a/gguf-py/gguf/gguf_reader.py
+++ b/gguf-py/gguf/gguf_reader.py
@@ -22,6 +22,7 @@ if __name__ == "__main__":
     sys.path.insert(0, str(Path(__file__).parent.parent))
 
 from gguf.constants import (
+    GGML_MAX_DIMS,
     GGML_QUANT_SIZES,
     GGUF_DEFAULT_ALIGNMENT,
     GGUF_MAGIC,
@@ -266,6 +267,8 @@ class GGUFReader:
         # Get Tensor Dimensions Count
         n_dims = self._get(offs, np.uint32)
         offs += int(n_dims.nbytes)
+        if n_dims[0] > GGML_MAX_DIMS:
+            raise ValueError(f'Tensor dimensions count {n_dims[0]} exceeds GGML_MAX_DIMS ({GGML_MAX_DIMS})')
 
         # Get Tensor Dimension Array
         dims = self._get(offs, np.uint64, n_dims[0])
@@ -326,7 +329,10 @@ class GGUFReader:
                 raise ValueError(f'Found duplicated tensor with name {tensor_name}')
             tensor_names.add(tensor_name)
             ggml_type = GGMLQuantizationType(raw_dtype[0])
-            n_elems = int(np.prod(dims))
+            # use Python ints: np.prod on uint64 wraps silently on overflow
+            n_elems = 1
+            for dim in dims.tolist():
+                n_elems *= int(dim)
             np_dims = tuple(reversed(dims.tolist()))
             block_size, type_size = GGML_QUANT_SIZES[ggml_type]
             n_bytes = n_elems * type_size // block_size

diff --git a/gguf-py/tests/test_gguf_reader_validation.py b/gguf-py/tests/test_gguf_reader_validation.py
new file mode 100644 (file)
index 0000000..98f30a9
--- /dev/null
+++ b/gguf-py/tests/test_gguf_reader_validation.py
@@ -0,0 +1,37 @@
+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)