[Triv2/Fix] Add error handling codes to check too large input data

This patch adds error handling codes to check too larget input data.

Signed-off-by: Dongju Chae <dongju.chae@samsung.com>

diff --git a/src/core/ne-handler.cc b/src/core/ne-handler.cc
index a762ced..d7954ba 100644 (file)
--- a/src/core/ne-handler.cc
+++ b/src/core/ne-handler.cc
@@ -1422,12 +1422,21 @@ TrinityVision2::run (npu_input_opmode opmode, const Model *model,
 
   /** extract input data */
   for (uint32_t idx = 0; idx < input->num_buffers; idx++) {
+    size_t max_seg_size = segt->getInputSegment(idx)->getSize();
+    uint32_t seg_offset = segt->getInputSegmentOffset(idx);
+
+    if (input->bufs[idx].size + seg_offset > max_seg_size) {
+      logerr (TAG, "Too large input data provided: max segment size (%zu)\n",
+          max_seg_size);
+      return -ERANGE;
+    }
+
     if (!segt->getInputSegment(idx)->isExternal ()) {
       auto func = std::bind (TrinityVision2::manipulateData, model, idx, true,
           std::placeholders::_1, std::placeholders::_2, std::placeholders::_3);
       int status = comm_.extractGenericBuffer (
           &input->bufs[idx],
-          segt->getInputSegment(idx)->getData() + segt->getInputSegmentOffset(idx),
+          segt->getInputSegment(idx)->getData() + seg_offset,
           func);
       if (status != 0) {
         logerr (TAG, "Failed to feed input segment: %d\n", status);