tfrecord (vahidk/tfrecord) CWE-770 unbounded allocation PoC
Proof-of-concept file for a huntr Model Format Vulnerability report against
vahidk/tfrecord (independent Python TFRecord reader, PyPI package
tfrecord, used as a PyTorch IterableDataset).
evil.tfrecord is a 12-byte file: an 8-byte little-endian uint64 length
field claiming 500 MiB (524288000), followed by 4 bytes of junk in place
of the CRC field (the reader never checks the CRC's value, only that 4
bytes were read).
read_records() (tfrecord/reader.py:61-85) does
datum_bytes = datum_bytes.zfill(int(length * 1.5)) - an unconditional
buffer allocation sized from the claimed length - before checking
whether the file actually contains that many remaining bytes
(file.readinto(...) only raises afterward, on short reads). The format
itself protects against this (writer.py writes a CRC32C of the length
field specifically so a reader can validate it before trusting it), but
reader.py reads the CRC bytes twice and never checks their value - only
the byte count read.
Verified live (fresh clone, commit aa3a2e3975cd3be3d47e68be476a2b5c608dde1b)
under a 300MB ulimit -v cap:
$ ulimit -v 300000
$ python3 -c "from tfrecord.reader import tfrecord_iterator; next(tfrecord_iterator('evil.tfrecord'))"
MemoryError
Reachable via the documented public API:
TFRecordDataset.__iter__ (tfrecord/torch/dataset.py, the standard way to
use this package with a PyTorch DataLoader) -> tfrecord_loader() ->
tfrecord_iterator() -> read_records().
Impact: Denial of Service - an attacker who can supply or influence an
untrusted .tfrecord file (loaded via the standard PyTorch dataset API) can
force unbounded memory allocation from a file as small as 12 bytes.