EnigmaConsultant/tf-tensor-bundle-nullderef-poc

TensorFlow tensor_bundle checkpoint reader: NULL-pointer-dereference DoS via oversized DT_STRING shape

Target: github.com/tensorflow/tensorflow — tensorflow/core/util/tensor_bundle/tensor_bundle.cc (BundleReader::GetValue) and tensorflow/core/framework/tensor.cc (Tensor::TotalBytes / Helper<tstring>::TotalBytes / TypedAllocator::Allocate).

Tested against: tensorflow-cpu==2.21.0 (v2.21.0-rc1-5-ga481b10260d), installed from PyPI, unmodified. Root cause confirmed present in the current tensorflow/tensorflow master branch source as of 2026-07-06.

Summary

Loading an attacker-supplied TensorFlow checkpoint whose metadata (.index) file declares a DT_STRING tensor with an astronomically large 1-D shape crashes the host process with a SIGSEGV (NULL pointer dereference). No corresponding .data-00000-of-00001 shard file is required at all — the crash happens purely while TensorFlow validates the tensor's metadata, before it ever tries to read the tensor's byte range from disk. The malicious .index file in this repo is 142 bytes.

This is reachable through the standard, widely used checkpoint-inspection API: tf.train.load_checkpoint(prefix).get_tensor(key) (this is what tensorflow.python.training.py_checkpoint_reader.CheckpointReader.get_tensor / NewCheckpointReader / the inspect_checkpoint.py CLI tool / many downstream libraries that peek at checkpoint contents by key use under the hood).

Root cause

1. BundleReader::GetValue() (tensor_bundle.cc) constructs a fresh Tensor straight from the untrusted BundleEntryProto's dtype/shape fields:

   const TensorShape stored_shape(TensorShape(entry.shape()));
   if (val->NumElements() == 0) {
     ret = new Tensor(entry.dtype(), stored_shape);
   }
   

   TensorShape::IsValid() only rejects shapes whose element count overflows int64; it happily accepts a single dimension of, e.g., 2^62 + 7 elements.

2. Tensor::Tensor(Allocator*, DataType, TensorShape) (tensor.cc) allocates the backing buffer via TypedAllocator::Allocate<T>(), which by design returns nullptr when the tensor is too large to allocate (typed_allocator.h: "May return NULL if the tensor has too many elements to represent in a single allocation."). The Tensor constructor does not check this and unconditionally wraps the null-backed Buffer<T> — the resulting Tensor looks "constructed" but is not actually backed by memory.

3. Back in GetValue(), right after constructing that Tensor, for DT_STRING entries the code computes a sanity bound by calling ret->TotalBytes():

   const size_t lower_bound = ret->NumElements() + ret->TotalBytes() -
                               sizeof(tstring) * ret->NumElements();
   

   Tensor::TotalBytes() dispatches to Helper<tstring>::TotalBytes(buf_, shape_.num_elements()). That specialization's second parameter is declared as a 32-bit int (unlike the generic template, which uses int64_t), so the 64-bit element count is silently truncated. It then unconditionally dereferences the (null) string array:

   static int64_t TotalBytes(TensorBuffer* in, int n) {
     int64_t tot = in->size();
     const tstring* p = in->base<const tstring>();   // == nullptr
     for (int i = 0; i < n; ++i, ++p) tot += p->size();  // crash: p->size() on nullptr
     return tot;
   }
   

   With n > 0 (guaranteed by choosing the low 32 bits of the shape to be a small positive number) this is a guaranteed NULL pointer dereference — no OOM, no huge malloc attempt, no timing dependency. It happens before entry.size(), entry.offset(), entry.crc32c(), or the .data shard are ever consulted.

Attacker model / reachability

The attacker only needs to control the checkpoint's .index metadata file (a simple leveldb/sstable-style table of serialized BundleEntryProtos) — a totally standard "malicious model/checkpoint shared with a victim" scenario (e.g. a checkpoint uploaded to a model hub, or bundled inside a SavedModel/ckpt directory a user is asked to load or merely inspect). get_tensor()/load_checkpoint() is exactly the API many tools and libraries use to peek at a checkpoint's contents (list/print/convert weights) without first building or restoring a full graph, so the crash is reachable without the victim doing anything beyond pointing a checkpoint inspection call at the attacker's file.

Note: tf.raw_ops.RestoreV2 (the op behind tf.train.Checkpoint.restore()) goes through OpKernelContext::allocate_output, which does check for a failed allocation and returns a clean ResourceExhaustedError instead of crashing — so the vulnerable surface is specifically the direct BundleReader::Lookup / CheckpointReader::GetTensor C++ API (get_tensor, get_variable_to_dtype_map is safe, only get_tensor on a DT_STRING key triggers it), not the RestoreV2 graph op.

Reproduction

pip install tensorflow-cpu==2.21.0 crc32c
python3 build_malicious_ckpt.py /tmp/poison_ckpt/model.ckpt   # writes only model.ckpt.index (142 bytes)
python3 poc_trigger_fh.py /tmp/poison_ckpt/model.ckpt

Observed output (see crash_evidence.log for a full transcript):

Loading checkpoint: /tmp/poison_ckpt/model.ckpt
dtype map: {'poison': tf.string}
Calling get_tensor('poison') -- expecting SIGSEGV from the real BundleReader/Tensor code path
Fatal Python error: Segmentation fault

Current thread 0x00007f4ea2d82200 (most recent call first):
  File ".../tensorflow/python/training/py_checkpoint_reader.py", line 66 in get_tensor
  File "poc_trigger_fh.py", line 10 in <module>

Shell exit code is 139 (128 + SIGSEGV) in both poc_trigger.py (no faulthandler) and poc_trigger_fh.py (with faulthandler.enable(), which confirms the crash is a genuine hardware signal, not a caught Python exception).

Files

• build_malicious_ckpt.py — stand-alone leveldb/sstable-table encoder (re-implemented from xla/tsl/lib/io/{format,table_builder,block_builder}.cc + tensor_bundle.proto) that hand-crafts the malicious 142-byte .index metadata file. No TensorFlow write APIs are used to build the malicious file — it is built entirely from raw bytes to demonstrate a real attacker only needs to control the file format, not any TensorFlow API.
• poison_model.ckpt.index — the pre-built malicious checkpoint metadata (142 bytes, no .data shard needed).
• poc_trigger.py, poc_trigger_fh.py — trigger scripts via tf.train.load_checkpoint(...).get_tensor().
• poc_restorev2.py — comparison script showing the (safe) tf.raw_ops.RestoreV2 behavior for context.
• crash_evidence.log — captured transcript of the actual crash on tensorflow-cpu==2.21.0.

Impact

Denial of service: any process that inspects/loads an attacker-supplied checkpoint via tf.train.load_checkpoint(...).get_tensor() (or equivalent CheckpointReader usage) crashes immediately and unrecoverably (SIGSEGV, not a catchable Python exception) on a 142-byte malicious input file, with no .data shard required.

Scope note (for the triager)

TensorFlow's own SECURITY.md states that "loading untrusted checkpoints or graphs is equivalent to running untrusted code" and that memory corruption is only considered a TensorFlow-side security issue when reachable through a "production-grade, benign model." We flag this transparently: TensorFlow upstream would likely triage a checkpoint-triggered crash as out-of-policy for a TF-side CVE/advisory. We are reporting it here because huntr's TensorFlow checkpoint-reader bounty target treats "a malicious/untrusted model or checkpoint file crashes the loader" as precisely the in-scope threat model for the ML-supply-chain bounty program (that is the premise of the model-file-format bounty category), and the crash is a concrete, deterministic, 100%-reproducible NULL dereference (not a fuzz-only theoretical OOM) triggered by 142 bytes with no real tensor payload.