diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py index ea241ada2..bf3e08380 100644 --- a/gguf-py/gguf/gguf_reader.py +++ b/gguf-py/gguf/gguf_reader.py @@ -32,6 +32,10 @@ from gguf.constants import ( GGUFEndian, ) +# limits mirroring ggml/src/gguf.cpp (not part of gguf.h) +GGUF_MAX_STRING_LENGTH = 1024 * 1024 * 1024 +GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024 + logger = logging.getLogger(__name__) READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION] @@ -167,6 +171,10 @@ class GGUFReader: offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64])) offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64])) tensor_count, kv_count = temp_counts + if tensor_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') + if kv_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') offs = self._build_fields(offs, kv_count) # Build Tensor Info Fields @@ -217,6 +225,10 @@ class GGUFReader: def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]: slen = self._get(offset, np.uint64) + if int(slen[0]) > GGUF_MAX_STRING_LENGTH: + raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}') + if offset + 8 + int(slen[0]) > self.data.nbytes: + raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}') return slen, self._get(offset + 8, np.uint8, slen[0]) def _get_field_parts( @@ -241,6 +253,8 @@ class GGUFReader: raw_itype = self._get(offs, np.uint32) offs += int(raw_itype.nbytes) alen = self._get(offs, np.uint64) + if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') offs += int(alen.nbytes) aparts: list[npt.NDArray[Any]] = [raw_itype, alen] data_idxs: list[int] = []