Bounds-check the FlatTensor root table offset before reading it (pytorch#22118)

### Summary

`FlatTensorDataMap::load()` reads the flatbuffer root table without
first bounds-checking the root offset:

```cpp
const flat_tensor_flatbuffer::FlatTensor* flat_tensor =
    flat_tensor_flatbuffer::GetFlatTensor(flat_tensor_data->data());
// then immediately:
flat_tensor->named_data();   // walks the root vtable
flat_tensor->segments();
```

`GetFlatTensor()` interprets the `uoffset_t` at the start of the buffer
as the root-table offset and returns `buf + offset`. The first field
access then walks that table's vtable. Nothing on this path runs full
flatbuffer verification, so a corrupt or truncated PTD file with a bad
root offset makes these accessors dereference memory **outside the
buffer** (an OOB read, flagged by ASan).

The program loader already guards against exactly this in
`runtime/executor/program.cpp`. This change mirrors that guard for PTD
files.

### What changed

After the existing identifier and alignment checks, before
`GetFlatTensor()`:
- Confirm the buffer is at least a flatbuffer header long (`uoffset_t` +
file identifier).
- Confirm the root offset is `>= kMinBufferSize` and leaves room for a
vtable `soffset_t` within the buffer.
- Return `InvalidExternalData` otherwise (the same error the surrounding
checks use).

FlatTensor is **not** size-prefixed — both writers finish the buffer
plain (`builder.Finish` in C++, default `flatc` in Python) and the
reader uses `GetFlatTensor` (not the size-prefixed variant) — so the
root offset is at byte 0, exactly as in the program case.

### Scope

This is a pre-existing hardening, independent of the schema-version work
in pytorch#22114 (the version gate this protects is additive; the OOB path
exists today via `named_data()`/`segments()`).

### Test plan

New `FlatTensorDataMapTest.RejectsOutOfBoundsRootOffset`: copies a valid
PTD into a max-aligned buffer, overwrites only the root offset with a
value past the end (leaving identifier + alignment intact), and asserts
`load()` returns `InvalidExternalData`. Without the check, that same
input is an out-of-bounds vtable read under ASan.

> Note: I was unable to build/run the C++ suite in my local environment,
so I'm relying on this PR's CI (which builds the runtime and runs
`extension/flat_tensor` tests, including under ASan) as the
authoritative check. The change mirrors an existing, tested guard in
`program.cpp`.

Author: shoumikhin

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -289,6 +289,38 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
       flat_tensor_data->data(),
       kMinimumAlignment);
 
+  // Verify that the root table offset is within bounds before dereferencing it.
+  // GetFlatTensor() below reads the root offset from the first bytes of the
+  // buffer and returns a pointer at buf + root_offset; reading any field then
+  // walks that table's vtable. Nothing here runs full flatbuffer verification,
+  // so a corrupt offset would otherwise let version()/named_data()/segments()
+  // dereference memory outside the buffer.
+  //
+  // Minimum size: root offset (uoffset_t) + file identifier, i.e. the
+  // flatbuffer header. The identifier check above already implies at least this
+  // many bytes, but check explicitly so the bound below cannot underflow.
+  constexpr size_t kMinBufferSize =
+      sizeof(flatbuffers::uoffset_t) + flatbuffers::kFileIdentifierLength;
+  ET_CHECK_OR_RETURN_ERROR(
+      flat_tensor_data->size() >= kMinBufferSize,
+      InvalidExternalData,
+      "FlatTensor data size %zu is too small (minimum %zu)",
+      flat_tensor_data->size(),
+      kMinBufferSize);
+  uint32_t root_offset =
+      flatbuffers::ReadScalar<flatbuffers::uoffset_t>(flat_tensor_data->data());
+  // The root table is at buf + root_offset. It must not point into the header
+  // and must leave room for at least a vtable offset (soffset_t) at its
+  // position.
+  ET_CHECK_OR_RETURN_ERROR(
+      root_offset >= kMinBufferSize &&
+          root_offset <=
+              flat_tensor_data->size() - sizeof(flatbuffers::soffset_t),
+      InvalidExternalData,
+      "FlatTensor root table offset %u is invalid for data size %zu",
+      root_offset,
+      flat_tensor_data->size());
+
   // Get pointer to root of flatbuffer table.
   const flat_tensor_flatbuffer::FlatTensor* flat_tensor =
       flat_tensor_flatbuffer::GetFlatTensor(flat_tensor_data->data());

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -266,3 +266,40 @@ TEST_F(FlatTensorDataMapTest, NewerSchemaVersionFailsToLoad) {
 
   EXPECT_EQ(data_map.error(), Error::InvalidExternalData);
 }
+
+TEST_F(FlatTensorDataMapTest, RejectsOutOfBoundsRootOffset) {
+  // A valid file loads.
+  Result<FreeableBuffer> valid = data_map_loader_->load(
+      0,
+      data_map_loader_->size().get(),
+      DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Constant));
+  ASSERT_EQ(valid.error(), Error::Ok);
+
+  // Copy it into a max-aligned buffer so we can corrupt the root table offset
+  // without tripping the earlier alignment check. std::vector<uint8_t> only
+  // guarantees 1-byte alignment, so over-allocate and offset to an aligned
+  // start. The first 4 bytes of a (non-size-prefixed) flatbuffer are the
+  // uoffset_t pointing at the root table; the next 4 are the file identifier.
+  constexpr size_t kAlignment = alignof(std::max_align_t);
+  const size_t size = valid->size();
+  std::unique_ptr<uint8_t[]> storage(new uint8_t[size + kAlignment]);
+  const size_t offset =
+      (kAlignment - (reinterpret_cast<uintptr_t>(storage.get()) % kAlignment)) %
+      kAlignment;
+  uint8_t* corrupt = storage.get() + offset;
+  std::memcpy(corrupt, valid->data(), size);
+
+  // Overwrite only the root offset with a value far past the end of the buffer,
+  // leaving the identifier and alignment intact so the corruption is caught by
+  // the root-offset bounds check rather than the identifier or size checks.
+  const uint32_t bad_offset = static_cast<uint32_t>(size) + 0x1000u;
+  std::memcpy(corrupt, &bad_offset, sizeof(bad_offset));
+
+  // Without the bounds check, GetFlatTensor() would return a pointer outside
+  // the buffer and the first field read would walk an out-of-bounds vtable
+  // (a use that ASan flags). With it, load() returns cleanly.
+  BufferDataLoader corrupt_loader(corrupt, size);
+  Result<FlatTensorDataMap> corrupt_map =
+      FlatTensorDataMap::load(&corrupt_loader);
+  ASSERT_EQ(corrupt_map.error(), Error::InvalidExternalData);
+}