!9227 verify shape and index rationality to icsl

From: @lyvette
Reviewed-by: @hangangqiang
Signed-off-by: @hangangqiang

mindspore/lite/src/inner_context.cc

@@ -64,6 +64,10 @@ int InnerContext::IsValid() {
     MS_LOG(ERROR) << "Device list is empty.";
     return RET_NOT_SUPPORT;
   }
+  if (!IsCpuEnabled()) {
+    MS_LOG(ERROR) << "CPU is not supported.";
+    return RET_NOT_SUPPORT;
+  }
 #ifndef SUPPORT_GPU
   if (IsGpuEnabled()) {
     MS_LOG(ERROR) << "GPU is not supported.";

mindspore/lite/src/lite_kernel.cc

@@ -94,6 +94,11 @@ int LiteKernel::PreProcess() {
   for (auto *output : outputs) {
     MS_ASSERT(output != nullptr);
 
+    if (output->ElementsNum() >= MAX_MALLOC_SIZE / static_cast<int>(sizeof(int64_t))) {
+      MS_LOG(ERROR) << "The size of output tensor is too big";
+      return RET_ERROR;
+    }
+
     auto ret = output->MallocData();
     if (ret != RET_OK) {
       MS_LOG(ERROR) << "MallocData failed";

mindspore/lite/src/model_common.cc

@@ -15,16 +15,15 @@
  */
 #include "src/model_common.h"
 #include "include/version.h"
+#include "src/ops/while.h"
 #ifndef PRIMITIVE_WRITEABLE
 #include "src/ops/ops_register.h"
 #endif
 
 namespace mindspore::lite {
 bool ConvertNodes(const schema::MetaGraph *meta_graph, Model *model) {
-  MS_ASSERT(model != nullptr);
-  MS_ASSERT(meta_graph != nullptr);
-  if (meta_graph->nodes() == nullptr) {
-    MS_LOG(ERROR) << "meta_graph is invalid, please check your model file.";
+  if (model == nullptr || meta_graph == nullptr || meta_graph->nodes() == nullptr) {
+    MS_LOG(ERROR) << "model or meta_graph is invalid, please check your model file.";
     return false;
   }
   for (size_t i = 0; i < meta_graph->nodes()->size(); ++i) {
@@ -34,9 +33,13 @@ bool ConvertNodes(const schema::MetaGraph *meta_graph, Model *model) {
       return false;
     }
     auto c_node = meta_graph->nodes()->GetAs<schema::CNode>(i);
-    MS_ASSERT(c_node != nullptr);
+    if (c_node == nullptr || c_node->primitive() == nullptr || c_node->name() == nullptr ||
+        c_node->inputIndex() == nullptr) {
+      MS_LOG(ERROR) << "c_node is invalid.";
+      delete node;
+      return false;
+    }
     auto src_prim = c_node->primitive();
-    MS_ASSERT(src_prim != nullptr);
 #ifdef PRIMITIVE_WRITEABLE
     node->primitive_ = PrimitiveC::Create(const_cast<schema::Primitive *>(src_prim));
 #else
@@ -49,10 +52,8 @@ bool ConvertNodes(const schema::MetaGraph *meta_graph, Model *model) {
       return false;
     }
     node->primitive_->set_quant_type(c_node->quantType());
-    MS_ASSERT(c_node->name() != nullptr);
     node->name_ = c_node->name()->c_str();
     node->node_type_ = c_node->nodeType();
-    MS_ASSERT(c_node->inputIndex() != nullptr);
     auto count = c_node->inputIndex()->size();
     for (uint32_t j = 0; j < count; ++j) {
       node->input_indices_.push_back(c_node->inputIndex()->Get(j));
@@ -69,17 +70,15 @@ bool ConvertNodes(const schema::MetaGraph *meta_graph, Model *model) {
 }
 
 bool ConvertTensors(const schema::MetaGraph *meta_graph, Model *model) {
-  MS_ASSERT(model != nullptr);
-  MS_ASSERT(meta_graph != nullptr);
-  if (meta_graph->allTensors() == nullptr) {
-    MS_LOG(ERROR) << "meta_graph is invalid, please check your model file.";
+  if (model == nullptr || meta_graph == nullptr || meta_graph->allTensors() == nullptr) {
+    MS_LOG(ERROR) << "model or meta_graph is invalid, please check your model file.";
     return false;
   }
   auto tensor_count = meta_graph->allTensors()->size();
   for (uint32_t i = 0; i < tensor_count; ++i) {
     auto *tensor = meta_graph->allTensors()->GetAs<schema::Tensor>(i);
     if (tensor == nullptr) {
-      MS_LOG(ERROR) << i << "th tensor in model is nullptr";
+      MS_LOG(ERROR) << i << "the tensor in model is nullptr";
       return false;
     }
     model->all_tensors_.push_back(const_cast<mindspore::schema::Tensor *>(tensor));
@@ -88,31 +87,33 @@ bool ConvertTensors(const schema::MetaGraph *meta_graph, Model *model) {
 }
 
 int ConvertSubGraph(const schema::SubGraph *sub_graph, Model *model) {
-  MS_ASSERT(model != nullptr);
-  MS_ASSERT(sub_graph != nullptr);
+  if (model == nullptr || sub_graph == nullptr) {
+    MS_LOG(ERROR) << "model or sub_graph is null.";
+    return RET_ERROR;
+  }
+  if (sub_graph->name() == nullptr || sub_graph->inputIndices() == nullptr || sub_graph->outputIndices() == nullptr ||
+      sub_graph->nodeIndices() == nullptr || sub_graph->tensorIndices() == nullptr) {
+    MS_LOG(ERROR) << "sub_graph is invalid.";
+    return RET_ERROR;
+  }
   auto *sub_graph_temp = new (std::nothrow) Model::SubGraph();
   if (sub_graph_temp == nullptr) {
     MS_LOG(ERROR) << "new subGraph fail!";
     return RET_ERROR;
   }
-  MS_ASSERT(sub_graph->name() != nullptr);
   sub_graph_temp->name_ = sub_graph->name()->c_str();
-  MS_ASSERT(sub_graph->inputIndices() != nullptr);
   auto in_count = sub_graph->inputIndices()->size();
   for (uint32_t i = 0; i < in_count; ++i) {
     sub_graph_temp->input_indices_.push_back(sub_graph->inputIndices()->Get(i));
   }
-  MS_ASSERT(sub_graph->outputIndices() != nullptr);
   auto out_count = sub_graph->outputIndices()->size();
   for (uint32_t i = 0; i < out_count; ++i) {
     sub_graph_temp->output_indices_.push_back(sub_graph->outputIndices()->Get(i));
   }
-  MS_ASSERT(sub_graph->nodeIndices() != nullptr);
   auto node_count = sub_graph->nodeIndices()->size();
   for (uint32_t i = 0; i < node_count; ++i) {
     sub_graph_temp->node_indices_.push_back(sub_graph->nodeIndices()->Get(i));
   }
-  MS_ASSERT(sub_graph->tensorIndices() != nullptr);
   auto tensor_count = sub_graph->tensorIndices()->size();
   for (uint32_t i = 0; i < tensor_count; ++i) {
     sub_graph_temp->tensor_indices_.push_back(sub_graph->tensorIndices()->Get(i));
@@ -122,8 +123,10 @@ int ConvertSubGraph(const schema::SubGraph *sub_graph, Model *model) {
 }
 
 int MetaGraphMappingSubGraph(const mindspore::schema::MetaGraph *meta_graph, Model *model) {
-  MS_ASSERT(model != nullptr);
-  MS_ASSERT(meta_graph != nullptr);
+  if (model == nullptr || meta_graph == nullptr) {
+    MS_LOG(ERROR) << "model or meta_graph is null.";
+    return RET_ERROR;
+  }
   if (meta_graph->inputIndex() == nullptr || meta_graph->outputIndex() == nullptr || meta_graph->nodes() == nullptr ||
       meta_graph->allTensors() == nullptr) {
     MS_LOG(ERROR) << "meta_graph is invalid, please check your model file.";
@@ -137,22 +140,18 @@ int MetaGraphMappingSubGraph(const mindspore::schema::MetaGraph *meta_graph, Mod
   if (meta_graph->name() != nullptr) {
     sub_graph_temp->name_ = meta_graph->name()->c_str();
   }
-  MS_ASSERT(meta_graph->inputIndex() != nullptr);
   auto in_count = meta_graph->inputIndex()->size();
   for (uint32_t i = 0; i < in_count; ++i) {
     sub_graph_temp->input_indices_.push_back(meta_graph->inputIndex()->Get(i));
   }
-  MS_ASSERT(meta_graph->outputIndex() != nullptr);
   auto out_count = meta_graph->outputIndex()->size();
   for (uint32_t i = 0; i < out_count; ++i) {
     sub_graph_temp->output_indices_.push_back(meta_graph->outputIndex()->Get(i));
   }
-  MS_ASSERT(meta_graph->nodes() != nullptr);
   auto node_count = meta_graph->nodes()->size();
   for (uint32_t i = 0; i < node_count; ++i) {
     sub_graph_temp->node_indices_.push_back(i);
   }
-  MS_ASSERT(meta_graph->allTensors() != nullptr);
   auto tensor_count = meta_graph->allTensors()->size();
   for (uint32_t i = 0; i < tensor_count; ++i) {
     sub_graph_temp->tensor_indices_.push_back(i);
@@ -161,6 +160,76 @@ int MetaGraphMappingSubGraph(const mindspore::schema::MetaGraph *meta_graph, Mod
   return RET_OK;
 }
 
+int NodeVerify(const Model &model) {
+  auto tensor_size = model.all_tensors_.size();
+  uint32_t subGraph_size = model.sub_graphs_.size();
+
+  for (auto &node : model.all_nodes_) {
+    if (node == nullptr || node->primitive_ == nullptr) {
+      MS_LOG(ERROR) << "node or its primitive_ is null.";
+      return RET_ERROR;
+    }
+    if (std::any_of(node->input_indices_.begin(), node->input_indices_.end(),
+                    [&tensor_size](const uint32_t &idx) { return idx >= tensor_size; })) {
+      MS_LOG(ERROR) << "Index of node->input_indices_ is beyond size.";
+      return RET_ERROR;
+    }
+    if (std::any_of(node->output_indices_.begin(), node->output_indices_.end(),
+                    [&tensor_size](const uint32_t &idx) { return idx >= tensor_size; })) {
+      MS_LOG(ERROR) << "Index of node->output_indices_ is beyond size.";
+      return RET_ERROR;
+    }
+
+    auto prim = node->primitive_;
+    if (prim->Type() == schema::PrimitiveType_While) {
+      auto whileOp = reinterpret_cast<mindspore::lite::While *>(const_cast<mindspore::lite::PrimitiveC *>(prim));
+      if (whileOp == nullptr) {
+        MS_LOG(ERROR) << "whileOp is null.";
+        return RET_ERROR;
+      }
+      if (static_cast<uint32_t>(whileOp->GetBodySubgraphIndex()) >= subGraph_size ||
+          static_cast<uint32_t>(whileOp->GetCondSubgraphIndex()) >= subGraph_size) {
+        MS_LOG(ERROR) << "index of subGraph is beyond subGraph_size.";
+        return RET_ERROR;
+      }
+    }
+  }
+  return RET_OK;
+}
+
+int SubGraphVerify(const Model &model) {
+  auto tensor_size = model.all_tensors_.size();
+  auto node_size = model.all_nodes_.size();
+
+  for (auto &graph : model.sub_graphs_) {
+    if (graph == nullptr) {
+      MS_LOG(ERROR) << "graph is null.";
+      return RET_ERROR;
+    }
+    if (std::any_of(graph->input_indices_.begin(), graph->input_indices_.end(),
+                    [&tensor_size](const uint32_t &idx) { return idx >= tensor_size; })) {
+      MS_LOG(ERROR) << "Index of graph->input_indices_ is beyond tensor_size.";
+      return RET_ERROR;
+    }
+    if (std::any_of(graph->output_indices_.begin(), graph->output_indices_.end(),
+                    [&tensor_size](const uint32_t &idx) { return idx >= tensor_size; })) {
+      MS_LOG(ERROR) << "Index of graph->output_indices_ is beyond tensor_size.";
+      return RET_ERROR;
+    }
+    if (std::any_of(graph->tensor_indices_.begin(), graph->tensor_indices_.end(),
+                    [&tensor_size](const uint32_t &idx) { return idx >= tensor_size; })) {
+      MS_LOG(ERROR) << "Index of graph->tensor_indices_ is beyond tensor_size.";
+      return RET_ERROR;
+    }
+    if (std::any_of(graph->node_indices_.begin(), graph->node_indices_.end(),
+                    [&node_size](const uint32_t &idx) { return idx >= node_size; })) {
+      MS_LOG(ERROR) << "Index of graph->node_indices_ is beyond node_size.";
+      return RET_ERROR;
+    }
+  }
+  return RET_OK;
+}
+
 Model *ImportFromBuffer(const char *model_buf, size_t size, bool take_buf) {
   if (model_buf == nullptr) {
     MS_LOG(ERROR) << "The model buf is nullptr";
@@ -245,6 +314,6 @@ Model *ImportFromBuffer(const char *model_buf, size_t size, bool take_buf) {
     delete model;
     return nullptr;
   }
-  return model;
+  return NodeVerify(*model) == RET_OK && SubGraphVerify(*model) == RET_OK ? model : nullptr;
 }
 }  // namespace mindspore::lite

mindspore/lite/src/model_common.h

@@ -28,6 +28,10 @@ int ConvertSubGraph(const schema::SubGraph *sub_graph, Model *model);
 
 int MetaGraphMappingSubGraph(const mindspore::schema::MetaGraph *meta_graph, Model *model);
 
+int NodeVerify(const Model &model);
+
+int SubGraphVerify(const Model &model);
+
 Model *ImportFromBuffer(const char *model_buf, size_t size, bool take_buf);
 }  // namespace mindspore::lite
 #endif  // MINDSPORE_LITE_SRC_MODEL_COMMON_H_

mindspore/lite/src/scheduler.cc

@@ -123,6 +123,13 @@ int Scheduler::InferShape(const lite::Model *model, std::vector<Tensor *> *tenso
       MS_LOG(ERROR) << "InferShape failed, name: " << node->name_ << ", type: "
                     << schema::EnumNamePrimitiveType(static_cast<schema::PrimitiveType>(primitive->Type()));
       return RET_INFER_ERR;
+    } else {
+      for (auto &output : outputs) {
+        if (output->ElementsNum() >= MAX_MALLOC_SIZE / static_cast<int>(sizeof(int64_t))) {
+          MS_LOG(ERROR) << "The size of output tensor is too big";
+          return RET_ERROR;
+        }
+      }
     }
   }
 

mindspore/lite/tools/converter/parser/onnx/onnx_relu_parser.cc

@@ -86,23 +86,27 @@ STATUS OnnxPReluParser::Parse(const onnx::GraphProto &onnx_graph, const onnx::No
     }
   }
 
-  const onnx::TensorProto *slope = &params[0];
-  if (slope == nullptr) {
-    MS_LOG(ERROR) << "input error: params[0] is null";
-    return RET_ERROR;
-  }
-  const auto slope_raw_data = reinterpret_cast<const float *>(slope->raw_data().data());
-  const int64_t slope_size = slope->raw_data().size() / sizeof(float);
-  if (slope_size == 1) {
-    attr->slope.push_back(*slope_raw_data);
-    attr->channelShared = true;
-  } else {
-    attr->slope.resize(slope_size);
-    attr->channelShared = false;
-    if (memcpy_s(attr->slope.data(), slope_size * sizeof(float), slope_raw_data, slope_size * sizeof(float)) != 0) {
-      MS_LOG(ERROR) << "memcpy_s failed";
+  if (!params.empty()) {
+    const onnx::TensorProto *slope = &params[0];
+    if (slope == nullptr) {
+      MS_LOG(ERROR) << "input error: params[0] is null";
       return RET_ERROR;
     }
+    const auto slope_raw_data = reinterpret_cast<const float *>(slope->raw_data().data());
+    const int64_t slope_size = slope->raw_data().size() / sizeof(float);
+    if (slope_size == 1) {
+      attr->slope.push_back(*slope_raw_data);
+      attr->channelShared = true;
+    } else {
+      attr->slope.resize(slope_size);
+      attr->channelShared = false;
+      if (memcpy_s(attr->slope.data(), slope_size * sizeof(float), slope_raw_data, slope_size * sizeof(float)) != EOK) {
+        MS_LOG(ERROR) << "memcpy_s failed";
+        return RET_ERROR;
+      }
+    }
+  } else {
+    MS_LOG(WARNING) << "The slope pf prelu is null, which may cause errors.";
   }
 
   op->primitive->value.type = schema::PrimitiveType_PReLU;