CVE-2022-23593: Improper Check for Unusual or Exceptional Conditions (CWE-754)
value: Tensorflow is an Open Source Machine Learning Framework. The `simplifyBroadcast` function in the MLIR-TFRT infrastructure in TensorFlow is vulnerable to a segfault (hence, denial of service), if called with scalar shapes. If all shapes are scalar, then `maxRank` is 0, so we build an empty `SmallVector`. The fix will be included in TensorFlow 2.8.0. This is the only affected version.
CWE description: The product does not check or incorrectly checks for unusual or exceptional conditions that are not expected to occur frequently during day to day operation of the product.
CVSS v3 base score: 7.5
Published: 2022-02-04T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/compiler/mlir/tfrt/jit/transforms/tf_cpurt_symbolic_shape_optimization.cc, tensorflow/compiler/mlir/tfrt/python_tests/regression_tests/scalar_broadcast.mlir
Commit date: 2021-12-21 01:12:06 -0800
Commit message: Avoid Segfault for scalar shapes.

Calling tensor::FromElementsOp with an empty vector of elements and no type
causes a segfault. We need to let the FromElementsOp know which scalar type it
should have.
Also add back the DynamicBroadcastInDimOp canonicalization patterns, which
previously prevented this bug from happening.
Add a regression test that demonstrates the bug.

PiperOrigin-RevId: 417561444
Change-Id: I6d1d6cfb71aabbad6102422625a00bbe253ac95a

Diff: diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_cpurt_symbolic_shape_optimization.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_cpurt_symbolic_shape_optimization.cc
index 628141a6cd46..0a8fc42e80a7 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_cpurt_symbolic_shape_optimization.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_cpurt_symbolic_shape_optimization.cc
@@ -155,10 +155,14 @@ llvm::Optional<Value> simplifyBroadcast(ShapeComponentAnalysis& analysis,
     auto found_shape = analysis.GetValueInfo(shape.value());
     if (!found_shape) return {};
     shapes_found.push_back(*found_shape);
     maxRank = std::max(maxRank, found_shape->size());
   }
+  if (maxRank == 0) {
+    return Value(builder->create<tensor::FromElementsOp>(
+        loc, shapes[0].getType(), SmallVector<Value>()));
+  }
 
   SmallVector<const ShapeComponentAnalysis::SymbolicExpr*> joined_dimensions(
       maxRank);
   SmallVector<std::pair<Value, int64_t>> shape_and_rank_for_dim(maxRank);
   for (const auto &shape : llvm::enumerate(shapes_found)) {
diff --git a/tensorflow/compiler/mlir/tfrt/python_tests/regression_tests/scalar_broadcast.mlir b/tensorflow/compiler/mlir/tfrt/python_tests/regression_tests/scalar_broadcast.mlir
new file mode 100644
index 000000000000..33b785ee6fb1
--- /dev/null
+++ b/tensorflow/compiler/mlir/tfrt/python_tests/regression_tests/scalar_broadcast.mlir
@@ -0,0 +1,11 @@
+builtin.func @test(%V__0 : tensor<i1> { python_test_attrs.static_type = tensor<i1> }, %V__1 : tensor<f32> { python_test_attrs.static_type = tensor<f32> }, %V__2 : tensor<f32> { python_test_attrs.static_type = tensor<f32> }) -> tensor<f32> {
+  %0 = "tf.Cast"(%V__0) : (tensor<i1>) -> tensor<i1>
+  %1 = "tf.Selu"(%V__2) : (tensor<f32>) -> tensor<f32>
+  %2 = "tf.NextAfter"(%1, %V__2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
+  %3 = "tf.Elu"(%2) : (tensor<f32>) -> tensor<f32>
+  %4 = "tf.Cosh"(%3) : (tensor<f32>) -> tensor<f32>
+  %5 = "tf.Elu"(%4) : (tensor<f32>) -> tensor<f32>
+  %6 = "tf.Div"(%V__1, %5) : (tensor<f32>, tensor<f32>) -> tensor<f32>
+  %7 = "tf.Select"(%0, %6, %V__1) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
+  return %7 : tensor<f32>
+}
\ No newline at end of file

Fixed code:   if (maxRank == 0) {
    return Value(builder->create<tensor::FromElementsOp>(
        loc, shapes[0].getType(), SmallVector<Value>()));
  }
builtin.func @test(%V__0 : tensor<i1> { python_test_attrs.static_type = tensor<i1> }, %V__1 : tensor<f32> { python_test_attrs.static_type = tensor<f32> }, %V__2 : tensor<f32> { python_test_attrs.static_type = tensor<f32> }) -> tensor<f32> {
  %0 = "tf.Cast"(%V__0) : (tensor<i1>) -> tensor<i1>
  %1 = "tf.Selu"(%V__2) : (tensor<f32>) -> tensor<f32>
  %2 = "tf.NextAfter"(%1, %V__2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = "tf.Elu"(%2) : (tensor<f32>) -> tensor<f32>
  %4 = "tf.Cosh"(%3) : (tensor<f32>) -> tensor<f32>
  %5 = "tf.Elu"(%4) : (tensor<f32>) -> tensor<f32>
  %6 = "tf.Div"(%V__1, %5) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %7 = "tf.Select"(%0, %6, %V__1) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
  return %7 : tensor<f32>
}