From fd532c26117b1a5b70f9a348cb375202fa18b0d8 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Tomasz=20Do=C5=82bniak?= Date: Wed, 25 Nov 2020 15:41:25 +0100 Subject: [PATCH] Fix the potential unlimited mem allocation problem (#3347) --- inference-engine/src/readers/onnx_reader/onnx_model_validator.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/inference-engine/src/readers/onnx_reader/onnx_model_validator.cpp b/inference-engine/src/readers/onnx_reader/onnx_model_validator.cpp index 4f8d8b9..fa7660d 100644 --- a/inference-engine/src/readers/onnx_reader/onnx_model_validator.cpp +++ b/inference-engine/src/readers/onnx_reader/onnx_model_validator.cpp @@ -89,13 +89,16 @@ namespace onnx { uint32_t decode_varint(std::istream& model) { std::vector bytes; - bytes.reserve(4); + // max 4 bytes for a single value because this function returns a 32-bit long decoded varint + const size_t MAX_VARINT_BYTES = 4u; + // optimization to avoid allocations during push_back calls + bytes.reserve(MAX_VARINT_BYTES); char key_component = 0; model.get(key_component); // keep reading all bytes from the stream which have the MSB on - while (key_component & 0x80) { + while (key_component & 0x80 && bytes.size() < MAX_VARINT_BYTES) { // drop the most significant bit const char component = key_component & ~0x80; bytes.push_back(component); -- 2.7.4