Check input and output tensor length (#4847)

Check correctness of length for input and output tensor setting based on NNAPI spec

Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>

diff --git a/runtimes/neurun/frontend/nnapi/execution.cc b/runtimes/neurun/frontend/nnapi/execution.cc
index 4463c8e..ac63f65 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/execution.cc
+++ b/runtimes/neurun/frontend/nnapi/execution.cc
@@ -21,6 +21,7 @@
 #include "wrapper/compilation.h"
 #include "wrapper/execution.h"
 #include "wrapper/event.h"
+#include "wrapper/NNAPIConvert.h"
 #include "util/logging.h"
 
 //
@@ -104,6 +105,12 @@ int ANeuralNetworksExecution_setInput(ANeuralNetworksExecution *execution, int32
       VERBOSE(NNAPI::Execution) << "setInput: Shape mismatch" << std::endl;
       return ANEURALNETWORKS_BAD_DATA;
     }
+
+    if (NNAPIConvert::calculateSizeFromType(type) != length)
+    {
+      VERBOSE(NNAPI::Execution) << "setInput: Invalid length" << std::endl;
+      return ANEURALNETWORKS_BAD_DATA;
+    }
   }
   else
   {
@@ -112,6 +119,12 @@ int ANeuralNetworksExecution_setInput(ANeuralNetworksExecution *execution, int32
       VERBOSE(NNAPI::Execution) << "setInput: Unspecified dimension value" << std::endl;
       return ANEURALNETWORKS_BAD_DATA;
     }
+
+    if (execution->getOperandSize(operand_index) != length)
+    {
+      VERBOSE(NNAPI::Execution) << "setInput: Invalid length" << std::endl;
+      return ANEURALNETWORKS_BAD_DATA;
+    }
   }
 
   if (!execution->setInput(index, type, buffer, length))
@@ -171,6 +184,12 @@ int ANeuralNetworksExecution_setOutput(ANeuralNetworksExecution *execution, int3
       VERBOSE(NNAPI::Execution) << "setOutput: Shape mismatch" << std::endl;
       return ANEURALNETWORKS_BAD_DATA;
     }
+
+    if (NNAPIConvert::calculateSizeFromType(type) != length)
+    {
+      VERBOSE(NNAPI::Execution) << "setOutput: Invalid length" << std::endl;
+      return ANEURALNETWORKS_BAD_DATA;
+    }
   }
   else
   {
@@ -179,6 +198,12 @@ int ANeuralNetworksExecution_setOutput(ANeuralNetworksExecution *execution, int3
       VERBOSE(NNAPI::Execution) << "setOutput: Unspecified dimension value" << std::endl;
       return ANEURALNETWORKS_BAD_DATA;
     }
+
+    if (execution->getOperandSize(operand_index) != length)
+    {
+      VERBOSE(NNAPI::Execution) << "setInput: Invalid length" << std::endl;
+      return ANEURALNETWORKS_BAD_DATA;
+    }
   }
 
   if (!execution->setOutput(index, type, buffer, length))

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.cc b/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.cc
index 25e09e5..e1f194f 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.cc
+++ b/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.cc
@@ -16,6 +16,8 @@
 
 #include "NNAPIConvert.h"
 
+#include <numeric>
+
 using namespace ::neurun::model;
 
 operand::DataType NNAPIConvert::getDataType(OperandCode type)
@@ -42,3 +44,13 @@ operand::Shape NNAPIConvert::getShape(const ANeuralNetworksOperandType *type)
 
   return shape;
 }
+
+size_t NNAPIConvert::calculateSizeFromType(const ANeuralNetworksOperandType *type)
+{
+  auto shape = getShape(type);
+  auto data_type = getDataType((OperandCode)(type->type));
+
+  const auto &dims = shape.dims();
+  return std::accumulate(dims.begin(), dims.end(), operand::sizeOfDataType(data_type),
+                         std::multiplies<size_t>());
+}

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.h b/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.h
index 9192d55..3c8c22a 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.h
+++ b/runtimes/neurun/frontend/nnapi/wrapper/NNAPIConvert.h
@@ -54,6 +54,13 @@ public:
    * @return    neurun's internal operand shape
    */
   static ::neurun::model::operand::Shape getShape(const ANeuralNetworksOperandType *type);
+
+  /**
+   * @brief     Calcaulate operand size from NNAPI type
+   * @param[in] type  NNAPI's operand type
+   * @return    Operand size
+   */
+  static size_t calculateSizeFromType(const ANeuralNetworksOperandType *type);
 };
 
 #endif // __NEURUN_NNAPI_CONVERT_H__

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/execution.cc b/runtimes/neurun/frontend/nnapi/wrapper/execution.cc
index 5b1e6db..a7fdbea 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/execution.cc
+++ b/runtimes/neurun/frontend/nnapi/wrapper/execution.cc
@@ -93,6 +93,11 @@ bool ANeuralNetworksExecution::haveUnspecifiedDims(
   return ((operand_shape.element_nums() == 0) ? true : false);
 }
 
+size_t ANeuralNetworksExecution::getOperandSize(const neurun::model::operand::Index index) noexcept
+{
+  return _executor->model().operands.at(index).operandSize();
+}
+
 bool ANeuralNetworksExecution::setInput(uint32_t index, const ANeuralNetworksOperandType *type,
                                         const void *buffer, size_t length) noexcept
 {

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/execution.h b/runtimes/neurun/frontend/nnapi/wrapper/execution.h
index e491327..66db7ed 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/execution.h
+++ b/runtimes/neurun/frontend/nnapi/wrapper/execution.h
@@ -44,6 +44,7 @@ public:
   bool compareShape(const ANeuralNetworksOperandType *type,
                     const neurun::model::operand::Index index) noexcept;
   bool haveUnspecifiedDims(const neurun::model::operand::Index index) noexcept;
+  size_t getOperandSize(const neurun::model::operand::Index index) noexcept;
 
 private:
   std::shared_ptr<neurun::exec::IExecutor> _executor;