Skip to content

Commit 0db103b

Browse files
authored
Use FlatTensorHeader::kHeaderOffset for the FlatTensor size check (pytorch#22137)
### Summary `FlatTensorDataMap::load` checks that the buffer is at least as long as the flatbuffer header before it reads the root table offset out of it. It computed that length as `sizeof(flatbuffers::uoffset_t) + flatbuffers::kFileIdentifierLength`. `kFileIdentifierLength` only exists at namespace scope in newer flatbuffers releases. In flatbuffers 1.12 the same constant exists only as a member of `FlatBufferBuilder`, so a build against an older flatbuffers fails to compile: ``` flat_tensor_data_map.cpp:303:40: error: no member named 'kFileIdentifierLength' in namespace 'flatbuffers' ``` The submodule pinned in this repository is v24.3.25, where the namespace scope name does exist, so CI here cannot see the problem. ### What this changes The number being computed is 8: a 4 byte root table offset followed by a 4 byte file identifier. `FlatTensorHeader::kHeaderOffset` already names exactly that number, because the FlatTensor header begins right after the flatbuffer header ends. The matching bounds check in `Program::load` already uses `ExtendedHeader::kHeaderOffset` for the same reason. So this swaps the expression for `FlatTensorHeader::kHeaderOffset`. The bound does not move, behavior is unchanged, and no include is added, `flat_tensor_header.h` was already included by this file. ### Test plan Compiled `extension/flat_tensor/flat_tensor_data_map.cpp` with `-std=c++17 -Wall` against the pinned flatbuffers v24.3.25. Clean, no warnings and no errors. Checked that the fix is load bearing by compiling both spellings against both flatbuffers versions: | | old expression | new expression | | --- | --- | --- | | flatbuffers 1.12.0 | fails with the error above | compiles | | flatbuffers 24.3.25 | compiles | compiles | The probe also `static_assert`s that the new expression equals the old one and equals 8, so the two spellings really do produce the same bound.
1 parent cbcdfc1 commit 0db103b

1 file changed

Lines changed: 5 additions & 5 deletions

File tree

‎extension/flat_tensor/flat_tensor_data_map.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -296,11 +296,11 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
296296
// so a corrupt offset would otherwise let version()/named_data()/segments()
297297
// dereference memory outside the buffer.
298298
//
299-
// Minimum size: root offset (uoffset_t) + file identifier, i.e. the
300-
// flatbuffer header. The identifier check above already implies at least this
301-
// many bytes, but check explicitly so the bound below cannot underflow.
302-
constexpr size_t kMinBufferSize =
303-
sizeof(flatbuffers::uoffset_t) + flatbuffers::kFileIdentifierLength;
299+
// Minimum size: root offset + file identifier, i.e. the flatbuffer header
300+
// before the FlatTensor header begins. The identifier check above already
301+
// implies at least this many bytes, but check explicitly so the bound below
302+
// cannot underflow.
303+
constexpr size_t kMinBufferSize = FlatTensorHeader::kHeaderOffset;
304304
ET_CHECK_OR_RETURN_ERROR(
305305
flat_tensor_data->size() >= kMinBufferSize,
306306
InvalidExternalData,

0 commit comments

Comments
 (0)