Fix allocation-size-too-big crash in prepare_input_tensors (#8233)

(Adapted from an LLM-suggested fix for a fuzzer-discovered crash)

The crash is an allocation-size-too-big error that occurs when the `prepare_input_tensors` function attempts to allocate an excessively large amount of memory for the `inputs` array. This crash is caused by the function's inability to handle large numbers of inputs, resulting in an attempt to allocate a huge amount of memory that exceeds the system's limits.

The root cause of the crash is the lack of bounds checking on the `num_inputs` variable, which allows the function to attempt to allocate an arbitrarily large amount of memory. This is exacerbated by the fact that the function allocates memory for each input tensor separately, without checking the total size of all tensors before allocating memory for the `inputs` array.

The patch fixes the crash by adding bounds checking on the `num_inputs` variable and calculating the total size of all tensors before allocating memory for the `inputs` array.

Differential Revision: D68876117

Author: dbort

extension/runner_util/inputs.cpp

@@ -22,15 +22,36 @@ using executorch::runtime::TensorInfo;
 namespace executorch {
 namespace extension {
 
-Result<BufferCleanup> prepare_input_tensors(Method& method) {
+Result<BufferCleanup> prepare_input_tensors(
+    Method& method,
+    PrepareInputTensorsOptions options) {
   MethodMeta method_meta = method.method_meta();
   size_t num_inputs = method_meta.num_inputs();
-  size_t num_allocated = 0;
+
+  // A large number of small allocations could exhaust the heap even if the
+  // total size is smaller than the limit.
+  ET_CHECK_OR_RETURN_ERROR(
+      num_inputs <= options.max_inputs,
+      InvalidProgram,
+      "Too many inputs: %zu > %zu",
+      num_inputs,
+      options.max_inputs);
+
+  // Allocate memory for the inputs array
   void** inputs = (void**)malloc(num_inputs * sizeof(void*));
+  ET_CHECK_OR_RETURN_ERROR(
+      inputs != nullptr,
+      MemoryAllocationFailed,
+      "malloc(%zd) failed",
+      num_inputs * sizeof(void*));
 
+  // Allocate memory for each input tensor.
+  size_t total_size = 0;
+  size_t num_allocated = 0;
   for (size_t i = 0; i < num_inputs; i++) {
     auto tag = method_meta.input_tag(i);
     if (!tag.ok()) {
+      // The BufferCleanup will free the inputs when it goes out of scope.
       BufferCleanup cleanup({inputs, num_allocated});
       return tag.error();
     }
@@ -40,10 +61,29 @@ Result<BufferCleanup> prepare_input_tensors(Method& method) {
     }
     Result<TensorInfo> tensor_meta = method_meta.input_tensor_meta(i);
     if (!tensor_meta.ok()) {
+      BufferCleanup cleanup({inputs, num_allocated});
       return tensor_meta.error();
     }
     // This input is a tensor. Allocate a buffer for it.
-    void* data_ptr = malloc(tensor_meta->nbytes());
+    size_t tensor_size = tensor_meta->nbytes();
+    total_size += tensor_size;
+    if (total_size > options.max_total_allocation_size) {
+      ET_LOG(
+          Error,
+          "Allocating %zu bytes for input %zu would exceed "
+          "max_total_allocation_size %zu",
+          tensor_size,
+          i,
+          options.max_total_allocation_size);
+      BufferCleanup cleanup({inputs, num_allocated});
+      return Error::InvalidProgram;
+    }
+    void* data_ptr = malloc(tensor_size);
+    if (data_ptr == nullptr) {
+      ET_LOG(Error, "malloc(%zu) failed for input %zu", tensor_size, i);
+      BufferCleanup cleanup({inputs, num_allocated});
+      return Error::MemoryAllocationFailed;
+    }
     inputs[num_allocated++] = data_ptr;
 
     // Create the tensor and set it as the input.
@@ -52,11 +92,11 @@ Result<BufferCleanup> prepare_input_tensors(Method& method) {
     if (err != Error::Ok) {
       ET_LOG(
           Error, "Failed to prepare input %zu: 0x%" PRIx32, i, (uint32_t)err);
-      // The BufferCleanup will free the inputs when it goes out of scope.
       BufferCleanup cleanup({inputs, num_allocated});
       return err;
     }
   }
+
   return BufferCleanup({inputs, num_allocated});
 }
 

extension/runner_util/inputs.h

@@ -51,18 +51,37 @@ class BufferCleanup final {
   executorch::runtime::Span<void*> buffers_;
 };
 
+/// Defines options for `prepare_input_tensors()`.
+struct PrepareInputTensorsOptions {
+  /**
+   * The maximum total size in bytes of all input tensors. If the total size of
+   * all inputs exceeds this, an error is returned. This prevents allocating too
+   * much memory if the PTE file is malformed.
+   */
+  size_t max_total_allocation_size = 1024 * 1024 * 1024;
+
+  /**
+   * The maximum number of inputs to allocate. If the number of inputs exceeds
+   * this, an error is returned. This prevents allocating too much memory if the
+   * PTE file is malformed.
+   */
+  size_t max_inputs = 1024;
+};
+
 /**
  * Allocates input tensors for the provided Method, filling them with ones. Does
  * not modify inputs that are not Tensors.
  *
  * @param[in] method The Method that owns the inputs to prepare.
+ * @param[in] options Extra options for preparing the inputs.
  *
  * @returns On success, an object that owns any allocated tensor memory. It must
  *     remain alive when calling `method->execute()`.
  * @returns An error on failure.
  */
 executorch::runtime::Result<BufferCleanup> prepare_input_tensors(
-    executorch::runtime::Method& method);
+    executorch::runtime::Method& method,
+    PrepareInputTensorsOptions options = {});
 
 namespace internal {
 /**

extension/runner_util/test/inputs_test.cpp

@@ -28,6 +28,7 @@ using executorch::runtime::EValue;
 using executorch::runtime::MemoryAllocator;
 using executorch::runtime::MemoryManager;
 using executorch::runtime::Method;
+using executorch::runtime::MethodMeta;
 using executorch::runtime::Program;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
@@ -100,6 +101,35 @@ TEST_F(InputsTest, Smoke) {
   // the pointers.
 }
 
+TEST_F(InputsTest, ExceedingInputCountLimitFails) {
+  // The smoke test above demonstrated that we can prepare inputs with the
+  // default limits. It should fail if we lower the max below the number of
+  // actual inputs.
+  MethodMeta method_meta = method_->method_meta();
+  size_t num_inputs = method_meta.num_inputs();
+  ASSERT_GE(num_inputs, 1);
+  executorch::extension::PrepareInputTensorsOptions options;
+  options.max_inputs = num_inputs - 1;
+
+  Result<BufferCleanup> input_buffers =
+      prepare_input_tensors(*method_, options);
+  ASSERT_NE(input_buffers.error(), Error::Ok);
+}
+
+TEST_F(InputsTest, ExceedingInputAllocationLimitFails) {
+  // The smoke test above demonstrated that we can prepare inputs with the
+  // default limits. It should fail if we lower the max below the actual
+  // allocation size.
+  executorch::extension::PrepareInputTensorsOptions options;
+  // The input tensors are float32, so 1 byte will always be smaller than any
+  // non-empty input tensor.
+  options.max_total_allocation_size = 1;
+
+  Result<BufferCleanup> input_buffers =
+      prepare_input_tensors(*method_, options);
+  ASSERT_NE(input_buffers.error(), Error::Ok);
+}
+
 TEST(BufferCleanupTest, Smoke) {
   // Returns the size of the buffer at index `i`.
   auto test_buffer_size = [](size_t i) {