Verify pte size at runtime

runtime/executor/program.cpp

@@ -67,6 +67,7 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
   // See if the program size is in the header.
   size_t program_size = 0;
   size_t segment_base_offset = 0;
+  size_t segment_data_size = 0;
   {
     EXECUTORCH_SCOPE_PROF("Program::check_header");
     Result<FreeableBuffer> header = loader->load(
@@ -82,6 +83,24 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
       // The header has the program size.
       program_size = eh->program_size;
       segment_base_offset = eh->segment_base_offset;
+      segment_data_size = eh->segment_data_size;
+
+      // segment_data_size was added in ET 1.0 release. For BC, only check the
+      // expected file size when there are no segments or when segment_data_size
+      // is positive (0-value may indicate no segments)
+      if ((segment_data_size == 0 && segment_base_offset == 0) ||
+          segment_data_size > 0) {
+        size_t expected = segment_base_offset == 0
+            ? program_size
+            : segment_base_offset + segment_data_size;
+        size_t actual = loader->size().get();
+        ET_CHECK_OR_RETURN_ERROR(
+            expected <= actual,
+            InvalidProgram,
+            "File size is too small. Expected file size from extended header is %zu, actual file size from data loader is %zu",
+            expected,
+            actual);
+      }
     } else if (eh.error() == Error::NotFound) {
       // No header; the program consumes the whole file, and there are no
       // segments.

runtime/executor/test/program_test.cpp

@@ -574,3 +574,22 @@ TEST_F(ProgramTest, LoadFromMutableSegment) {
       &program.get(), 500, 1, 1, buffer);
   EXPECT_NE(err, Error::Ok);
 }
+
+TEST_F(ProgramTest, LoadAndCheckPTESize) {
+  // Load the serialized ModuleAddMul data, with constants in the segment.
+  const char* linear_path = std::getenv("ET_MODULE_ADD_MUL_PATH");
+  Result<FileDataLoader> linear_loader = FileDataLoader::from(linear_path);
+  ASSERT_EQ(linear_loader.error(), Error::Ok);
+  Result<Program> program = Program::load(&linear_loader.get());
+  ASSERT_EQ(program.error(), Error::Ok);
+
+  // Create a truncated file.
+  Result<FreeableBuffer> truncated_file = linear_loader->load(
+      0, 200, DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
+  ASSERT_EQ(truncated_file.error(), Error::Ok);
+
+  BufferDataLoader truncated_loader =
+      BufferDataLoader(truncated_file->data(), 200);
+  Result<Program> truncated_program = Program::load(&truncated_loader);
+  ASSERT_EQ(truncated_program.error(), Error::InvalidProgram);
+}