Fix stack buffer overflow when XNNPACK tensor has too many dims (#3233)

Summary:
Pull Request resolved: #3233

Noticed this overflow when I was looking through the XNNPACK backend.

I am not very familiar with executorch or XNNPACK, so please be critical in review! In particular, my test is crashing in xnn_delete_runtime and I don't know why; it was previously getting a UBSAN failure for the stack buffer overflow before I patched XNNExecutor.cpp

Reviewed By: digantdesai, mcr229

Differential Revision: D56450593

fbshipit-source-id: d0f2b4d5dc44e08f9679560d4a6adde7f1559173

Author: swolchok

backends/xnnpack/runtime/XNNExecutor.cpp

@@ -87,6 +87,12 @@ __ET_NODISCARD Error XNNExecutor::prepare_args(EValue** args) {
     if (i < input_ids_.size()) {
       size_t num_dims = tensor->dim();
       size_t dims[XNN_MAX_TENSOR_DIMS];
+      ET_CHECK_OR_RETURN_ERROR(
+          num_dims <= XNN_MAX_TENSOR_DIMS,
+          InvalidArgument,
+          "XNNPACK backend accepts tensors with at most %d dims, but got %zu",
+          XNN_MAX_TENSOR_DIMS,
+          num_dims);
       for (int d = 0; d < num_dims; ++d) {
         dims[d] = tensor->size(d);
       }

backends/xnnpack/targets.bzl

@@ -38,9 +38,11 @@ def define_common_targets():
         preprocessor_flags = [
             # "-DENABLE_XNNPACK_PROFILING",
         ],
+        exported_deps = [
+            "//executorch/runtime/backend:interface",
+        ],
         deps = [
             third_party_dep("XNNPACK"),
-            "//executorch/runtime/backend:interface",
             "//executorch/backends/xnnpack/serialization:xnnpack_flatbuffer_header",
             "//executorch/backends/xnnpack/threadpool:threadpool",
             "//executorch/runtime/core/exec_aten/util:tensor_util",

backends/xnnpack/test/runtime/test_xnnexecutor.cpp

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/xnnpack/runtime/XNNExecutor.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <gtest/gtest.h>
+#include <xnnpack/subgraph.h>
+
+using torch::executor::Error;
+using torch::executor::EValue;
+using torch::executor::testing::TensorFactory;
+using torch::executor::xnnpack::delegate::XNNExecutor;
+
+TEST(XNNExecutorTest, ArgumentWithTooManyDimensions) {
+  XNNExecutor executor;
+  xnn_subgraph_t subgraph = nullptr;
+  xnn_runtime_t rt = nullptr;
+  et_pal_init();
+  ASSERT_EQ(xnn_initialize(nullptr), xnn_status_success);
+  ASSERT_EQ(xnn_create_subgraph(2, 0, &subgraph), xnn_status_success);
+  std::unique_ptr<xnn_subgraph, decltype(&xnn_delete_subgraph)> auto_subgraph(
+      subgraph, xnn_delete_subgraph);
+
+  auto input_id = XNN_INVALID_NODE_ID;
+  std::vector<size_t> dims = {
+      1,
+  };
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_quantized_tensor_value(
+          subgraph,
+          xnn_datatype_qint8,
+          0,
+          1,
+          dims.size(),
+          dims.data(),
+          nullptr,
+          /*external_id=*/0,
+          /*flags=*/XNN_VALUE_FLAG_EXTERNAL_INPUT,
+          &input_id));
+  ASSERT_NE(input_id, XNN_INVALID_NODE_ID);
+
+  auto output_id = XNN_INVALID_NODE_ID;
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_quantized_tensor_value(
+          subgraph,
+          xnn_datatype_qint8,
+          0,
+          1,
+          dims.size(),
+          dims.data(),
+          nullptr,
+          /*external_id=*/0,
+          /*flags=*/XNN_VALUE_FLAG_EXTERNAL_OUTPUT,
+          &output_id));
+  ASSERT_NE(output_id, XNN_INVALID_NODE_ID);
+
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_clamp(subgraph, 1, 2, input_id, output_id, 0));
+
+  ASSERT_EQ(xnn_create_runtime(subgraph, &rt), xnn_status_success);
+  EXPECT_EQ(
+      executor.initialize(
+          rt,
+          {
+              0,
+          },
+          {
+              1,
+          }),
+      Error::Ok);
+  TensorFactory<exec_aten::ScalarType::Int> tf;
+  auto input_tensor = tf.make({1, 1, 1, 1, 1, 1, 1, 1, 1}, {42});
+  ASSERT_EQ(input_tensor.dim(), 9);
+  auto output_tensor = tf.make(
+      {
+          1,
+      },
+      {
+          1,
+      });
+  EValue input_ev(input_tensor);
+  EValue output_ev(output_tensor);
+  std::array<EValue*, 2> args = {&input_ev, &output_ev};
+  // Check for invalid number of dimensions should fail without stack overflow.
+  EXPECT_EQ(executor.prepare_args(args.data()), Error::InvalidArgument);
+}

backends/xnnpack/test/targets.bzl

@@ -1,3 +1,4 @@
+load("@fbsource//xplat/executorch/backends/xnnpack/third-party:third_party_libs.bzl", "third_party_dep")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 def define_common_targets():
@@ -17,3 +18,14 @@ def define_common_targets():
             "//executorch/backends/xnnpack:dynamic_quant_utils",
         ],
     )
+
+    runtime.cxx_test(
+        name = "xnnexecutor_test",
+        srcs = ["runtime/test_xnnexecutor.cpp"],
+        deps = [
+            third_party_dep("XNNPACK"),
+            "//executorch/runtime/core/exec_aten/testing_util:tensor_util",
+            "//executorch/runtime/core/exec_aten/util:scalar_type_util",
+            "//executorch/backends/xnnpack:xnnpack_backend",
+        ],
+    )