CVE-2022-23576: Integer Overflow or Wraparound (CWE-190)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `OpLevelCostEstimator::CalculateOutputSize` is vulnerable to an integer overflow if an attacker can create an operation which would involve tensors with large enough number of elements. We can have a large enough number of dimensions in `output_shape.dim()` or just a small number of dimensions being large enough to cause an overflow in the multiplication. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
CWE description: The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
CVSS v3 base score: 6.5
Published: 2022-02-04T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/core/grappler/costs/BUILD, tensorflow/core/grappler/costs/op_level_cost_estimator.cc
Commit date: 2021-11-09 14:11:04 -0800
Commit message: Prevent integer overflow in `OpLevelCostEstimator::CalculateOutputSize`.

In order to not change the API, we return a negative value in case of overflow. A better fix is to change the API to return a status instead.

PiperOrigin-RevId: 408701427
Change-Id: Idf31e7f0bf18ca824d084fdd355e1f653f145c20

Diff: diff --git a/tensorflow/core/grappler/costs/BUILD b/tensorflow/core/grappler/costs/BUILD
index 66e0c6cd7c5f1..b267a0f2787f3 100644
--- a/tensorflow/core/grappler/costs/BUILD
+++ b/tensorflow/core/grappler/costs/BUILD
@@ -338,10 +338,11 @@ cc_library(
         "//third_party/eigen3",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/clusters:utils",
+        "//tensorflow/core/util:overflow",
     ] + tf_protos_grappler(),
 )
 
 tf_cc_test(
     name = "op_level_cost_estimator_test",
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index 10acbf54595da..5e328542ab145 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -25,10 +25,11 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/grappler/clusters/utils.h"
 #include "tensorflow/core/grappler/costs/op_context.h"
 #include "tensorflow/core/grappler/costs/utils.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 namespace grappler {
 
 // TODO(dyoon): update op to Predict method map for TF ops with V2 or V3 suffix.
@@ -1605,11 +1606,18 @@ int64_t OpLevelCostEstimator::CalculateOutputSize(const OpInfo& op_info,
     int64_t output_size = DataTypeSize(BaseType(dt));
     int num_dims = std::max(1, original_output_shape.dim_size());
     auto output_shape = MaybeGetMinimumShape(original_output_shape, num_dims,
                                              found_unknown_shapes);
     for (const auto& dim : output_shape.dim()) {
-      output_size *= dim.size();
+      int64_t new_output_size =
+          MultiplyWithoutOverflow(output_size, dim.size());
+      if (new_output_size < 0) {
+        VLOG(1) << "Overflow encountered when estimating cost, multiplying "
+                << output_size << " with " << dim.size();
+        return -1;
+      }
+      output_size = new_output_size;
     }
     total_output_size += output_size;
     VLOG(1) << "Output Size: " << output_size
             << " Total Output Size:" << total_output_size;
   }

Vulnerable code:       output_size *= dim.size();
Fixed code:         "//tensorflow/core/util:overflow",
#include "tensorflow/core/util/overflow.h"
      int64_t new_output_size =
          MultiplyWithoutOverflow(output_size, dim.size());
      if (new_output_size < 0) {
        VLOG(1) << "Overflow encountered when estimating cost, multiplying "
                << output_size << " with " << dim.size();
        return -1;
      }
      output_size = new_output_size;