From 966dc9fb162103c0286ac98882fae639f0185b06 Mon Sep 17 00:00:00 2001
From: chenjianping <jpc.chen@huawei.com>
Date: Fri, 18 Sep 2020 20:55:56 +0800
Subject: [PATCH] malloc data nullptr check

---
 mindspore/lite/internal/src/common/string.cc  |  6 ++--
 .../lite/internal/src/kernel/fp32/matmul.cc   |  1 +
 .../lite/internal/src/kernel/fp32/reduce.cc   | 14 +++++-----
 mindspore/lite/internal/src/ms_tensor.cc      | 28 +++++++++++++++++--
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/mindspore/lite/internal/src/common/string.cc b/mindspore/lite/internal/src/common/string.cc
index 1bc675d027d..44852b8702b 100644
--- a/mindspore/lite/internal/src/common/string.cc
+++ b/mindspore/lite/internal/src/common/string.cc
@@ -199,7 +199,7 @@ void String::clear() noexcept {
 String &String::operator+=(const String &str) {
   size_t new_size = size_ + str.size_;
   char *tmp = reinterpret_cast<char *>(malloc(sizeof(char) * (new_size + 1)));
-  if (buffer_ == nullptr) {
+  if (tmp == nullptr) {
     MS_C_EXCEPTION("malloc data failed");
   }
   memcpy(tmp, this->buffer_, size_ + 1);
@@ -218,7 +218,7 @@ String &String::operator+=(const char *str) {
   size_t str_size = strlen(str);
   size_t new_size = size_ + str_size;
   char *tmp = reinterpret_cast<char *>(malloc(sizeof(char) * (new_size + 1)));
-  if (buffer_ == nullptr) {
+  if (tmp == nullptr) {
     MS_C_EXCEPTION("malloc data failed");
   }
   memcpy(tmp, this->buffer_, size_ + 1);
@@ -232,7 +232,7 @@ String &String::operator+=(const char *str) {
 
 String &String::operator+=(const char ch) {
   char *tmp = reinterpret_cast<char *>(malloc(sizeof(char) * (size_ + 2)));
-  if (buffer_ == nullptr) {
+  if (tmp == nullptr) {
     MS_C_EXCEPTION("malloc data failed");
   }
   memcpy(tmp, this->buffer_, size_ + 1);
diff --git a/mindspore/lite/internal/src/kernel/fp32/matmul.cc b/mindspore/lite/internal/src/kernel/fp32/matmul.cc
index dc0729e72eb..1f4694ac91c 100644
--- a/mindspore/lite/internal/src/kernel/fp32/matmul.cc
+++ b/mindspore/lite/internal/src/kernel/fp32/matmul.cc
@@ -155,6 +155,7 @@ int DoMatMul(const TensorPtrVector &in_tensors, const TensorPtrVector &out_tenso
   kernel_data->a_c12_ptr_
     = reinterpret_cast<float *>(allocator->Malloc(params->batch * params->row_12_ * params->deep_ * sizeof(float)));
   if (kernel_data->a_c12_ptr_ == NULL) {
+    FreeMatMulKernelData(kernel_data, allocator);
     return RET_MEMORY_FAILED;
   }
   memset(kernel_data->a_c12_ptr_, 0, params->row_12_ * params->deep_ * sizeof(float));
diff --git a/mindspore/lite/internal/src/kernel/fp32/reduce.cc b/mindspore/lite/internal/src/kernel/fp32/reduce.cc
index 9d6cbb679a1..235b611993c 100644
--- a/mindspore/lite/internal/src/kernel/fp32/reduce.cc
+++ b/mindspore/lite/internal/src/kernel/fp32/reduce.cc
@@ -58,12 +58,14 @@ int MallocTmpBuffer(std::vector<float *> *data_buffers, const ShapeVector &shape
   return RET_OK;
 }
 
-int FreeTmpBuffer(std::vector<float *> *data_buffers, mindspore::lite::Allocator *allocator) {
+void FreeTmpBuffer(std::vector<float *> *data_buffers, mindspore::lite::Allocator *allocator) {
+  if (data_buffers == nullptr) {
+    return;
+  }
   for (int i = 0; i < data_buffers->size(); ++i) {
     allocator->Free(data_buffers->at(i));
   }
   data_buffers->clear();
-  return RET_OK;
 }
 
 int RunReduce(Reducer reducer, std::vector<float *> data_buffers, float *in_data, float *out_data, Int32Vector axes,
@@ -220,13 +222,11 @@ int DoReduce(const TensorPtrVector &in_tensors, const TensorPtrVector &out_tenso
   }
   status = RunReduce(reducer, data_buffers, reinterpret_cast<float *>(in_tensors[0]->data_),
                      reinterpret_cast<float *>(out_tensors[0]->data_), axes, in_tensors[0]->shape_);
-  if (status != RET_OK) {
-    return status;
-  }
 
-  status = FreeTmpBuffer(&data_buffers, allocator);
+  FreeTmpBuffer(&data_buffers, allocator);
+
   if (status != RET_OK) {
-    return status;
+    return RET_ERROR;
   }
   return RET_OK;
 }
diff --git a/mindspore/lite/internal/src/ms_tensor.cc b/mindspore/lite/internal/src/ms_tensor.cc
index c5cd3afe93d..682f8a789e5 100644
--- a/mindspore/lite/internal/src/ms_tensor.cc
+++ b/mindspore/lite/internal/src/ms_tensor.cc
@@ -17,6 +17,7 @@
 #include "internal/include/string.h"
 #include "internal/include/vector.h"
 #include "internal/include/ms_tensor.h"
+#include "internal/src/lite_log.h"
 
 MSTensor *CreateTensor(TypeId data_type, const ShapeVector &shape) {
   MSTensor *tensor = new MSTensor();
@@ -28,7 +29,12 @@ MSTensor *CreateTensor(TypeId data_type, const ShapeVector &shape) {
   return tensor;
 }
 
-void DestroyTensor(MSTensor *ptr) { delete ptr; }
+void DestroyTensor(MSTensor *ptr) {
+  if (ptr == nullptr) {
+    return;
+  }
+  delete ptr;
+}
 
 int MSTensor::ElementsNum() const {
   int result = 1;
@@ -206,14 +212,30 @@ int MSTensor::ElementsC4Num() const {
 
 void *MSTensor::operator new(size_t sz) {
   void *storage = malloc(sz);
+  if (storage == nullptr) {
+    MS_C_EXCEPTION("malloc tensor fail!");
+  }
   return storage;
 }
 
 void *MSTensor::operator new[](size_t sz) {
   void *storage = malloc(sz);
+  if (storage == nullptr) {
+    MS_C_EXCEPTION("malloc tensor array fail!");
+  }
   return storage;
 }
 
-void MSTensor::operator delete(void *ptr, size_t sz) { free(ptr); }
+void MSTensor::operator delete(void *ptr, size_t sz) {
+  if (ptr == nullptr) {
+    return;
+  }
+  free(ptr);
+}
 
-void MSTensor::operator delete[](void *ptr, size_t sz) { free(ptr); }
+void MSTensor::operator delete[](void *ptr, size_t sz) {
+  if (ptr == nullptr) {
+    return;
+  }
+  free(ptr);
+}