From 09cfff9860801b295ae74b62e499a5aa78c3732b Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Wed, 10 Sep 2025 16:22:03 -0700
Subject: [PATCH] Verify pte size at runtime

Pull Request resolved: https://github.com/pytorch/executorch/pull/14092

Check that the file size is at least the size of what we expect from the extended header.

ie. header_offset + program_size + segment_data_size.
ghstack-source-id: 308927511

Differential Revision: [D81938296](https://our.internmc.facebook.com/intern/diff/D81938296/)
---
 runtime/executor/program.cpp           | 19 +++++++++++++++++++
 runtime/executor/test/program_test.cpp | 19 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
index e58c8a96aa7..344e3c7177a 100644
--- a/runtime/executor/program.cpp
+++ b/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.
diff --git a/runtime/executor/test/program_test.cpp b/runtime/executor/test/program_test.cpp
index 962bf8f548a..3afb71b3565 100644
--- a/runtime/executor/test/program_test.cpp
+++ b/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);
+}