flat_tensor check size (pytorch#14188)

Summary:

Add size check similar to D81938296, for flat tensor

Differential Revision: D82168471

Author: lucylq

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -73,7 +73,8 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
           segments->size());
       // Validate the segment.
       ET_CHECK_OR_RETURN_ERROR(
-          segments->Get(segment_index)->offset() < segment_end_offset,
+          (segments->Get(segment_index)->offset() +
+           segments->Get(segment_index)->size()) < segment_end_offset,
           InvalidExternalData,
           "Invalid segment offset %" PRIu64
           " is larger than the segment_base_offset + segment_data_size %" PRIu64
@@ -206,15 +207,21 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
   }
   Result<FlatTensorHeader> fh =
       FlatTensorHeader::Parse(header->data(), header->size());
-  if (fh.error() == Error::NotFound) {
-    // No header, throw error.
-    ET_LOG(Error, "No FlatTensorHeader found.");
-    return fh.error();
-  } else if (fh.error() != Error::Ok) {
-    // corruption, throw error.
-    ET_LOG(Error, "Flat tensor header may be corrupt.");
-    return fh.error();
-  }
+
+  ET_CHECK_OR_RETURN_ERROR(
+      fh.ok(),
+      InvalidExternalData,
+      "Failed to parse FlatTensor header with error code %u. File may be corrupt.",
+      static_cast<uint32_t>(fh.error()));
+
+  size_t expected_size = fh->segment_base_offset + fh->segment_data_size;
+  size_t actual_size = loader->size().get();
+  ET_CHECK_OR_RETURN_ERROR(
+      expected_size == actual_size,
+      InvalidExternalData,
+      "File size is too small; file may be corrupted or truncated. Expected %zu from flat_tensor header, received %zu from data loader",
+      expected_size,
+      actual_size);
 
   // Load flatbuffer data as a segment.
   Result<FreeableBuffer> flat_tensor_data = loader->load(

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/extension/data_loader/buffer_data_loader.h>
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
@@ -17,14 +18,15 @@
 #include <gtest/gtest.h>
 
 using namespace ::testing;
+using executorch::extension::BufferDataLoader;
+using executorch::extension::FileDataLoader;
 using executorch::extension::FlatTensorDataMap;
 using executorch::extension::FlatTensorHeader;
 using executorch::runtime::DataLoader;
 using executorch::runtime::Error;
 using executorch::runtime::FreeableBuffer;
 using executorch::runtime::Result;
 using executorch::runtime::TensorLayout;
-using torch::executor::util::FileDataLoader;
 
 class FlatTensorDataMapTest : public ::testing::Test {
  protected:
@@ -51,7 +53,7 @@ TEST_F(FlatTensorDataMapTest, LoadFlatTensorDataMap) {
   EXPECT_EQ(data_map.error(), Error::Ok);
 }
 
-TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetMetadata) {
+TEST_F(FlatTensorDataMapTest, GetMetadata) {
   Result<FlatTensorDataMap> data_map =
       FlatTensorDataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);
@@ -93,7 +95,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetMetadata) {
   EXPECT_EQ(const_c_res.error(), Error::NotFound);
 }
 
-TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetData) {
+TEST_F(FlatTensorDataMapTest, GetData) {
   Result<FlatTensorDataMap> data_map =
       FlatTensorDataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);
@@ -114,7 +116,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetData) {
   EXPECT_EQ(data_c_res.error(), Error::NotFound);
 }
 
-TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_Keys) {
+TEST_F(FlatTensorDataMapTest, GetKeys) {
   Result<FlatTensorDataMap> data_map =
       FlatTensorDataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);
@@ -138,7 +140,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_Keys) {
   EXPECT_EQ(key2_res.error(), Error::InvalidArgument);
 }
 
-TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_LoadInto) {
+TEST_F(FlatTensorDataMapTest, LoadInto) {
   Result<FlatTensorDataMap> data_map =
       FlatTensorDataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);
@@ -160,3 +162,23 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_LoadInto) {
   }
   free(data);
 }
+
+TEST_F(FlatTensorDataMapTest, LoadAndCheckSize) {
+  Result<FlatTensorDataMap> data_map =
+      FlatTensorDataMap::load(data_map_loader_.get());
+  EXPECT_EQ(data_map.error(), Error::Ok);
+
+  // Truncate the file.
+  size_t trunc_size = data_map_loader_->size().get() - 8;
+  Result<FreeableBuffer> truncated_file = data_map_loader_->load(
+      0,
+      trunc_size,
+      DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
+  ASSERT_EQ(truncated_file.error(), Error::Ok);
+
+  BufferDataLoader truncated_loader =
+      BufferDataLoader(truncated_file->data(), trunc_size);
+  Result<FlatTensorDataMap> truncated_program =
+      FlatTensorDataMap::load(&truncated_loader);
+  ASSERT_EQ(truncated_program.error(), Error::InvalidExternalData);
+}

extension/flat_tensor/test/targets.bzl

@@ -46,6 +46,7 @@ def define_common_targets(is_fbcode=False):
             ],
             deps = [
                 "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/extension/data_loader:buffer_data_loader",
                 "//executorch/extension/flat_tensor:flat_tensor_data_map",
                 "//executorch/runtime/core:named_data_map",
                 "//executorch/runtime/core/exec_aten:lib",