From b85ffa11f479d51775d74ca2caf46ce5def7a4a5 Mon Sep 17 00:00:00 2001 From: yangruoqi713 Date: Mon, 1 Nov 2021 11:21:15 +0800 Subject: [PATCH] [MSLITE][DEVELOP] include stdint.h for data_type, elementNum & Size overflow check --- include/api/data_type.h | 2 ++ mindspore/lite/src/runtime/kernel/arm/base/carry_data.cc | 2 ++ .../lite/src/runtime/kernel/arm/base/constant_of_shape.cc | 1 + .../lite/src/runtime/kernel/arm/base/convolution_base.cc | 1 + .../runtime/kernel/arm/base/detection_post_process_base.cc | 2 ++ .../runtime/kernel/arm/base/group_convolution_creator.cc | 1 + mindspore/lite/src/runtime/kernel/arm/base/one_hot_base.cc | 2 +- .../lite/src/runtime/kernel/arm/base/quant_dtype_cast.cc | 6 ++---- .../src/runtime/kernel/arm/base/random_standard_normal.cc | 1 + mindspore/lite/src/runtime/kernel/arm/base/reduce_base.cc | 3 ++- mindspore/lite/src/runtime/kernel/arm/base/reshape_base.cc | 1 + mindspore/lite/src/runtime/kernel/arm/base/select.cc | 1 + mindspore/lite/src/runtime/kernel/arm/base/stack_base.cc | 1 + 13 files changed, 18 insertions(+), 6 deletions(-) diff --git a/include/api/data_type.h b/include/api/data_type.h index 61eb1d51f2b..053244241b8 100644 --- a/include/api/data_type.h +++ b/include/api/data_type.h @@ -16,6 +16,8 @@ #ifndef MINDSPORE_INCLUDE_API_DATA_TYPE_H_ #define MINDSPORE_INCLUDE_API_DATA_TYPE_H_ +#include + namespace mindspore { enum class DataType : int { kTypeUnknown = 0, diff --git a/mindspore/lite/src/runtime/kernel/arm/base/carry_data.cc b/mindspore/lite/src/runtime/kernel/arm/base/carry_data.cc index 5665927a0ac..347751c0912 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/carry_data.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/carry_data.cc @@ -85,9 +85,11 @@ int CarryDataKernel::MoveTensorData(lite::Tensor *dst_tensor, lite::Tensor *src_ CHECK_NULL_RETURN(src_tensor->data()); CHECK_NULL_RETURN(dst_tensor->data()); // need replace with increase data ref count + MS_CHECK_FALSE(src_tensor->Size() == 0, RET_ERROR); memcpy(dst_tensor->data(), src_tensor->data(), src_tensor->Size()); return RET_OK; } + #ifndef CONTROLFLOW_TENSORLIST_CLIP int CarryDataKernel::MoveTensorListData(lite::TensorList *dst_tensorlist, lite::TensorList *src_tensorlist) { // shape may change, because tensors.size() can be change in RunGraph diff --git a/mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.cc b/mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.cc index 5fc176c80b6..743dd4f96af 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.cc @@ -71,6 +71,7 @@ int ConstantOfShapeCPUKernel::Run() { CHECK_NULL_RETURN(output); param_->data_type_ = output->data_type(); param_->element_size_ = output->ElementsNum(); + MS_CHECK_GT(param_->element_size_, 0, RET_ERROR); output_ptr_ = output->data(); CHECK_NULL_RETURN(output_ptr_); diff --git a/mindspore/lite/src/runtime/kernel/arm/base/convolution_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/convolution_base.cc index 9f13d7387af..df41144b72e 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/convolution_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/convolution_base.cc @@ -139,6 +139,7 @@ int ConvolutionBaseCPUKernel::InitConvWeightBias() { } if (in_tensors_.size() == kInputSize2) { + MS_CHECK_FALSE(in_tensors_.at(kBiasIndex)->Size() == 0, RET_ERROR); memcpy(bias_data_, origin_bias_, in_tensors_.at(kBiasIndex)->Size()); } else { MS_ASSERT(in_tensors_.size() == kInputSize1); diff --git a/mindspore/lite/src/runtime/kernel/arm/base/detection_post_process_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/detection_post_process_base.cc index 6c85a6d6438..926a15ad33f 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/detection_post_process_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/detection_post_process_base.cc @@ -49,6 +49,7 @@ int DetectionPostProcessBaseCPUKernel::Prepare() { params_->selected_ = nullptr; params_->anchors_ = nullptr; auto anchor_tensor = in_tensors_.at(2); + MS_CHECK_GT(anchor_tensor->ElementsNum(), 0, RET_ERROR); CHECK_NULL_RETURN(anchor_tensor->data()); if (anchor_tensor->data_type() == kNumberTypeInt8) { auto quant_param = anchor_tensor->quant_params().front(); @@ -78,6 +79,7 @@ int DetectionPostProcessBaseCPUKernel::Prepare() { MS_LOG(ERROR) << "Malloc anchor failed"; return RET_ERROR; } + MS_CHECK_FALSE(anchor_tensor->Size() == 0, RET_ERROR); memcpy(params_->anchors_, anchor_tensor->data(), anchor_tensor->Size()); } else { MS_LOG(ERROR) << "unsupported anchor data type " << anchor_tensor->data_type(); diff --git a/mindspore/lite/src/runtime/kernel/arm/base/group_convolution_creator.cc b/mindspore/lite/src/runtime/kernel/arm/base/group_convolution_creator.cc index e1ee7c4ca25..aaee63fe52a 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/group_convolution_creator.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/group_convolution_creator.cc @@ -76,6 +76,7 @@ lite::Tensor *CreateConstTensor(const lite::Tensor *tensor, const std::vectorSize() == 0, nullptr); uint8_t *new_tensor_data = reinterpret_cast(tensor->data()) + index * new_tensor->Size(); memcpy(new_tensor->data(), reinterpret_cast(new_tensor_data), new_tensor->Size()); return new_tensor; diff --git a/mindspore/lite/src/runtime/kernel/arm/base/one_hot_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/one_hot_base.cc index 57d711839ba..fb6116b5d15 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/one_hot_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/one_hot_base.cc @@ -82,7 +82,7 @@ int OneHotCPUKernel::ReSize() { return RET_ERROR; } inner_size_ = indices->ElementsNum() / outer_size_; - + MS_CHECK_GT(inner_size_, 0, RET_ERROR); return RET_OK; } diff --git a/mindspore/lite/src/runtime/kernel/arm/base/quant_dtype_cast.cc b/mindspore/lite/src/runtime/kernel/arm/base/quant_dtype_cast.cc index 1c6d02f584a..55ed42f8991 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/quant_dtype_cast.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/quant_dtype_cast.cc @@ -62,10 +62,7 @@ int QuantDTypeCastCPUKernel::ReSize() { MS_ASSERT(in_tensor != nullptr); num_unit_ = static_cast(in_tensor->ElementsNum()); thread_n_num_ = MSMIN(thread_num_, num_unit_); - if (thread_n_num_ == 0) { - MS_LOG(ERROR) << "div zero"; - return RET_ERROR; - } + MS_CHECK_GT(thread_n_num_, 0, RET_ERROR); thread_n_stride_ = UP_DIV(num_unit_, thread_n_num_); return RET_OK; } @@ -178,6 +175,7 @@ int QuantDTypeCastCPUKernel::Run() { if (int8_ptr_ == nullptr || int8_out_ptr_ == nullptr) { return RET_NULL_PTR; } + MS_CHECK_GT(in_tensors_[0]->ElementsNum(), 0, RET_ERROR); float32_ptr_ = new (std::nothrow) float[in_tensors_[0]->ElementsNum()]; if (float32_ptr_ == nullptr) { MS_LOG(ERROR) << "new float[] failed"; diff --git a/mindspore/lite/src/runtime/kernel/arm/base/random_standard_normal.cc b/mindspore/lite/src/runtime/kernel/arm/base/random_standard_normal.cc index 4a69bd0fcc3..c101b3df4bd 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/random_standard_normal.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/random_standard_normal.cc @@ -44,6 +44,7 @@ int RandomStandardNormalCPUKernel::Run() { std::default_random_engine engine{static_cast(random_seed)}; std::normal_distribution nums(0, 1.0); auto all_data_nums = out_tensors_[0]->ElementsNum(); + MS_CHECK_GT(all_data_nums, 0, RET_ERROR); auto out_data = out_tensors_[0]->data(); MS_ASSERT(out_data != nullptr); auto output = reinterpret_cast(out_data); diff --git a/mindspore/lite/src/runtime/kernel/arm/base/reduce_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/reduce_base.cc index 2afc2ff97ba..0cefe05dde7 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/reduce_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/reduce_base.cc @@ -106,7 +106,7 @@ int ReduceBaseCPUKernel::Prepare() { MS_CHECK_FALSE_MSG((axes_tensor->data_type() != kNumberTypeInt && axes_tensor->data_type() != kNumberTypeInt32), RET_ERROR, "The data type of axes tensor should be int32"); num_axes_ = axes_tensor->ElementsNum(); - if (axes_tensor->ElementsNum() > MAX_SHAPE_SIZE) { + if (num_axes_ <= 0 && num_axes_ > MAX_SHAPE_SIZE) { MS_LOG(ERROR) << "input axes invalid."; return RET_ERROR; } @@ -121,6 +121,7 @@ int ReduceBaseCPUKernel::Prepare() { return RET_ERROR; } } else { + MS_CHECK_FALSE(axes_tensor->Size() == 0, RET_ERROR); memcpy(axes_, axes_tensor->data(), axes_tensor->Size()); } } else { diff --git a/mindspore/lite/src/runtime/kernel/arm/base/reshape_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/reshape_base.cc index 40760a5d1b7..dfbab2bfdbd 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/reshape_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/reshape_base.cc @@ -44,6 +44,7 @@ int ReshapeBaseCPUKernel::Run() { op_parameter_->is_train_session_) { CHECK_NULL_RETURN(out_tensor->data()); CHECK_NULL_RETURN(in_tensor->data()); + MS_CHECK_FALSE(in_tensor->Size() == 0, RET_ERROR); memcpy(out_tensor->data(), in_tensor->data(), in_tensor->Size()); return RET_OK; } diff --git a/mindspore/lite/src/runtime/kernel/arm/base/select.cc b/mindspore/lite/src/runtime/kernel/arm/base/select.cc index 449c1b46425..e5f4c09a154 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/select.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/select.cc @@ -75,6 +75,7 @@ int SelectCPUKernel::Run() { } MS_ASSERT(in_tensors_.at(1)->Size() == out_tensors_.at(0)->Size()); auto size = in_tensors_.at(1)->ElementsNum(); + MS_CHECK_GT(size, 0, RET_ERROR); auto condition = static_cast(bool_tensor->data()); auto input1 = static_cast(in_tensors_.at(kFirstIdx)->data()); auto input2 = static_cast(in_tensors_.at(kSecondIdx)->data()); diff --git a/mindspore/lite/src/runtime/kernel/arm/base/stack_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/stack_base.cc index fc38e0a251d..39e95c9137d 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/stack_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/stack_base.cc @@ -68,6 +68,7 @@ int StackBaseCPUKernel::ReSize() { copy_size_ = GetCopyNum(input0_shape, axis_, input0_shape.size()) * data_type_size_; outer_size_ = GetOuterSize(input0_shape, axis_); } + MS_CHECK_GT(copy_size_, 0, RET_ERROR); return RET_OK; }