Check size before memcpy

Author: lucylq

extension/runner_util/inputs.cpp

@@ -86,17 +86,47 @@ Result<BufferCleanup> prepare_input_tensors(
             Debug, "Verifying and setting input for non-tensor input %zu", i);
 
         if (tag.get() == Tag::Int) {
-          int64_t int_input;
-          std::memcpy(&int_input, buffer, buffer_size);
-          err = method.set_input(runtime::EValue(int_input), i);
+          if (buffer_size != sizeof(int64_t)) {
+            ET_LOG(
+                Error,
+                "Int input at index %zu has size %zu, expected sizeof(int64_t) %zu",
+                i,
+                buffer_size,
+                sizeof(int64_t));
+            err = Error::InvalidArgument;
+          } else {
+            int64_t int_input;
+            std::memcpy(&int_input, buffer, buffer_size);
+            err = method.set_input(runtime::EValue(int_input), i);
+          }
         } else if (tag.get() == Tag::Double) {
-          double double_input;
-          std::memcpy(&double_input, buffer, buffer_size);
-          err = method.set_input(runtime::EValue(double_input), i);
+          if (buffer_size != sizeof(double)) {
+            ET_LOG(
+                Error,
+                "Double input at index %zu has size %zu, expected sizeof(double) %zu",
+                i,
+                buffer_size,
+                sizeof(double));
+            err = Error::InvalidArgument;
+          } else {
+            double double_input;
+            std::memcpy(&double_input, buffer, buffer_size);
+            err = method.set_input(runtime::EValue(double_input), i);
+          }
         } else if (tag.get() == Tag::Bool) {
           bool bool_input;
-          std::memcpy(&bool_input, buffer, buffer_size);
-          err = method.set_input(runtime::EValue(bool_input), i);
+          if (buffer_size != sizeof(bool)) {
+            ET_LOG(
+                Error,
+                "Bool input at index %zu has size %zu, expected sizeof(bool) %zu",
+                i,
+                buffer_size,
+                sizeof(bool));
+            err = Error::InvalidArgument;
+          } else {
+            std::memcpy(&bool_input, buffer, buffer_size);
+            err = method.set_input(runtime::EValue(bool_input), i);
+          }
         } else {
           ET_LOG(
               Error,

extension/runner_util/test/inputs_test.cpp

@@ -186,3 +186,42 @@ TEST(BufferCleanupTest, Smoke) {
   // complaint.
   bc2.reset();
 }
+
+TEST_F(InputsTest, DoubleInputWrongSizeFails) {
+  MethodMeta method_meta = method_->method_meta();
+
+  // ModuleAdd has 3 inputs: tensor, tensor, double (alpha)
+  ASSERT_EQ(method_meta.num_inputs(), 3);
+
+  // Verify input 2 is a Double
+  auto tag = method_meta.input_tag(2);
+  ASSERT_TRUE(tag.ok());
+  ASSERT_EQ(tag.get(), Tag::Double);
+
+  // Create input_buffers with wrong size for the Double input
+  std::vector<std::pair<char*, size_t>> input_buffers;
+
+  // Allocate correct buffers for tensors (inputs 0 and 1)
+  auto tensor0_meta = method_meta.input_tensor_meta(0);
+  auto tensor1_meta = method_meta.input_tensor_meta(1);
+  ASSERT_TRUE(tensor0_meta.ok());
+  ASSERT_TRUE(tensor1_meta.ok());
+
+  std::vector<char> buf0(tensor0_meta->nbytes());
+  std::vector<char> buf1(tensor1_meta->nbytes());
+
+  // ModuleAdd expects alpha=1.0. Need to set this correctly, otherwise
+  // set_input fails validation before the buffer overflow happens.
+  double alpha = 1.0;
+  // Double is size 8; use a larger buffer to invoke overflow.
+  char large_buffer[16];
+  memcpy(large_buffer, &alpha, sizeof(double));
+
+  input_buffers.push_back({buf0.data(), buf0.size()});
+  input_buffers.push_back({buf1.data(), buf1.size()});
+  input_buffers.push_back({large_buffer, sizeof(large_buffer)});
+
+  Result<BufferCleanup> result =
+      prepare_input_tensors(*method_, {}, input_buffers);
+  EXPECT_EQ(result.error(), Error::InvalidArgument);
+}