Vulnerability: Abort caused by allocating a vector that is too large in Tensorflow ### Impact
During shape inference, TensorFlow can [allocate a large vector](https://github.com/tensorflow/tensorflow/blob/a1320ec1eac186da1d03f033109191f715b2b130/tensorflow/core/framework/shape_inference.cc#L788-L790) based on a value from a tensor controlled by the user:

```cc
  const auto num_dims = Value(shape_dim);
  std::vector<DimensionHandle> dims;
  dims.reserve(num_dims);
``` 
  
### Patches           
We have patched the issue in GitHub commit [1361fb7e29449629e1df94d44e0427ebec8c83c7](https://github.com/tensorflow/tensorflow/commit/1361fb7e29449629e1df94d44e0427ebec8c83c7).

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.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.
Vulnerability ID: GHSA-627q-g293-49q7
CWE: CWE-400
Severity score: 6.5
Project: https://github.com/tensorflow/tensorflow
Published: 2022-02-07
File: tensorflow/core/framework/shape_inference.cc
Commit message: Fix abort caused by allocating a too large vector.

We need to make sure that the number of dimensions in a shape is within limits.

PiperOrigin-RevId: 408997911
Change-Id: If59e1c23f2ec9c2d4ff4d8632fd62b2a7773a4eb
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 limitations under the License. ==============================================================================*/ #include "tensorflow/core/framework/shape_inference.h"  #include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/full_type_util.h" #include "tensorflow/core/framework/node_def.pb.h" Status InferenceContext::InternalMakeShapeFromTensor(       return ReturnUnknownShape(out);     }     const auto num_dims = Value(shape_dim);     std::vector<DimensionHandle> dims;     dims.reserve(num_dims);     for (int i = 0; i < num_dims; i++) dims.push_back(UnknownDim());

Fixed version (after fix):
 limitations under the License. ==============================================================================*/ #include "tensorflow/core/framework/shape_inference.h"  #include <cstdint>  #include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/full_type_util.h" #include "tensorflow/core/framework/node_def.pb.h" Status InferenceContext::InternalMakeShapeFromTensor(       return ReturnUnknownShape(out);     }     const auto num_dims = Value(shape_dim);     // TODO(mihaimaruseac): Should be `TensorShape::MaxDimensions()` as we are     // not able to materialize shapes with more than this number of dimensions     // but then shape inference would fail for operations such as     // `tf.range`/`tf.ones`, etc. where the shape is not really materialized,     // only used during the inference. Hence, just prevent doing a `reserve`     // with a very large argument.     const int64_t max_dimensions = 1 << 20;     if (num_dims >= max_dimensions) {       return errors::Internal(           "Cannot create a tensor with ", num_dims,           " dimensions, as these would be more than maximum of ",           max_dimensions);     }     std::vector<DimensionHandle> dims;     dims.reserve(num_dims);     for (int i = 0; i < num_dims; i++) dims.push_back(UnknownDim());

Patch:
patch: @@ -14,6 +14,8 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/framework/shape_inference.h"
 
+#include <cstdint>
+
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/full_type_util.h"
 #include "tensorflow/core/framework/node_def.pb.h"
@@ -789,6 +791,19 @@ Status InferenceContext::InternalMakeShapeFromTensor(
       return ReturnUnknownShape(out);
     }
     const auto num_dims = Value(shape_dim);
+    // TODO(mihaimaruseac): Should be `TensorShape::MaxDimensions()` as we are
+    // not able to materialize shapes with more than this number of dimensions
+    // but then shape inference would fail for operations such as
+    // `tf.range`/`tf.ones`, etc. where the shape is not really materialized,
+    // only used during the inference. Hence, just prevent doing a `reserve`
+    // with a very large argument.
+    const int64_t max_dimensions = 1 << 20;
+    if (num_dims >= max_dimensions) {
+      return errors::Internal(
+          "Cannot create a tensor with ", num_dims,
+          " dimensions, as these would be more than maximum of ",
+          max_dimensions);
+    }
     std::vector<DimensionHandle> dims;
     dims.reserve(num_dims);
     for (int i = 0; i < num_dims; i++) dims.push_back(UnknownDim());
path: tensorflow/core/framework/shape_inference.cc