Fix integer overflow in RKNPU implicit bias allocation

onnxruntime/core/providers/rknpu/onnx_converter.cc

@@ -1,6 +1,7 @@
 // Copyright 2020 rock-chips.com Inc.
 
 #include <fstream>
+#include <limits>
 #include <map>
 #include <numeric>
 #include <string>
@@ -10,7 +11,9 @@
 #include <algorithm>
 #include <memory>
 #include <vector>
+#include "core/common/common.h"
 #include "core/common/logging/logging.h"
+#include "core/common/safeint.h"
 #include "onnx_converter.h"
 #include "node_attr_helper.h"
 
@@ -119,12 +122,45 @@ OnnxConverter::CreateRknnTensor(const std::string& name,
   return graph_->CreateTensor(attr, (void*)data);
 }
 
+static uint32_t ToRknpuDim(int64_t dim, const std::string& name) {
+  ORT_ENFORCE(dim >= 0 && dim <= static_cast<int64_t>(std::numeric_limits<uint32_t>::max()),
+              "RKNPU: tensor dimension out of uint32_t range (name=", name, ", dim=", dim, ")");
+
+  return static_cast<uint32_t>(dim);
+}
+
+// Allocates a zero-initialized buffer holding `count` elements of `element_size`
+// bytes each, used as an implicit (all-zero) bias when a Conv/Gemm node omits
+// its bias input. A malicious model can specify a weight dimension large enough
+// that `element_size * count` overflows size_t and wraps to a tiny allocation,
+// which would cause a heap buffer overflow when the buffer is later consumed.
+// SafeInt throws on overflow, and the allocation result is null-checked. Both
+// failure paths report element_size/count (and the byte count) to aid diagnosis.
+static void* AllocZeroedBias(size_t element_size, uint32_t count) {
+  size_t num_bytes = 0;
+  try {
+    num_bytes = SafeInt<size_t>(element_size) * count;
+  } catch (const std::exception&) {
+    ORT_THROW("RKNPU: implicit bias size overflow (element_size=", element_size,
+              ", count=", count, ")");
+  }
+  void* ptr = malloc(num_bytes);
+  ORT_ENFORCE(ptr != nullptr || num_bytes == 0,
+              "RKNPU: failed to allocate ", num_bytes,
+              " bytes for implicit bias (element_size=", element_size,
+              ", count=", count, ")");
+  if (ptr != nullptr) {
+    memset(ptr, 0, num_bytes);
+  }
+  return ptr;
+}
+
 void OnnxConverter::HandleInitializer() {
   for (const auto& tensor : model_proto_.graph().initializer()) {
     const std::string name = tensor.name();
     std::vector<uint32_t> dims;
     for (const auto dim : tensor.dims()) {
-      dims.push_back(static_cast<uint32_t>(dim));
+      dims.push_back(ToRknpuDim(dim, name));
     }
     if (tensor.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
       const char* ptr = tensor.float_data().empty()
@@ -186,7 +222,7 @@ std::vector<std::shared_ptr<rk::nn::Tensor>> OnnxConverter::GetInputOfOnnxModel(
     for (const auto& dim : input.type().tensor_type().shape().dim()) {
       if (dim.value_case() ==
           ONNX_NAMESPACE::TensorShapeProto_Dimension::kDimValue) {
-        shape.push_back(static_cast<uint32_t>(dim.dim_value()));
+        shape.push_back(ToRknpuDim(dim.dim_value(), input.name()));
       } else {
         throw std::invalid_argument(
             "The input of graph doesn't have dim_value");
@@ -267,7 +303,7 @@ Shaper::Shape GetShape(const ONNX_NAMESPACE::ModelProto& model_proto,
 
       for (const auto& dim : value_info.type().tensor_type().shape().dim()) {
         if (dim.has_dim_value()) {
-          shape.push_back(dim.dim_value());
+          shape.push_back(ToRknpuDim(dim.dim_value(), value_info.name()));
         } else {
           break;
         }
@@ -548,7 +584,7 @@ std::vector<std::vector<int>> OnnxConverter::GetSupportedNodes(
     const std::string name = tensor.name();
     std::vector<uint32_t> dims;
     for (const auto dim : tensor.dims()) {
-      dims.push_back(static_cast<uint32_t>(dim));
+      dims.push_back(ToRknpuDim(dim, name));
     }
     tensor_dims_[name] = dims;
   }
@@ -944,8 +980,7 @@ void OnnxConverter::AddLayerConvImpl(const std::string& input,
     }
   } else {
     uint32_t dim = shaper_[weight][0];
-    void* ptr = (void*)malloc(sizeof(float) * dim);
-    memset(ptr, 0, sizeof(float) * dim);
+    void* ptr = AllocZeroedBias(sizeof(float), dim);
     free_list_.push_back(ptr);
 
     std::vector<uint32_t> dims = {dim};
@@ -1053,8 +1088,7 @@ void OnnxConverter::AddLayerQLinearConvImpl(const string& input,
     }
   } else {
     uint32_t dim = shaper_[weight][0];
-    void* ptr = (void*)malloc(sizeof(int32_t) * dim);
-    memset(ptr, 0, sizeof(int32_t) * dim);
+    void* ptr = AllocZeroedBias(sizeof(int32_t), dim);
     free_list_.push_back(ptr);
 
     std::vector<uint32_t> dims = {dim};
@@ -1142,8 +1176,7 @@ void OnnxConverter::AddLayerDepthwiseConvImpl(
     }
   } else {
     uint32_t dim = shaper_[weight][0];
-    void* ptr = (void*)malloc(sizeof(float) * dim);
-    memset(ptr, 0, sizeof(float) * dim);
+    void* ptr = AllocZeroedBias(sizeof(float), dim);
     free_list_.push_back(ptr);
 
     std::vector<uint32_t> dims = {dim};
@@ -1376,8 +1409,7 @@ void OnnxConverter::AddLayerFC(const std::string& input,
     }
   } else {
     uint32_t dim = shaper_[weight][0];
-    void* ptr = (void*)malloc(sizeof(float) * dim);
-    memset(ptr, 0, sizeof(float) * dim);
+    void* ptr = AllocZeroedBias(sizeof(float), dim);
     free_list_.push_back(ptr);
 
     std::vector<uint32_t> dims = {dim};