Add additional ANEURALNETWORKS_BAD_DATA check into updateDimensionInfo (#910)

This commit adds additional ANEURALNETWORKS_BAD_DATA check into updateDimensionInfo.
These codes have copied from android NN master branch(af63a36)

Fix : #805

Signed-off-by: sjsujinkim <sjsujin.kim@samsung.com>

diff --git a/src/runtime/ref/nn/depend/libhidl/base/include/hidl/HidlSupport.h b/src/runtime/ref/nn/depend/libhidl/base/include/hidl/HidlSupport.h
index 1d5184e..bc8d56e 100644 (file)
--- a/src/runtime/ref/nn/depend/libhidl/base/include/hidl/HidlSupport.h
+++ b/src/runtime/ref/nn/depend/libhidl/base/include/hidl/HidlSupport.h
@@ -313,6 +313,9 @@ struct hidl_vec {
         static_assert(hidl_vec<T>::kOffsetOfBuffer == 0, "wrong offset");
     }
 
+    // Note, does not initialize primitive types.
+    hidl_vec(size_t size) : hidl_vec() { resize(size); }
+
     hidl_vec(const hidl_vec<T> &other) : hidl_vec() {
         *this = other;
     }

diff --git a/src/runtime/ref/nn/runtime/ExecutionBuilder.cpp b/src/runtime/ref/nn/runtime/ExecutionBuilder.cpp
index 37ad99e..333836e 100644 (file)
--- a/src/runtime/ref/nn/runtime/ExecutionBuilder.cpp
+++ b/src/runtime/ref/nn/runtime/ExecutionBuilder.cpp
@@ -60,17 +60,30 @@ int ModelArgumentInfo::setFromMemory(const Operand& operand, const ANeuralNetwor
 
 int ModelArgumentInfo::updateDimensionInfo(const Operand& operand,
                                            const ANeuralNetworksOperandType* newType) {
+    nnAssert(dimensions.empty());
     if (newType == nullptr) {
-        dimensions = hidl_vec<uint32_t>();
+        for (auto i : operand.dimensions) {
+            if (i == 0) {
+                LOG(ERROR) << "Setting input/output with unspecified dimensions";
+                return ANEURALNETWORKS_BAD_DATA;
+            }
+        }
+        dimensions = operand.dimensions;
     } else {
         uint32_t count = newType->dimensionCount;
         if (static_cast<OperandType>(newType->type) != operand.type ||
             count != operand.dimensions.size()) {
-            LOG(ERROR) << "ANeuralNetworksExecution_setInput/Output incompatible types";
+            LOG(ERROR) << "Setting input/output with incompatible types";
             return ANEURALNETWORKS_BAD_DATA;
         }
+        dimensions = hidl_vec<uint32_t>(count);
         for (uint32_t i = 0; i < count; i++) {
-            dimensions[i] = newType->dimensions[i];
+            if (operand.dimensions[i] != 0 && operand.dimensions[i] != newType->dimensions[i]) {
+                LOG(ERROR) << "Overriding a fully specified dimension is disallowed";
+                return ANEURALNETWORKS_BAD_DATA;
+            } else {
+                dimensions[i] = newType->dimensions[i];
+            }
         }
     }
     return ANEURALNETWORKS_NO_ERROR;