From 12c92b2b1fe2b2a0d217c91dc0dfe663e7473560 Mon Sep 17 00:00:00 2001 From: hyeonseok lee Date: Thu, 16 Mar 2023 14:30:12 +0900 Subject: [PATCH] [ahub] fix ahub issue - cast unsigned value to signed value and check its value it overrides sign-bit. Signed-off-by: hyeonseok lee --- nntrainer/compiler/tflite_interpreter.cpp | 8 +++++++- nntrainer/models/neuralnet.cpp | 2 +- nntrainer/tensor/tensor.cpp | 17 +++++++++++++---- nntrainer/utils/util_func.cpp | 16 ++++++++++++++-- 4 files changed, 35 insertions(+), 8 deletions(-) diff --git a/nntrainer/compiler/tflite_interpreter.cpp b/nntrainer/compiler/tflite_interpreter.cpp index 3140245..8e53e20 100644 --- a/nntrainer/compiler/tflite_interpreter.cpp +++ b/nntrainer/compiler/tflite_interpreter.cpp @@ -57,7 +57,13 @@ void builder2file(const flatbuffers::FlatBufferBuilder &builder, NNTR_THROW_IF(!os.good(), std::invalid_argument) << FUNC_TAG << "failed to open, reason: " << strerror_r(errno, error_buf, error_buflen); - os.write((char *)builder.GetBufferPointer(), builder.GetSize()); + + std::streamsize sz = static_cast(builder.GetSize()); + NNTR_THROW_IF(sz < 0, std::invalid_argument) + << FUNC_TAG << "builder size: " << builder.GetSize() + << " is too big. It cannot be represented by std::streamsize"; + + os.write((char *)builder.GetBufferPointer(), sz); os.close(); } diff --git a/nntrainer/models/neuralnet.cpp b/nntrainer/models/neuralnet.cpp index feb1284..a9148bf 100644 --- a/nntrainer/models/neuralnet.cpp +++ b/nntrainer/models/neuralnet.cpp @@ -407,7 +407,7 @@ void NeuralNetwork::save(const std::string &file_path, } if (opt && istrequal(opt->getType(), "adam")) { std::string adam = "adam"; - model_file.write(adam.c_str(), adam.size()); + model_file.write(adam.c_str(), 4); for (auto iter = model_graph.cbegin(); iter != model_graph.cend(); iter++) { (*iter)->save(model_file, true); diff --git a/nntrainer/tensor/tensor.cpp b/nntrainer/tensor/tensor.cpp index 63f7e48..f76b13f 100644 --- a/nntrainer/tensor/tensor.cpp +++ b/nntrainer/tensor/tensor.cpp @@ -1600,8 +1600,12 @@ void Tensor::save(std::ostream &file) { NNTR_THROW_IF(!contiguous, std::invalid_argument) << getName() << " is not contiguous, cannot save."; - checkedWrite(file, (char *)getData(), bytes(), - "[Tensor::save] operation failed"); + std::streamsize sz = static_cast(bytes()); + NNTR_THROW_IF(sz < 0, std::invalid_argument) + << "save size: " << bytes() + << " is too big. It cannot be represented by std::streamsize"; + + checkedWrite(file, (char *)getData(), sz, "[Tensor::save] operation failed"); putData(); } @@ -1609,8 +1613,13 @@ void Tensor::read(std::ifstream &file) { NNTR_THROW_IF(!contiguous, std::invalid_argument) << getName() << " is not contiguous, cannot read."; - checkedRead(file, (char *)getData(), bytes(), - "[Tensor::read] operation failed"); + std::streamsize sz = static_cast(bytes()); + + NNTR_THROW_IF(sz < 0, std::invalid_argument) + << "read size: " << bytes() + << " is too big. It cannot be represented by std::streamsize"; + + checkedRead(file, (char *)getData(), sz, "[Tensor::read] operation failed"); putData(); } diff --git a/nntrainer/utils/util_func.cpp b/nntrainer/utils/util_func.cpp index 2deaadf..0c5b864 100644 --- a/nntrainer/utils/util_func.cpp +++ b/nntrainer/utils/util_func.cpp @@ -98,8 +98,14 @@ std::string readString(std::ifstream &file, const char *error_msg) { size_t size; checkedRead(file, (char *)&size, sizeof(size), error_msg); + + std::streamsize sz = static_cast(size); + NNTR_THROW_IF(sz < 0, std::invalid_argument) + << "read string size: " << sz + << " is too big. It cannot be represented by std::streamsize"; + str.resize(size); - checkedRead(file, (char *)&str[0], size, error_msg); + checkedRead(file, (char *)&str[0], sz, error_msg); return str; } @@ -109,7 +115,13 @@ void writeString(std::ofstream &file, const std::string &str, size_t size = str.size(); checkedWrite(file, (char *)&size, sizeof(size), error_msg); - checkedWrite(file, (char *)&str[0], size, error_msg); + + std::streamsize sz = static_cast(size); + NNTR_THROW_IF(sz < 0, std::invalid_argument) + << "write string size: " << size + << " is too big. It cannot be represented by std::streamsize"; + + checkedWrite(file, (char *)&str[0], sz, error_msg); } bool endswith(const std::string &target, const std::string &suffix) { -- 2.7.4