From 2792f239c397c799b3a87d45e19c77ca0244194d Mon Sep 17 00:00:00 2001
From: jingyanwangms <47403504+jingyanwangms@users.noreply.github.com>
Date: Sun, 11 Aug 2024 03:28:41 -0700
Subject: [PATCH] Security Fuzz Test Fixes (#21608)

### Description
Fix address sanitizer and memory access Bug 1, 4, 5, 7, 8 found in
security fuzz test

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
---
 onnxruntime/core/framework/tensorprotoutils.cc             | 1 +
 onnxruntime/core/optimizer/unsqueeze_elimination.cc        | 4 ++++
 onnxruntime/core/providers/cpu/quantization/qlinearconv.cc | 4 ++--
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc
index cbd53298ab2ad..42f491825462c 100644
--- a/onnxruntime/core/framework/tensorprotoutils.cc
+++ b/onnxruntime/core/framework/tensorprotoutils.cc
@@ -1358,6 +1358,7 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n
 common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node,
                                               const std::filesystem::path& model_path,
                                               ONNX_NAMESPACE::TensorProto& tensor) {
+  ORT_ENFORCE(node.output_size() == 1, "NodeProto for Constant should have 1 output. Got:", node.output_size());
   return ConstantNodeProtoToTensorProto(node, model_path, tensor, node.output(0));
 }
 
diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.cc b/onnxruntime/core/optimizer/unsqueeze_elimination.cc
index 4efc8018f0217..d52cc82af02bb 100644
--- a/onnxruntime/core/optimizer/unsqueeze_elimination.cc
+++ b/onnxruntime/core/optimizer/unsqueeze_elimination.cc
@@ -40,6 +40,10 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect&
   // Generate new dims.
   InlinedVector<int64_t> new_dims(output_rank, 0);
   for (int64_t axis : axes) {
+    if (static_cast<size_t>(axis) >= new_dims.size()) {
+      LOGS(logger, WARNING) << "UnsqueezeElimination cannot remove node due to invalid axes" << node.Name();
+      return Status::OK();
+    }
     new_dims[static_cast<size_t>(axis)] = 1;
   }
 
diff --git a/onnxruntime/core/providers/cpu/quantization/qlinearconv.cc b/onnxruntime/core/providers/cpu/quantization/qlinearconv.cc
index 21a256eee6f14..7797cbe678bd4 100644
--- a/onnxruntime/core/providers/cpu/quantization/qlinearconv.cc
+++ b/onnxruntime/core/providers/cpu/quantization/qlinearconv.cc
@@ -380,8 +380,8 @@ Status QLinearConv<ActType>::PrePack(const Tensor& tensor, int input_idx, Alloca
   const int64_t M = shape[0];
   const int64_t C = shape[1];
 
-  // Verify that the total number of output channels is a multiple of the group count.
-  if (M % conv_attrs_.group != 0) {
+  // Verify that conv_attrs_.group is not 0 and the total number of output channels is a multiple of the group count.
+  if (conv_attrs_.group == 0 || M % conv_attrs_.group != 0) {
     return Status::OK();
   }