CVE-2023-25675: Incorrect Comparison (CWE-697)
value: TensorFlow is an open source machine learning platform. When running versions prior to 2.12.0 and 2.11.1 with XLA, `tf.raw_ops.Bincount` segfaults when given a parameter `weights` that is neither the same shape as parameter `arr` nor a length-0 tensor. A fix is included in TensorFlow 2.12.0 and 2.11.1.
CWE description: The product compares two entities in a security-relevant context, but the comparison is incorrect, which may lead to resultant weaknesses.
CVSS v3 base score: 7.5
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/compiler/tests/BUILD, tensorflow/compiler/tests/bincount_op_test.py, tensorflow/compiler/tf2xla/kernels/bincount_op.cc
Commit date: 2023-02-01 21:52:11 -0800
Commit message: [Tensorflow] Fix security vulnerability with DenseBincountOp

PiperOrigin-RevId: 506514542

Diff: diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index 772962844a540..c078689a743cf 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -2415,5 +2415,21 @@ tf_xla_py_test(
         "//tensorflow/python:framework",
         "//tensorflow/python:platform_test",
         "//tensorflow/python:training",
     ],
 )
+
+tf_xla_py_test(
+    name = "bincount_op_test",
+    size = "small",
+    srcs = ["bincount_op_test.py"],
+    enable_mlir_bridge = False,
+    python_version = "PY3",
+    shard_count = 10,
+    tags = [
+        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
+    ],
+    deps = [
+        ":xla_test",
+        "//tensorflow/python:platform_test",
+    ],
+)
diff --git a/tensorflow/compiler/tests/bincount_op_test.py b/tensorflow/compiler/tests/bincount_op_test.py
new file mode 100644
index 0000000000000..79e8a7e91b8a0
--- /dev/null
+++ b/tensorflow/compiler/tests/bincount_op_test.py
@@ -0,0 +1,40 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for bincount using the XLA JIT."""
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.platform import googletest
+
+
+class BincountTest(xla_test.XLATestCase):
+
+  def testInputRank0(self):
+    with self.session():
+      with self.test_scope():
+        bincount = gen_math_ops.bincount(arr=6, size=804, weights=[52, 351])
+
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          (
+              "`weights` must be the same shape as `arr` or a length-0"
+              " `Tensor`, in which case it acts as all weights equal to 1."
+          ),
+      ):
+        self.evaluate(bincount)
+
+
+if __name__ == "__main__":
+  googletest.main()
diff --git a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
index ba95c77f61250..5cac4d98c50a9 100644
--- a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
@@ -60,25 +60,19 @@ class DenseBincountOp : public XlaOpKernel {
     auto zero = xla::Zero(ctx->builder(), dtype);
     auto one = xla::One(ctx->builder(), dtype);
     StatusOr<xla::Shape> input_shape_or = ctx->builder()->GetShape(input);
     OP_REQUIRES_OK(ctx, input_shape_or.status());
     auto input_shape = input_shape_or.value();
-    auto size = input_shape.dimensions(0);
 
-    if (!size) {
-      output = xla::Broadcast(zero, {output_size});
-      ctx->SetOutput(0, output);
-      return;
-    }
     auto rank = input_shape.rank();
 
     OP_REQUIRES(ctx, rank <= 2,
                 errors::InvalidArgument(
                     "Shape must be at most rank 2 but is rank ", rank));
-
     xla::XlaOp weights = ctx->Input(2);
     StatusOr<xla::Shape> weights_shape_or = ctx->builder()->GetShape(weights);
+
     OP_REQUIRES_OK(ctx, weights_shape_or.status());
 
     auto weights_shape = weights_shape_or.value();
     OP_REQUIRES(ctx,
                 xla::ShapeUtil::CompatibleIgnoringElementType(weights_shape,
@@ -89,15 +83,24 @@ class DenseBincountOp : public XlaOpKernel {
                     "`weights` must be the same shape as `arr` or a length-0 "
                     "`Tensor`, in which case it acts as all weights equal to "
                     "1. Received ",
                     weights_shape.DebugString()));
 
+    auto size = input_shape.dimensions(0);
+
+    if (!size) {
+      output = xla::Broadcast(zero, {output_size});
+      ctx->SetOutput(0, output);
+      return;
+    }
+
     auto weights_size = weights_shape.dimensions(0);
     bool has_weights = false;
     if (weights_size) {
       has_weights = true;
     }
+
     xla::Shape output_shape = xla::ShapeUtil::MakeShape(dtype, {output_size});
     xla::ScatterDimensionNumbers scatter_dnums;
     scatter_dnums.set_index_vector_dim(1);
     scatter_dnums.add_inserted_window_dims(0);
     scatter_dnums.add_scatter_dims_to_operand_dims(0);

Vulnerable code:     auto size = input_shape.dimensions(0);
    if (!size) {
      output = xla::Broadcast(zero, {output_size});
      ctx->SetOutput(0, output);
      return;
    }

Fixed code: 
tf_xla_py_test(
    name = "bincount_op_test",
    size = "small",
    srcs = ["bincount_op_test.py"],
    enable_mlir_bridge = False,
    python_version = "PY3",
    shard_count = 10,
    tags = [
        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
    ],
    deps = [
        ":xla_test",
        "//tensorflow/python:platform_test",
    ],
)
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for bincount using the XLA JIT."""
from tensorflow.compiler.tests import xla_test
from tensorflow.python.framework import errors
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.platform import googletest


class BincountTest(xla_test.XLATestCase):

  def testInputRank0(self):
    with self.session():
      with self.test_scope():
        bincount = gen_math_ops.bincount(arr=6, size=804, weights=[52, 351])

      with self.assertRaisesRegex(
          errors.InvalidArgumentError,
          (
              "`weights` must be the same shape as `arr` or a length-0"
              " `Tensor`, in which case it acts as all weights equal to 1."
          ),
      ):
        self.evaluate(bincount)


if __name__ == "__main__":
  googletest.main()

    auto size = input_shape.dimensions(0);

    if (!size) {
      output = xla::Broadcast(zero, {output_size});
      ctx->SetOutput(0, output);
      return;
    }