From b58802c1cc57a984f65162a71f8d16b325493dee Mon Sep 17 00:00:00 2001 From: yvette Date: Wed, 2 Dec 2020 15:24:09 +0800 Subject: [PATCH] verify shape and index rationality --- mindspore/lite/src/inner_context.cc | 4 + mindspore/lite/src/lite_kernel.cc | 5 + mindspore/lite/src/model_common.cc | 123 ++++++++++++++---- mindspore/lite/src/model_common.h | 4 + mindspore/lite/src/scheduler.cc | 7 + .../converter/parser/onnx/onnx_relu_parser.cc | 34 ++--- 6 files changed, 135 insertions(+), 42 deletions(-) diff --git a/mindspore/lite/src/inner_context.cc b/mindspore/lite/src/inner_context.cc index 76ce01761b..f6c6f3a373 100644 --- a/mindspore/lite/src/inner_context.cc +++ b/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."; diff --git a/mindspore/lite/src/lite_kernel.cc b/mindspore/lite/src/lite_kernel.cc index 7f490a578e..30a6a5360f 100644 --- a/mindspore/lite/src/lite_kernel.cc +++ b/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(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"; diff --git a/mindspore/lite/src/model_common.cc b/mindspore/lite/src/model_common.cc index 272bdae957..06b861716b 100644 --- a/mindspore/lite/src/model_common.cc +++ b/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(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(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(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(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(const_cast(prim)); + if (whileOp == nullptr) { + MS_LOG(ERROR) << "whileOp is null."; + return RET_ERROR; + } + if (static_cast(whileOp->GetBodySubgraphIndex()) >= subGraph_size || + static_cast(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 diff --git a/mindspore/lite/src/model_common.h b/mindspore/lite/src/model_common.h index 2162328e01..dfcb21e477 100644 --- a/mindspore/lite/src/model_common.h +++ b/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_ diff --git a/mindspore/lite/src/scheduler.cc b/mindspore/lite/src/scheduler.cc index 41c2caff33..7b2eab5365 100644 --- a/mindspore/lite/src/scheduler.cc +++ b/mindspore/lite/src/scheduler.cc @@ -123,6 +123,13 @@ int Scheduler::InferShape(const lite::Model *model, std::vector *tenso MS_LOG(ERROR) << "InferShape failed, name: " << node->name_ << ", type: " << schema::EnumNamePrimitiveType(static_cast(primitive->Type())); return RET_INFER_ERR; + } else { + for (auto &output : outputs) { + if (output->ElementsNum() >= MAX_MALLOC_SIZE / static_cast(sizeof(int64_t))) { + MS_LOG(ERROR) << "The size of output tensor is too big"; + return RET_ERROR; + } + } } } diff --git a/mindspore/lite/tools/converter/parser/onnx/onnx_relu_parser.cc b/mindspore/lite/tools/converter/parser/onnx/onnx_relu_parser.cc index 2079a900c9..54d9154613 100644 --- a/mindspore/lite/tools/converter/parser/onnx/onnx_relu_parser.cc +++ b/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 = ¶ms[0]; - if (slope == nullptr) { - MS_LOG(ERROR) << "input error: params[0] is null"; - return RET_ERROR; - } - const auto slope_raw_data = reinterpret_cast(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 = ¶ms[0]; + if (slope == nullptr) { + MS_LOG(ERROR) << "input error: params[0] is null"; return RET_ERROR; } + const auto slope_raw_data = reinterpret_cast(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;