fix: add FlatBuffer verification in BuildFromAllocation to prevent heap OOB read#120175
fix: add FlatBuffer verification in BuildFromAllocation to prevent heap OOB read#120175Ashutosh0x wants to merge 2 commits into
Conversation
…rleaveDatasetOp InterleaveDatasetOp::MakeDataset only checks that cycle_length > 0 but enforces no upper bound. When cycle_length is INT64_MAX, the Iterator constructor allocates a vector of that size, causing either an immediate crash (OOM/SIGSEGV) or integer overflow in the modulo operation at line 171: (cycle_index_ + 1) % cycle_length_. Add a reasonable upper bound (1M) for both cycle_length and block_length with separate constants (kMaxCycleLength, kMaxBlockLength). Values above this threshold are rejected with InvalidArgument error. Added unit tests for both overflow cases. Fixes tensorflow#116198
…ap OOB read BuildFromAllocation() calls ValidateModelBuffers() which follows FlatBuffer offsets (model_->buffers()) without bounds checking. A crafted .tflite file with an attacker-controlled root table offset causes heap-buffer-overflow via ReadScalar<int>() in GetVTable(). An 8-byte file is sufficient to trigger the crash, and the OOB read distance is steerable from 0 to 64KB+. Add flatbuffers::Verifier check before calling ValidateModelBuffers(), consistent with how VerifyAndBuildFromAllocation already verifies. Returns nullptr with error on invalid buffers instead of crashing. Fixes tensorflow#115308
There was a problem hiding this comment.
Code Review
This pull request enhances security and robustness across TensorFlow components. It adds flatbuffer structure verification in model_builder_base.h to prevent heap-buffer-overflows on crafted inputs. Additionally, it introduces a maximum limit of 1M for both cycle_length and block_length in interleave_dataset_op.cc to guard against integer overflows and excessive memory allocations, supported by new unit tests in interleave_dataset_op_test.cc. No review comments were provided, and the implementation is clean, robust, and well-tested, so I have no further feedback.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@cantonios This PR fixes a heap OOB read in BuildFromAllocation() — a crafted 8-byte .tflite file can trigger it since there's no FlatBuffer verification before ValidateModelBuffers() accesses the root table offset. ASan confirms OOB up to 64KB+. The fix just adds flatbuffers::Verifier the same way VerifyAndBuildFromAllocation() already does it. All checks pass. Can you take a look? |
Summary
Add
flatbuffers::Verifiercheck inBuildFromAllocation()before callingValidateModelBuffers()to prevent heap out-of-bounds read on crafted.tflitefiles.Vulnerability (#115308)
FlatBufferModelBase::BuildFromAllocation()callsValidateModelBuffers()which accessesmodel_->buffers()— this follows the FlatBuffer root table offset without any bounds checking. A crafted 8-byte.tflitefile with an attacker-controlled root table offset causesReadScalar<int>()inGetVTable()to read out-of-bounds heap memory.ASan confirms heap-buffer-overflow with steerable OOB distance (0 to 64KB+).
Reproduction
\\python
import struct
with open('/tmp/crash.tflite', 'wb') as f:
f.write(struct.pack('<I', 0x18) + b'TFL3')
import tensorflow as tf
tf.lite.Interpreter(model_path='/tmp/crash.tflite') # SIGSEGV
\\
Fix
Add
flatbuffers::Verifiervalidation inBuildFromAllocation()beforeValidateModelBuffers(), consistent with howVerifyAndBuildFromAllocation()already verifies. Returnsnullptrwith error on invalid buffers instead of crashing.File Changed
tensorflow/compiler/mlir/lite/core/model_builder_base.hFixes #115308