Fix allocation-size-too-big crash in prepare_input_tensors

Summary:
(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,37 @@ using executorch::runtime::TensorInfo;
 namespace executorch {
 namespace extension {
 
-Result<BufferCleanup> prepare_input_tensors(Method& method) {
+Result<BufferCleanup> prepare_input_tensors(
+    Method& method,
+    size_t max_total_allocation_size,
+    size_t max_inputs) {
   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 <= max_inputs,
+      InvalidProgram,
+      "Too many inputs: %zu > %zu",
+      num_inputs,
+      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 +62,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 > max_total_allocation_size) {
+      ET_LOG(
+          Error,
+          "Allocating %zu bytes for input %zu would exceed "
+          "max_total_allocation_size %zu",
+          tensor_size,
+          i,
+          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 +93,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

@@ -56,13 +56,22 @@ class BufferCleanup final {
  * not modify inputs that are not Tensors.
  *
  * @param[in] method The Method that owns the inputs to prepare.
+ * @param[in] max_total_allocation_size 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.
+ * @param[in] max_inputs 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.
  *
  * @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,
+    size_t max_total_allocation_size = 1024 * 1024 * 1024,
+    size_t max_inputs = 1024);
 
 namespace internal {
 /**