From 12b768c05d486c8b10fe5a7ffe62aee89492ca59 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 8 Sep 2025 14:54:59 -0700
Subject: [PATCH] Verify pte size at runtime

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.

Differential Revision: [D81938296](https://our.internmc.facebook.com/intern/diff/D81938296/)

[ghstack-poisoned]
---
 runtime/executor/program.cpp           | 20 ++++++++++++++++++++
 runtime/executor/test/program_test.cpp | 19 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
index e58c8a96aa7..cbf8cfa1f6e 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,25 @@ 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;
+
+      // Check the expected file size in two cases:
+      // 1. segment_base_offset == 0 && segment_data_size == 0: this indicates
+      // there are no segments.
+      // 2. segment_data_size > 0; a 0 value may indicate that the field was
+      // not set, e.g. on older PTE files.
+      if ((segment_data_size == 0 && segment_base_offset == 0) ||
+          segment_data_size > 0) {
+        size_t expected =
+            ExtendedHeader::kHeaderOffset + program_size + 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..d686cd497b2 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);
+
+  Result<BufferDataLoader> truncated_loader =
+      BufferDataLoader(truncated_file->data(), 200);
+  Result<Program> truncated_program = Program::load(&truncated_loader.get());
+  ASSERT_EQ(truncated_program.error(), Error::InvalidProgram);
+}