p60975123
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

check input model buffer size

This commit is contained in:
hangangqiang 2021-05-28 11:34:52 +08:00
parent 9fdc0035a0
commit 8e1b5f8220
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mindspore/lite/src/lite_model.cc
View File

@ -373,6 +373,10 @@ int LiteModel::ConstructModel() {
return ModelVerify() ? RET_OK : RET_ERROR;
}
namespace {
constexpr size_t kMaxModelBufferSize = static_cast<size_t>(1024) * 1024 * 1024 * 2;
}
Model *ImportFromBuffer(const char *model_buf, size_t size, bool take_buf) {
if (model_buf == nullptr) {
MS_LOG(ERROR) << "The model buf is nullptr";
@ -386,8 +390,8 @@ Model *ImportFromBuffer(const char *model_buf, size_t size, bool take_buf) {
if (take_buf) {
model->buf = const_cast<char *>(model_buf);
} else {
if (size == 0) {
MS_LOG(ERROR) << "malloc size is equal to 0";
if (size == 0 || size > kMaxModelBufferSize) {
MS_LOG(ERROR) << "Input model buffer size invalid, require (0, 2GB].";
delete (model);
return nullptr;
}