summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
author오형석/On-Device Lab(SR)/Staff Engineer/삼성전자 <hseok82.oh@samsung.com>2019-10-15 03:17:06 (GMT)
committer이춘석/On-Device Lab(SR)/Staff Engineer/삼성전자 <chunseok.lee@samsung.com>2019-10-15 03:17:06 (GMT)
commit2d23c9c447332a15efbf6fbe8014c724a5f3315b (patch)
tree840df965dc7a97a0f1ef3e9cfaf8ad4fbe26d534
parent6c678f535b80cc83b1e1996922d73eb6b22254da (diff)
downloadnnfw-2d23c9c447332a15efbf6fbe8014c724a5f3315b.zip
nnfw-2d23c9c447332a15efbf6fbe8014c724a5f3315b.tar.gz
nnfw-2d23c9c447332a15efbf6fbe8014c724a5f3315b.tar.bz2
[m2] Model file verification (#8134)
Verification check before tflite model loading Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>
-rw-r--r--runtimes/neurun/frontend/tflite/loader.cc13
1 files changed, 10 insertions, 3 deletions
diff --git a/runtimes/neurun/frontend/tflite/loader.cc b/runtimes/neurun/frontend/tflite/loader.cc
index 844365b..14acbfc 100644
--- a/runtimes/neurun/frontend/tflite/loader.cc
+++ b/runtimes/neurun/frontend/tflite/loader.cc
@@ -40,11 +40,13 @@ using namespace neurun;
void Loader::loadFromFile(const char *file_path)
{
- std::ifstream stream(file_path, std::fstream::ate | std::fstream::binary);
+ std::ifstream stream(file_path, std::fstream::in | std::fstream::binary);
+
+ stream.seekg(0, stream.end);
auto size = stream.tellg();
- stream.seekg(0);
+ stream.seekg(0, stream.beg);
- _buffer.reserve(size);
+ _buffer.resize(size);
stream.read(_buffer.data(), size);
stream.close();
@@ -653,6 +655,11 @@ void Loader::loadConstantTensor(const Buffer *buffer, const uint32_t &index)
void Loader::loadModel()
{
+ flatbuffers::Verifier verifier(reinterpret_cast<const std::uint8_t *>(_buffer.data()),
+ _buffer.size());
+ if (!tflite::VerifyModelBuffer(verifier))
+ throw std::runtime_error{"Invalid tflite model"};
+
const auto *model = GetModel(_buffer.data());
// Version unused
// const auto version = model->version();