CVE-2020-26267: Out-of-bounds Read (CWE-125)
value: In affected versions of TensorFlow the tf.raw_ops.DataFormatVecPermute API does not validate the src_format and dst_format attributes. The code assumes that these two arguments define a permutation of NHWC. This can result in uninitialized memory accesses, read outside of bounds and even crashes. This is fixed in versions 1.15.5, 2.0.4, 2.1.3, 2.2.2, 2.3.2, and 2.4.0.
CWE description: The product reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.8
Published: 2020-12-10T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/data_format_ops.cc, tensorflow/python/ops/nn_test.py
Commit date: 2020-12-07 11:25:52 -0800
Commit message: Validate that `DataFormat*` attributes form a permutation.

The `src_format` and `dst_format` attributes for the `DataFormatDimMap` and `DataFormatVecPermute` raw ops are supposed to determine a permutation. However, this was not validated and could result in unitialized memory accesses as well as writes outside of bounds and potential crashes.

While here, we also test that the format attributes have the needed length, add tests for all validation failure cases, remove unnecessary calls to `strings::StrCat`, and fix a few grammar errors.

This will be cherry-picked on the supported release branches.

PiperOrigin-RevId: 346135579
Change-Id: I1c76392382c89ad8f072d5bc93d70669851eb404

Diff: diff --git a/tensorflow/core/kernels/data_format_ops.cc b/tensorflow/core/kernels/data_format_ops.cc
index b52d4d6c8882..00724b630ee9 100644
--- a/tensorflow/core/kernels/data_format_ops.cc
+++ b/tensorflow/core/kernels/data_format_ops.cc
@@ -16,39 +16,79 @@ limitations under the License.
 // See docs in ../ops/nn_ops.cc.
 
 #define EIGEN_USE_THREADS
 
 #include "tensorflow/core/kernels/data_format_ops.h"
+
+#include <map>
+
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
 
 typedef Eigen::ThreadPoolDevice CPUDevice;
 typedef Eigen::GpuDevice GPUDevice;
 
+// Ensure that `src` and `dst` define a valid permutation.
+// Ops defined in this file assume that user specifies a permutation via two
+// string attributes. This check validates that these attributes properly define
+// it to prevent security vulnerabilities.
+static bool IsValidPermutation(const std::string& src, const std::string& dst) {
+  if (src.size() != dst.size()) {
+    return false;
+  }
+
+  std::map<char, bool> characters;
+
+  // Every character in `src` must be present only once
+  for (const auto c : src) {
+    if (characters[c]) {
+      return false;
+    }
+    characters[c] = true;
+  }
+
+  // Every character in `dst` must show up in `src` exactly once
+  for (const auto c : dst) {
+    if (!characters[c]) {
+      return false;
+    }
+    characters[c] = false;
+  }
+
+  // At this point, characters[] has been switched to true and false exactly
+  // once for all character in `src` (and `dst`) so we have a valid permutation
+  return true;
+}
+
 template <typename Device, typename T>
 class DataFormatDimMapOp : public OpKernel {
  public:
   explicit DataFormatDimMapOp(OpKernelConstruction* context)
       : OpKernel(context) {
     string src_format;
     OP_REQUIRES_OK(context, context->GetAttr("src_format", &src_format));
     string dst_format;
     OP_REQUIRES_OK(context, context->GetAttr("dst_format", &dst_format));
     OP_REQUIRES(context, src_format.size() == 4 || src_format.size() == 5,
-                errors::InvalidArgument(strings::StrCat(
-                    "Source format must of length 4 or 5, received "
+                errors::InvalidArgument(
+                    "Source format must be of length 4 or 5, received "
                     "src_format = ",
-                    src_format)));
+                    src_format));
+    OP_REQUIRES(context, dst_format.size() == 4 || dst_format.size() == 5,
+                errors::InvalidArgument("Destination format must be of length "
+                                        "4 or 5, received dst_format = ",
+                                        dst_format));
     OP_REQUIRES(
-        context, dst_format.size() == 4 || dst_format.size() == 5,
-        errors::InvalidArgument(strings::StrCat(
-            "Destination format must of length 4 or 5, received dst_format = ",
-            dst_format)));
+        context, IsValidPermutation(src_format, dst_format),
+        errors::InvalidArgument(
+            "Destination and source format must determine a permutation, got ",
+            src_format, " and ", dst_format));
     dst_idx_ = Tensor(DT_INT32, {static_cast<int64>(src_format.size())});
     for (int i = 0; i < src_format.size(); ++i) {
       for (int j = 0; j < dst_format.size(); ++j) {
         if (dst_format[j] == src_format[i]) {
           dst_idx_.vec<int>()(i) = j;
@@ -76,12 +116,26 @@ class DataFormatVecPermuteOp : public OpKernel {
  public:
   explicit DataFormatVecPermuteOp(OpKernelConstruction* context)
       : OpKernel(context) {
     string src_format;
     OP_REQUIRES_OK(context, context->GetAttr("src_format", &src_format));
+    OP_REQUIRES(context, src_format.size() == 4 || src_format.size() == 5,
+                errors::InvalidArgument(
+                    "Source format must be of length 4 or 5, received "
+                    "src_format = ",
+                    src_format));
     string dst_format;
     OP_REQUIRES_OK(context, context->GetAttr("dst_format", &dst_format));
+    OP_REQUIRES(context, dst_format.size() == 4 || dst_format.size() == 5,
+                errors::InvalidArgument("Destination format must be of length "
+                                        "4 or 5, received dst_format = ",
+                                        dst_format));
+    OP_REQUIRES(
+        context, IsValidPermutation(src_format, dst_format),
+        errors::InvalidArgument(
+            "Destination and source format must determine a permutation, got ",
+            src_format, " and ", dst_format));
     src_format_ = src_format;
     dst_format_ = dst_format;
   }
 
   void Compute(OpKernelContext* context) override {
@@ -125,10 +179,14 @@ class DataFormatVecPermuteOp : public OpKernel {
             [](const char dim) { return dim != 'H' && dim != 'W'; });
         format_str->erase(new_end, format_str->end());
       };
       keep_only_spatial_dimensions(&src_format_str);
       keep_only_spatial_dimensions(&dst_format_str);
+      OP_REQUIRES(context,
+                  src_format_str.size() == 2 && dst_format_str.size() == 2,
+                  errors::InvalidArgument(
+                      "Format specifier must contain H and W for 2D case"));
     }
     ComputeDstIndex(src_format_str, dst_format_str, input.dims(), &dst_idx);
 
     functor::DataFormatVecPermute<Device, T>()(context->eigen_device<Device>(),
                                                input.flat<T>(),
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index 851bfcb66de3..aaf2f77fb297 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -25,10 +25,11 @@ import numpy as np
 from six.moves import xrange  # pylint: disable=redefined-builtin
 
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gradient_checker
@@ -1258,19 +1259,60 @@ class DataFormatDimMapTest(test_lib.TestCase):
     y = nn_ops.data_format_dim_map(x, src_format="NDHWC", dst_format="WHDCN")
     with test_util.use_gpu():
       y_val = self.evaluate(y)
       self.assertAllEqual(y_val, y_val_expected)
 
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
   def testArbitraryASCII(self):
     x_val = [-4, -3, -2, -1, 0, 1, 2, 3]
     y_val_expected = [3, 2, 1, 0, 3, 2, 1, 0]
     x = constant_op.constant(x_val)
     y = nn_ops.data_format_dim_map(x, src_format="qwer", dst_format="rewq")
     with test_util.use_gpu():
       y_val = self.evaluate(y)
       self.assertAllEqual(y_val, y_val_expected)
 
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testInvalidLength(self):
+    x = [-4, -3, -2, -1, 0, 1, 2, 3]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Source format must be of length 4 or 5"):
+      op = nn_ops.data_format_dim_map(
+          x, src_format="12345678", dst_format="87654321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testDuplicateSrc(self):
+    x = [-4, -3, -2, -1, 0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_dim_map(x, src_format="1233", dst_format="4321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testDuplicateDst(self):
+    x = [-4, -3, -2, -1, 0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_dim_map(x, src_format="1234", dst_format="3321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testExtraSpecifiers(self):
+    x = [-4, -3, -2, -1, 0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_dim_map(x, src_format="1234", dst_format="5321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
 
 class DataFormatVectorPermuteTest(test_lib.TestCase):
 
   def testNHWCToNCHW(self):
     x_val = [7, 4, 9, 3]
@@ -1368,10 +1410,64 @@ class DataFormatVectorPermuteTest(test_lib.TestCase):
     y = nn_ops.data_format_vec_permute(x, src_format="NCHW", dst_format="NHWC")
     with test_util.use_gpu():
       y_val = self.evaluate(y)
       self.assertAllEqual(y_val, [[7, 4], [4, 5], [5, 1], [9, 3]])
 
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testInvalidLength(self):
+    x = [0, 1, 2, 3]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Source format must be of length 4 or 5"):
+      op = nn_ops.data_format_vec_permute(
+          x, src_format="12345678", dst_format="87654321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testDuplicateSrc(self):
+    x = [0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_vec_permute(
+          x, src_format="1233", dst_format="4321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testDuplicateDst(self):
+    x = [0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_vec_permute(
+          x, src_format="1234", dst_format="3321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def testExtraSpecifiers(self):
+    x = [0, 1, 2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Destination and source format must determine a permutation"):
+      op = nn_ops.data_format_vec_permute(
+          x, src_format="1234", dst_format="5321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
+  @test_util.disable_xla("XLA catches the error and rethrows as different one")
+  def test2DNoWH(self):
+    x = [[0, 1], [2, 3]]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Format specifier must contain H and W for 2D case"):
+      op = nn_ops.data_format_vec_permute(
+          x, src_format="1234", dst_format="4321")
+      with test_util.use_gpu():
+        self.evaluate(op)
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class AvgPoolTest(test_lib.TestCase):
 
   def test1DTensor(self):

Vulnerable code:                 errors::InvalidArgument(strings::StrCat(
                    "Source format must of length 4 or 5, received "
                    src_format)));
        context, dst_format.size() == 4 || dst_format.size() == 5,
        errors::InvalidArgument(strings::StrCat(
            "Destination format must of length 4 or 5, received dst_format = ",
            dst_format)));
Fixed code: 
#include <map>

#include "tensorflow/core/platform/errors.h"
// Ensure that `src` and `dst` define a valid permutation.
// Ops defined in this file assume that user specifies a permutation via two
// string attributes. This check validates that these attributes properly define
// it to prevent security vulnerabilities.
static bool IsValidPermutation(const std::string& src, const std::string& dst) {
  if (src.size() != dst.size()) {
    return false;
  }

  std::map<char, bool> characters;

  // Every character in `src` must be present only once
  for (const auto c : src) {
    if (characters[c]) {
      return false;
    }
    characters[c] = true;
  }

  // Every character in `dst` must show up in `src` exactly once
  for (const auto c : dst) {
    if (!characters[c]) {
      return false;
    }
    characters[c] = false;
  }

  // At this point, characters[] has been switched to true and false exactly
  // once for all character in `src` (and `dst`) so we have a valid permutation
  return true;
}

                errors::InvalidArgument(
                    "Source format must be of length 4 or 5, received "
                    src_format));
    OP_REQUIRES(context, dst_format.size() == 4 || dst_format.size() == 5,
                errors::InvalidArgument("Destination format must be of length "
                                        "4 or 5, received dst_format = ",
                                        dst_format));
        context, IsValidPermutation(src_format, dst_format),
        errors::InvalidArgument(
            "Destination and source format must determine a permutation, got ",
            src_format, " and ", dst_format));
    OP_REQUIRES(context, src_format.size() == 4 || src_format.size() == 5,
                errors::InvalidArgument(
                    "Source format must be of length 4 or 5, received "
                    "src_format = ",
                    src_format));
    OP_REQUIRES(context, dst_format.size() == 4 || dst_format.size() == 5,
                errors::InvalidArgument("Destination format must be of length "
                                        "4 or 5, received dst_format = ",
                                        dst_format));
    OP_REQUIRES(
        context, IsValidPermutation(src_format, dst_format),
        errors::InvalidArgument(
            "Destination and source format must determine a permutation, got ",
            src_format, " and ", dst_format));
      OP_REQUIRES(context,
                  src_format_str.size() == 2 && dst_format_str.size() == 2,
                  errors::InvalidArgument(
                      "Format specifier must contain H and W for 2D case"));
from tensorflow.python.framework import errors
  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testInvalidLength(self):
    x = [-4, -3, -2, -1, 0, 1, 2, 3]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Source format must be of length 4 or 5"):
      op = nn_ops.data_format_dim_map(
          x, src_format="12345678", dst_format="87654321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testDuplicateSrc(self):
    x = [-4, -3, -2, -1, 0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_dim_map(x, src_format="1233", dst_format="4321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testDuplicateDst(self):
    x = [-4, -3, -2, -1, 0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_dim_map(x, src_format="1234", dst_format="3321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testExtraSpecifiers(self):
    x = [-4, -3, -2, -1, 0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_dim_map(x, src_format="1234", dst_format="5321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testInvalidLength(self):
    x = [0, 1, 2, 3]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Source format must be of length 4 or 5"):
      op = nn_ops.data_format_vec_permute(
          x, src_format="12345678", dst_format="87654321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testDuplicateSrc(self):
    x = [0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_vec_permute(
          x, src_format="1233", dst_format="4321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testDuplicateDst(self):
    x = [0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_vec_permute(
          x, src_format="1234", dst_format="3321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def testExtraSpecifiers(self):
    x = [0, 1, 2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Destination and source format must determine a permutation"):
      op = nn_ops.data_format_vec_permute(
          x, src_format="1234", dst_format="5321")
      with test_util.use_gpu():
        self.evaluate(op)

  @test_util.disable_xla("XLA catches the error and rethrows as different one")
  def test2DNoWH(self):
    x = [[0, 1], [2, 3]]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Format specifier must contain H and W for 2D case"):
      op = nn_ops.data_format_vec_permute(
          x, src_format="1234", dst_format="4321")
      with test_util.use_gpu():
        self.evaluate(op)