From dd23c1a5730ccb9a4d17dc6bb1094a0c9625f317 Mon Sep 17 00:00:00 2001
From: yangruoqi713 <yangruoqi@huawei.com>
Date: Tue, 30 Nov 2021 15:43:45 +0800
Subject: [PATCH] [MSLITE][DEVELOP] fix bug of c++ api

---
 .../lite/src/cxx_api/model/model_impl.cc      |  8 ++
 mindspore/lite/src/cxx_api/types.cc           | 84 ++++++++++++-------
 2 files changed, 63 insertions(+), 29 deletions(-)

diff --git a/mindspore/lite/src/cxx_api/model/model_impl.cc b/mindspore/lite/src/cxx_api/model/model_impl.cc
index 58da64ebdda..c74d0df72c8 100644
--- a/mindspore/lite/src/cxx_api/model/model_impl.cc
+++ b/mindspore/lite/src/cxx_api/model/model_impl.cc
@@ -45,6 +45,14 @@ CreateTrainSessionProto *CreateTrainSessionCallbackHolder(CreateTrainSessionProt
 
 Status ModelImpl::Build(const void *model_data, size_t data_size, ModelType model_type,
                         const std::shared_ptr<Context> &ms_context) {
+  if (model_data == nullptr) {
+    MS_LOG(ERROR) << "The input model buffer is nullptr.";
+    return kLiteNullptr;
+  }
+  if (data_size == 0) {
+    MS_LOG(ERROR) << "The input model buffer size is 0.";
+    return kLiteInputParamInvalid;
+  }
   context_ = ms_context;
 
   auto *lite_context = new (std::nothrow) lite::InnerContext();
diff --git a/mindspore/lite/src/cxx_api/types.cc b/mindspore/lite/src/cxx_api/types.cc
index 2026129b10b..55d8a2da94e 100644
--- a/mindspore/lite/src/cxx_api/types.cc
+++ b/mindspore/lite/src/cxx_api/types.cc
@@ -62,8 +62,7 @@ class Buffer::Impl {
       return false;
     }
 
-    memcpy(MutableData(), data, data_len);
-
+    (void)memcpy(MutableData(), data, data_len);
     return true;
   }
 
@@ -94,30 +93,41 @@ MSTensor *MSTensor::CreateTensor(const std::vector<char> &name, enum DataType ty
     return nullptr;
   }
   if (data_len > 0 && data == nullptr) {
-    MS_LOG(ERROR) << "Mull data ptr of tensor.";
+    MS_LOG(ERROR) << "Null data ptr of tensor.";
     return nullptr;
   }
-  auto impl = Impl::CreateTensorImpl(CharToString(name), type, shape, nullptr, data_len);
+  if (data_len == 0 && data != nullptr) {
+    MS_LOG(ERROR) << "Data len doesn't match the data buffer size.";
+    return nullptr;
+  }
+
+  void *new_data = nullptr;
+  if (data != nullptr) {
+    new_data = malloc(data_len);
+    if (new_data == nullptr) {
+      MS_LOG(ERROR) << "Allocate data failed.";
+      return nullptr;
+    }
+    (void)memcpy(new_data, data, data_len);
+  }
+  auto impl = Impl::CreateTensorImpl(CharToString(name), type, shape, new_data, data_len);
   if (impl == nullptr) {
     MS_LOG(ERROR) << "Allocate tensor impl failed.";
+    if (new_data != nullptr) {
+      free(new_data);
+    }
     return nullptr;
   }
-  impl->set_own_data(true);
 
   auto ms_tensor = new (std::nothrow) MSTensor(impl);
   if (ms_tensor == nullptr) {
-    MS_LOG(ERROR) << "Allocate tensor impl failed.";
+    MS_LOG(ERROR) << "Allocate MSTensor failed.";
+    if (new_data != nullptr) {
+      free(new_data);
+    }
     return nullptr;
   }
-
-  if (data != nullptr) {
-    if (ms_tensor->MutableData() == nullptr) {
-      MS_LOG(ERROR) << "Allocate data failed.";
-      delete ms_tensor;
-      return nullptr;
-    }
-    ::memcpy(ms_tensor->MutableData(), data, data_len);
-  }
+  impl->set_own_data(true);
   return ms_tensor;
 }
 
@@ -191,27 +201,43 @@ MSTensor *MSTensor::Clone() const {
     MS_LOG(ERROR) << "Illegal data size of tensor.";
     return nullptr;
   }
-  auto impl = Impl::CreateTensorImpl(this->Name(), this->DataType(), this->Shape(), nullptr, data_len);
-  if (impl == nullptr) {
-    MS_LOG(ERROR) << "Allocate tensor impl failed.";
+  if (data_len > 0 && impl_->Data() == nullptr) {
+    MS_LOG(ERROR) << "Null data ptr of tensor.";
+    return nullptr;
+  }
+  if (data_len == 0 && impl_->Data() != nullptr) {
+    MS_LOG(ERROR) << "Data len doesn't match the data buffer size.";
+    return nullptr;
+  }
+
+  void *new_data = nullptr;
+  if (impl_->Data() != nullptr) {
+    new_data = malloc(data_len);
+    if (new_data == nullptr) {
+      MS_LOG(ERROR) << "Allocate data failed.";
+      return nullptr;
+    }
+    (void)memcpy(new_data, impl_->MutableData(), data_len);
+  }
+
+  auto impl = Impl::CreateTensorImpl(this->Name(), this->DataType(), this->Shape(), new_data, data_len);
+  if (impl == nullptr) {
+    MS_LOG(ERROR) << "Allocate tensor impl failed.";
+    if (new_data != nullptr) {
+      free(new_data);
+    }
     return nullptr;
   }
-  impl->set_own_data(true);
 
   auto ms_tensor = new (std::nothrow) MSTensor(impl);
   if (ms_tensor == nullptr) {
-    MS_LOG(ERROR) << "Allocate tensor impl failed.";
+    MS_LOG(ERROR) << "Allocate MSTensor failed.";
+    if (new_data != nullptr) {
+      free(new_data);
+    }
     return nullptr;
   }
-
-  if (impl_->Data() != nullptr) {
-    if (ms_tensor->MutableData() == nullptr) {
-      MS_LOG(ERROR) << "Allocate data failed.";
-      delete ms_tensor;
-      return nullptr;
-    }
-    ::memcpy(ms_tensor->MutableData(), impl_->MutableData(), data_len);
-  }
+  impl->set_own_data(true);
   return ms_tensor;
 }