hacnho/tensorflow-checkpoint-metadata-abort-poc

TensorFlow Checkpoint metadata abort PoC

This repository contains a proof of concept for TensorFlow Checkpoint reader process aborts caused by attacker-controlled metadata in a .ckpt.index file.

Summary

TensorFlow V2 checkpoints store tensor metadata as TensorBundle entries inside the .index SSTable. The reader treats several untrusted metadata conditions as fatal assertions instead of returning a recoverable error:

• tf.train.load_checkpoint(prefix) aborts when a checksum-valid .ckpt.index contains a malformed BundleEntryProto value.
• reader.get_tensor("weight/.ATTRIBUTES/VARIABLE_VALUE") aborts when the entry uses an invalid DataType enum such as 127.

The control checkpoint loads and reads normally. The malicious checkpoints keep valid SSTable block CRC32C trailers, so TensorFlow reaches the metadata parser instead of rejecting the file at the checksum layer.

Contents

• control/: valid TensorFlow checkpoint.
• malformed_bundle_entry/: malicious checkpoint that aborts during tf.train.load_checkpoint(prefix).
• invalid_dtype/: malicious checkpoint that loads, then aborts during reader.get_tensor(...).
• repro_checkpoint_metadata_abort.py: standalone generator and reproducer.
• evidence/reproduce-result-20260623.json: recorded control/crash output.
• evidence/source-checkpoint-reader-v2.21.0.txt: source snippet showing CHECK(entry.ParseFromString(...)).
• evidence/source-tensor-v2.21.0.txt: source snippet showing the fatal Unexpected type path.

Reproduction

python3 -m venv .venv
. .venv/bin/activate
pip install -r requirements.txt
python repro_checkpoint_metadata_abort.py --workdir reproduce-work --json-out reproduce-result.json

Expected result:

• control_load: exits 0.
• control_get_tensor: exits 0.
• malformed_bundle_entry_load: exits by SIGABRT with checkpoint_reader.cc:122] Check failed: entry.ParseFromString(...).
• invalid_dtype_load: exits 0.
• invalid_dtype_get_tensor: exits by SIGABRT with tensor.cc:1014] Unexpected type: 127.

Environment

• TensorFlow: tensorflow-cpu==2.21.0
• Python: 3.12
• Triggered format: TensorFlow Checkpoint V2 (checkpoint, ckpt-1.index, ckpt-1.data-00000-of-00001)
• Impact: local denial of service when a victim validates, indexes, restores, or reads an untrusted checkpoint.