CVE-2021-41205: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. In affected versions the shape inference functions for the `QuantizeAndDequantizeV*` operations can trigger a read outside of bounds of heap allocated array. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.1
Published: 2021-11-05T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/ops/array_ops.cc, tensorflow/core/ops/array_ops_test.cc
Commit date: 2021-10-11 19:07:32 -0700
Commit message: Address QuantizeAndDequantizeV* heap oob. Added additional checks for the 'axis' attribute.

PiperOrigin-RevId: 402446942
Change-Id: Id2f6b82e4e740d0550329be02621c46466b5a5b9

Diff: diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 775a4cc469664..64bd4f3847854 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -2861,11 +2861,14 @@ REGISTER_OP("QuantizeAndDequantizeV2")
       TF_RETURN_IF_ERROR(c->GetAttr("axis", &axis));
       const int minmax_rank = (axis == -1) ? 0 : 1;
       ShapeHandle minmax;
       TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
       TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
-      if (axis != -1) {
+      if (axis < -1) {
+        return errors::InvalidArgument("axis should be at least -1, got ",
+                                       axis);
+      } else if (axis != -1) {
         ShapeHandle input;
         TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
         DimensionHandle depth;
         TF_RETURN_IF_ERROR(
             c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
@@ -2893,11 +2896,14 @@ REGISTER_OP("QuantizeAndDequantizeV4")
       TF_RETURN_IF_ERROR(c->GetAttr("axis", &axis));
       const int minmax_rank = (axis == -1) ? 0 : 1;
       ShapeHandle minmax;
       TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
       TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
-      if (axis != -1) {
+      if (axis < -1) {
+        return errors::InvalidArgument("axis should be at least -1, got ",
+                                       axis);
+      } else if (axis != -1) {
         ShapeHandle input;
         TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
         DimensionHandle depth;
         TF_RETURN_IF_ERROR(
             c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
@@ -2921,11 +2927,14 @@ REGISTER_OP("QuantizeAndDequantizeV4Grad")
       TF_RETURN_IF_ERROR(c->GetAttr("axis", &axis));
       const int minmax_rank = (axis == -1) ? 0 : 1;
       ShapeHandle minmax;
       TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
       TF_RETURN_IF_ERROR(c->Merge(c->input(3), minmax, &minmax));
-      if (axis != -1) {
+      if (axis < -1) {
+        return errors::InvalidArgument("axis should be at least -1, got ",
+                                       axis);
+      } else if (axis != -1) {
         ShapeHandle input;
         TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
         DimensionHandle depth;
         TF_RETURN_IF_ERROR(
             c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
@@ -2954,11 +2963,14 @@ REGISTER_OP("QuantizeAndDequantizeV3")
       TF_RETURN_IF_ERROR(c->GetAttr("axis", &axis));
       const int minmax_rank = (axis == -1) ? 0 : 1;
       ShapeHandle minmax;
       TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
       TF_RETURN_IF_ERROR(c->Merge(c->input(2), minmax, &minmax));
-      if (axis != -1) {
+      if (axis < -1) {
+        return errors::InvalidArgument("axis should be at least -1, got ",
+                                       axis);
+      } else if (axis != -1) {
         ShapeHandle input;
         TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
         DimensionHandle depth;
         TF_RETURN_IF_ERROR(
             c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
diff --git a/tensorflow/core/ops/array_ops_test.cc b/tensorflow/core/ops/array_ops_test.cc
index b84ebd10123d4..3de484f70010c 100644
--- a/tensorflow/core/ops/array_ops_test.cc
+++ b/tensorflow/core/ops/array_ops_test.cc
@@ -1372,10 +1372,12 @@ TEST(ArrayOpsTest, QuantizeAndDequantizeV2_ShapeFn) {
 
   INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[1];[]");
   INFER_ERROR("Shapes must be equal rank, but are 1 and 0", op,
               "[1,2,?,4,5];[];[1]");
   INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[1];[1]");
+  (*op.node_def.mutable_attr())["axis"].set_i(-2);
+  INFER_ERROR("axis should be at least -1, got -2", op, "?;?;?");
 }
 
 TEST(ArrayOpsTest, SpaceToBatch_ShapeFn) {
   ShapeInferenceTestOp op("SpaceToBatch");
   op.input_tensors.resize(2);

Vulnerable code:       if (axis != -1) {
      if (axis != -1) {
      if (axis != -1) {
      if (axis != -1) {
Fixed code:       if (axis < -1) {
        return errors::InvalidArgument("axis should be at least -1, got ",
                                       axis);
      } else if (axis != -1) {
      if (axis < -1) {
        return errors::InvalidArgument("axis should be at least -1, got ",
                                       axis);
      } else if (axis != -1) {
      if (axis < -1) {
        return errors::InvalidArgument("axis should be at least -1, got ",
                                       axis);
      } else if (axis != -1) {
      if (axis < -1) {
        return errors::InvalidArgument("axis should be at least -1, got ",
                                       axis);
      } else if (axis != -1) {
  (*op.node_def.mutable_attr())["axis"].set_i(-2);
  INFER_ERROR("axis should be at least -1, got -2", op, "?;?;?");