fix: prevent integer overflow in memory size calculation for tensors (#16481)

chunseoklee · web-flow · commit 95fba2da1880 · 2026-04-21T09:56:40.000+09:00
* fix: prevent integer overflow in memory size calculation for tensors

This commit addresses potential integer overflow issues in memory size calculations by:
- Changing num_elements type from int32_t to int64_t
- Adding overflow checks before size calculations
- Using large_num_elements() instead of num_elements()
- Casting final size to size_t after overflow validation

The changes affect BuddyMemoryManager and several kernel implementations (ExpandDims, If, Reshape, TransposeConv) to ensure safe memory operations with large tensors.

ONE-DCO-1.0-Signed-off-by: Chunseok Lee &lt;chunseok.lee@samsung.com&gt;
diff --git a/compiler/luci-interpreter/src/BuddyMemoryManager.cpp b/compiler/luci-interpreter/src/BuddyMemoryManager.cpp
@@ -43,8 +43,16 @@ BuddyMemoryManager::BuddyMemoryManager(uint8_t *memory_start, int32_t memSize)
 void BuddyMemoryManager::allocate_memory(luci_interpreter::Tensor &tensor)
 {
   const size_t element_size = getDataTypeSize(tensor.element_type());
-  const int32_t num_elements = tensor.shape().num_elements();
-  auto size = num_elements * element_size;
+  const int64_t num_elements = tensor.shape().large_num_elements();
+
+  // Check for integer overflow in size calculation
+  if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+  {
+    throw std::runtime_error("Integer overflow in size calculation");
+  }
+
+  const int64_t total_size = num_elements * element_size;
+  auto size = static_cast<size_t>(total_size);
   auto footprint = size + sizeof(Block);
   auto l = (footprint & (footprint - 1)) == 0
              ? lowerLog2(footprint)
diff --git a/compiler/luci-interpreter/src/kernels/ExpandDims.cpp b/compiler/luci-interpreter/src/kernels/ExpandDims.cpp
@@ -80,8 +80,16 @@ void ExpandDims::execute() const
   auto *output_data = output()->data<void>();
 
   const size_t element_size = getDataTypeSize(input()->element_type());
-  const int32_t num_elements = input()->shape().num_elements();
-  std::memcpy(output_data, input_data, num_elements * element_size);
+  const int64_t num_elements = input()->shape().large_num_elements();
+
+  // Check for integer overflow in size calculation
+  if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+  {
+    throw std::runtime_error("Integer overflow in size calculation");
+  }
+
+  const int64_t total_size = num_elements * element_size;
+  std::memcpy(output_data, input_data, static_cast<size_t>(total_size));
 }
 
 } // namespace kernels
diff --git a/compiler/luci-interpreter/src/kernels/If.cpp b/compiler/luci-interpreter/src/kernels/If.cpp
@@ -83,10 +83,18 @@ void If::execute() const
     // TODO: Think about how allocate memory for output in main graph
     active_graph->configureAllocations(output(i));
 
-    const int32_t num_elements = output(i)->shape().num_elements();
+    const int64_t num_elements = output(i)->shape().large_num_elements();
     const std::size_t element_size = getDataTypeSize(output(i)->element_type());
+
+    // Check for integer overflow in size calculation
+    if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+    {
+      throw std::runtime_error("Integer overflow in size calculation");
+    }
+
+    const int64_t total_size = num_elements * element_size;
     std::memcpy(output(i)->data<void>(), graph_outputs[i]->data<void>(),
-                num_elements * element_size);
+                static_cast<size_t>(total_size));
   }
 }
 
diff --git a/compiler/luci-interpreter/src/kernels/Reshape.cpp b/compiler/luci-interpreter/src/kernels/Reshape.cpp
@@ -101,8 +101,16 @@ void Reshape::execute() const
   auto *output_data = output()->data<void>();
 
   const size_t element_size = getDataTypeSize(input()->element_type());
-  const int32_t num_elements = input()->shape().num_elements();
-  std::memcpy(output_data, input_data, num_elements * element_size);
+  const int64_t num_elements = input()->shape().large_num_elements();
+
+  // Check for integer overflow in size calculation
+  if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+  {
+    throw std::runtime_error("Integer overflow in size calculation");
+  }
+
+  const int64_t total_size = num_elements * element_size;
+  std::memcpy(output_data, input_data, static_cast<size_t>(total_size));
 }
 
 } // namespace kernels
diff --git a/compiler/luci-interpreter/src/kernels/TransposeConv.cpp b/compiler/luci-interpreter/src/kernels/TransposeConv.cpp
@@ -296,7 +296,17 @@ void TransposeConv::evalQuantizedS16() const
   int32_t activation_max{};
   calculateActivationRangeQuantized(Activation::NONE, output(), &activation_min, &activation_max);
 
-  std::memset(scratch_data, 0, scratch_tensor->shape().num_elements() * sizeof(int64_t));
+  const int64_t num_elements = scratch_tensor->shape().large_num_elements();
+  const size_t element_size = sizeof(int64_t);
+
+  // Check for integer overflow in size calculation
+  if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+  {
+    throw std::runtime_error("Integer overflow in size calculation");
+  }
+
+  const int64_t total_size = num_elements * element_size;
+  std::memset(scratch_data, 0, static_cast<size_t>(total_size));
 
   BroadcastableWrapper<ChannelQuantMultipliers> output_multipliers(_quant_multipliers);
   for (int32_t batch = 0; batch < batches; ++batch)
diff --git a/compiler/luci-interpreter/src/kernels/While.cpp b/compiler/luci-interpreter/src/kernels/While.cpp
@@ -35,9 +35,17 @@ void copy(const std::vector<const Tensor *> &src, const std::vector<Tensor *> &d
     LUCI_INTERPRETER_CHECK(dst[i]->element_type() == src[i]->element_type());
     dst[i]->resize(src[i]->shape());
 
-    const int32_t num_elements = src[i]->shape().num_elements();
+    const int64_t num_elements = src[i]->shape().large_num_elements();
     const std::size_t element_size = getDataTypeSize(src[i]->element_type());
-    std::memcpy(dst[i]->data<void>(), src[i]->data<void>(), num_elements * element_size);
+
+    // Check for integer overflow in size calculation
+    if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+    {
+      throw std::runtime_error("Integer overflow in size calculation");
+    }
+
+    const int64_t total_size = num_elements * element_size;
+    std::memcpy(dst[i]->data<void>(), src[i]->data<void>(), static_cast<size_t>(total_size));
   }
 }
 
diff --git a/compiler/luci-interpreter/src/loader/GraphLoader.cpp b/compiler/luci-interpreter/src/loader/GraphLoader.cpp
@@ -39,9 +39,16 @@ template <typename NodeT> Shape getNodeShape(const NodeT *node)
 template <DataType DT> const void *getNodeDataImpl(const luci::CircleConst *node, size_t *data_size)
 {
   const size_t element_size = getDataTypeSize(DT);
-  const int32_t num_elements = node->size<DT>();
+  const int64_t num_elements = node->size<DT>(); // Assuming size<DT>() uses large_num_elements()
 
-  *data_size = num_elements * element_size;
+  // Check for integer overflow in size calculation
+  if (num_elements < 0 || static_cast<uint64_t>(num_elements) > SIZE_MAX / element_size)
+  {
+    throw std::runtime_error("Integer overflow in size calculation");
+  }
+
+  const int64_t total_size = num_elements * element_size;
+  *data_size = static_cast<size_t>(total_size);
   if (*data_size > 0)
   {
     // FIXME There is no good way to get the pointer to the data currently.
diff --git a/compiler/luci/import/src/Nodes/CircleConst.cpp b/compiler/luci/import/src/Nodes/CircleConst.cpp
@@ -96,6 +96,23 @@ void copy_data<loco::DataType::STRING>(const VectorWrapper<uint8_t> &raw_data,
   }
   assert(offsets.size() == num_elements + 1);
 
+  // Validate STRING offsets as non-negative, monotonic, and bounded within data buffer
+  for (uint32_t i = 0; i < offsets.size(); ++i)
+  {
+    if (offsets[i] < 0)
+    {
+      throw std::runtime_error("String offset is negative");
+    }
+    if (i > 0 && offsets[i] < offsets[i - 1])
+    {
+      throw std::runtime_error("String offsets are not monotonic");
+    }
+    if (offsets[i] > static_cast<int32_t>(raw_data.size()))
+    {
+      throw std::runtime_error("String offset is out of bounds");
+    }
+  }
+
   const_node->size<loco::DataType::STRING>(num_elements);
   for (uint32_t i = 0; i < num_elements; ++i)
   {
diff --git a/runtime/onert/core/src/loader/BaseLoader.h b/runtime/onert/core/src/loader/BaseLoader.h
@@ -113,6 +113,13 @@ template <typename LoaderDomain> class BaseLoader
   // Get BuiltinOperator
   BuiltinOperator getBuiltinOperator(const Operator *op)
   {
+    // Enforce explicit bounds validation for opcode_index before every operator-code lookup
+    if (op->opcode_index() < 0 ||
+        static_cast<size_t>(op->opcode_index()) >= _domain_model->operator_codes()->size())
+    {
+      throw std::runtime_error("Invalid opcode_index: " + std::to_string(op->opcode_index()));
+    }
+
     auto const builtin_opcode = _domain_model->operator_codes()->Get(op->opcode_index());
     auto builtin_op = builtin_opcode->builtin_code();
     if (builtin_op < BuiltinOperator::BuiltinOperator_PLACEHOLDER_FOR_GREATER_OP_CODES)
