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
This commit is contained in:
@@ -11,6 +11,7 @@ GGUF_MAGIC = 0x46554747 # "GGUF"
|
|||||||
GGUF_VERSION = 3
|
GGUF_VERSION = 3
|
||||||
GGUF_DEFAULT_ALIGNMENT = 32
|
GGUF_DEFAULT_ALIGNMENT = 32
|
||||||
GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h
|
GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h
|
||||||
|
GGML_MAX_DIMS = 4 # GGML_MAX_DIMS from ggml.h
|
||||||
|
|
||||||
#
|
#
|
||||||
# metadata keys
|
# metadata keys
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ if __name__ == "__main__":
|
|||||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||||
|
|
||||||
from gguf.constants import (
|
from gguf.constants import (
|
||||||
|
GGML_MAX_DIMS,
|
||||||
GGML_QUANT_SIZES,
|
GGML_QUANT_SIZES,
|
||||||
GGUF_DEFAULT_ALIGNMENT,
|
GGUF_DEFAULT_ALIGNMENT,
|
||||||
GGUF_MAGIC,
|
GGUF_MAGIC,
|
||||||
@@ -266,6 +267,8 @@ class GGUFReader:
|
|||||||
# Get Tensor Dimensions Count
|
# Get Tensor Dimensions Count
|
||||||
n_dims = self._get(offs, np.uint32)
|
n_dims = self._get(offs, np.uint32)
|
||||||
offs += int(n_dims.nbytes)
|
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
|
# Get Tensor Dimension Array
|
||||||
dims = self._get(offs, np.uint64, n_dims[0])
|
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}')
|
raise ValueError(f'Found duplicated tensor with name {tensor_name}')
|
||||||
tensor_names.add(tensor_name)
|
tensor_names.add(tensor_name)
|
||||||
ggml_type = GGMLQuantizationType(raw_dtype[0])
|
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()))
|
np_dims = tuple(reversed(dims.tolist()))
|
||||||
block_size, type_size = GGML_QUANT_SIZES[ggml_type]
|
block_size, type_size = GGML_QUANT_SIZES[ggml_type]
|
||||||
n_bytes = n_elems * type_size // block_size
|
n_bytes = n_elems * type_size // block_size
|
||||||
|
|||||||
@@ -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)
|
||||||
Reference in New Issue
Block a user