From 2d23c9c447332a15efbf6fbe8014c724a5f3315b Mon Sep 17 00:00:00 2001 From: =?utf8?q?=EC=98=A4=ED=98=95=EC=84=9D/On-Device=20Lab=28SR=29/Staff?= =?utf8?q?=20Engineer/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= Date: Tue, 15 Oct 2019 12:17:06 +0900 Subject: [PATCH] [m2] Model file verification (#8134) Verification check before tflite model loading Signed-off-by: Hyeongseok Oh --- runtimes/neurun/frontend/tflite/loader.cc | 13 ++++++++++--- 1 file 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(_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(); -- 2.7.4