Fix integer overflow in RGG.

Some RGG tests crashed when enabling the integer_overflow sanitizer
because the code attempted to take the negative of a uint32_t value
before casting to int32_t. This CL fixed the issue.

Without the sanitizer, the code will behave correctly, so the fix will
not alter the test behavior.

Test: NNAPI CTS
Test: NeuralNetworksTest_static_fuzzing with integer_overflow sanitizer
Change-Id: Id9ba1752ca063b44c78abbbd71e611b1d0f84fdf
Merged-In: Id9ba1752ca063b44c78abbbd71e611b1d0f84fdf
(cherry picked from commit 9da0f6420d01f929c1bbe316a2ecee5e926bdd07)

7 files changed, 29 insertions, 16 deletions

diff --git a/nn/runtime/test/Android.bp b/nn/runtime/test/Android.bp
index f51925899..1db4fb53b 100644
--- a/nn/runtime/test/Android.bp
+++ b/nn/runtime/test/Android.bp
@@ -190,6 +190,9 @@ cc_test {
     header_libs: [
         "libneuralnetworks_private_headers",
     ],
+    sanitize: {
+        integer_overflow: true,
+    },
 }
 
 cc_fuzz {
diff --git a/nn/runtime/test/fuzzing/operation_signatures/ConcatSplit.cpp b/nn/runtime/test/fuzzing/operation_signatures/ConcatSplit.cpp
index 22020cfc3..1522cf460 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/ConcatSplit.cpp
+++ b/nn/runtime/test/fuzzing/operation_signatures/ConcatSplit.cpp
@@ -120,9 +120,9 @@ DEFINE_OPERATION_SIGNATURE(CONCAT_3_V1_3){
 
 // SPLIT with fixed number of splits.
 static void splitConstructor(uint32_t numSplits, uint32_t rank, RandomOperation* op) {
-    int32_t axis = getUniform<int32_t>(-rank, rank - 1);
+    int32_t axis = getRandomAxis(rank);
     op->inputs[1]->setScalarValue<int32_t>(axis);
-    if (axis < 0) axis += rank;
+    axis = toPositiveAxis(axis, rank);
 
     op->inputs[0]->dimensions.resize(rank);
     for (uint32_t i = 0; i < numSplits; i++) {
diff --git a/nn/runtime/test/fuzzing/operation_signatures/Normalization.cpp b/nn/runtime/test/fuzzing/operation_signatures/Normalization.cpp
index 89ba36040..5d1da4b5b 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/Normalization.cpp
+++ b/nn/runtime/test/fuzzing/operation_signatures/Normalization.cpp
@@ -24,7 +24,7 @@ static void softmaxConstructor(TestOperandType dataType, uint32_t rank, RandomOp
     sameDimensionOpConstructor(dataType, rank, op);
     // Generate value for "axis" parameter.
     if (op->inputs.size() > 2) {
-        op->inputs[2]->setScalarValue<int32_t>(getUniform<int32_t>(-rank, rank - 1));
+        op->inputs[2]->setScalarValue<int32_t>(getRandomAxis(rank));
     }
 }
 
@@ -92,7 +92,7 @@ static void l2normConstructor(TestOperandType dataType, uint32_t rank, RandomOpe
     sameDimensionOpConstructor(dataType, rank, op);
     // Generate value for "axis" parameter.
     if (op->inputs.size() > 1) {
-        op->inputs[1]->setScalarValue<int32_t>(getUniform<int32_t>(-rank, rank - 1));
+        op->inputs[1]->setScalarValue<int32_t>(getRandomAxis(rank));
     }
     // L2_NORMALIZATION may produce NaN output values with all zero inputs. We should not connect
     // the output tensor to the input of another operation.
@@ -160,7 +160,7 @@ static void localResponseNormConstructor(TestOperandType dataType, uint32_t rank
     sameDimensionOpConstructor(dataType, rank, op);
     // Generate value for "axis" parameter.
     if (op->inputs.size() > 5) {
-        op->inputs[5]->setScalarValue<int32_t>(getUniform<int32_t>(-rank, rank - 1));
+        op->inputs[5]->setScalarValue<int32_t>(getRandomAxis(rank));
     }
 }
 
diff --git a/nn/runtime/test/fuzzing/operation_signatures/OperationSignatureUtils.h b/nn/runtime/test/fuzzing/operation_signatures/OperationSignatureUtils.h
index 74a5ae4e3..8fa93327f 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/OperationSignatureUtils.h
+++ b/nn/runtime/test/fuzzing/operation_signatures/OperationSignatureUtils.h
@@ -165,6 +165,16 @@ inline void uniformFinalizer(RandomOperand* op) {
     }
 }
 
+// Get a random value between [-rank, rank) for the "axis" parameter of NNAPI operations.
+inline int32_t getRandomAxis(int32_t rank) {
+    return getUniform(-rank, rank - 1);
+}
+
+// Convert a potentially negative axis index to the equivalent positive axis index.
+inline int32_t toPositiveAxis(int32_t axis, int32_t rank) {
+    return axis >= 0 ? axis : axis + rank;
+}
+
 // A helper struct for DEFINE_OPERATION_SIGNATURE macro.
 struct OperationSignatureHelper {
     std::string name;
diff --git a/nn/runtime/test/fuzzing/operation_signatures/Reduce.cpp b/nn/runtime/test/fuzzing/operation_signatures/Reduce.cpp
index a3bad3578..994f086f6 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/Reduce.cpp
+++ b/nn/runtime/test/fuzzing/operation_signatures/Reduce.cpp
@@ -31,9 +31,9 @@ static void reduceOpConstructor(TestOperandType, uint32_t rank, RandomOperation*
     op->inputs[1]->dimensions = {numAxis};
     op->inputs[1]->resizeBuffer<int32_t>(numAxis);
     for (uint32_t i = 0; i < numAxis; i++) {
-        int32_t dim = getUniform<int32_t>(-rank, rank - 1);
+        int32_t dim = getRandomAxis(rank);
         op->inputs[1]->value<int32_t>(i) = dim;
-        reduce[dim < 0 ? dim + rank : dim] = true;
+        reduce[toPositiveAxis(dim, rank)] = true;
     }
 
     // This scalar may have two types: in MEAN it is INT32, in REDUCE_* it is BOOL
@@ -103,10 +103,10 @@ static void singleAxisReduceOpConstructor(TestOperandType, uint32_t rank, Random
     setFreeDimensions(op->inputs[0], rank);
     // "axis" must be in the range [-rank, rank).
     // Negative "axis" is used to specify axis from the end.
-    int32_t axis = getUniform<int32_t>(-rank, rank - 1);
+    int32_t axis = getRandomAxis(rank);
     op->inputs[1]->setScalarValue<int32_t>(axis);
     for (uint32_t i = 0; i < rank; i++) {
-        if (i != static_cast<uint32_t>(axis) && i != axis + rank) {
+        if (i != static_cast<uint32_t>(toPositiveAxis(axis, rank))) {
             op->outputs[0]->dimensions.emplace_back(op->inputs[0]->dimensions[i]);
         }
     }
diff --git a/nn/runtime/test/fuzzing/operation_signatures/Reshape.cpp b/nn/runtime/test/fuzzing/operation_signatures/Reshape.cpp
index b50af38de..63423c0eb 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/Reshape.cpp
+++ b/nn/runtime/test/fuzzing/operation_signatures/Reshape.cpp
@@ -425,10 +425,10 @@ static void channelShuffleConstructor(TestOperandType dataType, uint32_t rank,
                                       RandomOperation* op) {
     sameShapeOpConstructor(dataType, rank, op);
     // The number of groups must be a divisor of the target axis size.
-    int32_t axis = getUniform<int32_t>(-rank, rank - 1);
+    int32_t axis = getRandomAxis(rank);
     op->inputs[2]->setScalarValue<int32_t>(axis);
     int32_t numGroups = op->inputs[1]->value<int32_t>();
-    if (axis < 0) axis += rank;
+    axis = toPositiveAxis(axis, rank);
     (op->inputs[0]->dimensions[axis] % numGroups).setEqual(0);
 }
 
@@ -519,9 +519,9 @@ DEFINE_SQUEEZE_SIGNATURE(V1_3, TestOperandType::TENSOR_QUANT8_ASYMM_SIGNED);
 
 static void expandDimsConstructor(TestOperandType, uint32_t rank, RandomOperation* op) {
     // Generate values for the "axis" tensor.
-    int32_t axis = getUniform<int32_t>(-rank - 1, rank);
+    int32_t axis = getRandomAxis(rank + 1);
     op->inputs[1]->setScalarValue<int32_t>(axis);
-    if (axis < 0) axis += rank + 1;
+    if (axis < 0) axis += static_cast<int32_t>(rank + 1);
 
     setFreeDimensions(op->inputs[0], rank);
     for (uint32_t i = 0; i < rank; i++) {
diff --git a/nn/runtime/test/fuzzing/operation_signatures/Selection.cpp b/nn/runtime/test/fuzzing/operation_signatures/Selection.cpp
index ffcb58ede..515fc0c03 100644
--- a/nn/runtime/test/fuzzing/operation_signatures/Selection.cpp
+++ b/nn/runtime/test/fuzzing/operation_signatures/Selection.cpp
@@ -114,9 +114,9 @@ DEFINE_OPERATION_SIGNATURE(HASHTABLE_LOOKUP_V1_0){
 
 static void gatherConstructor(TestOperandType, uint32_t rank, RandomOperation* op) {
     // Generate value for "axis" scalar.
-    int32_t axis = getUniform<int32_t>(-rank, rank - 1);
+    int32_t axis = getRandomAxis(rank);
     op->inputs[1]->setScalarValue<int32_t>(axis);
-    if (axis < 0) axis += rank;
+    axis = toPositiveAxis(axis, rank);
 
     // Set dimensions for input and indices tensor.
     uint32_t indRank = getUniform<uint32_t>(1, 5);
@@ -137,7 +137,7 @@ static void gatherConstructor(TestOperandType, uint32_t rank, RandomOperation* o
 
 static void gatherFinalizer(RandomOperation* op) {
     int32_t axis = op->inputs[1]->value<int32_t>();
-    if (axis < 0) axis += op->inputs[0]->dimensions.size();
+    axis = toPositiveAxis(axis, op->inputs[0]->dimensions.size());
     uint32_t dimValue = op->inputs[0]->dimensions[axis].getValue();
     uint32_t numElements = op->inputs[2]->getNumberOfElements();
     for (uint32_t i = 0; i < numElements; i++) {