CVE-2022-23580: Improper Validation of Specified Quantity in Input (CWE-1284)
value: Tensorflow is an Open Source Machine Learning Framework. During shape inference, TensorFlow can allocate a large vector based on a value from a tensor controlled by the user. 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 receives input that is expected to specify a quantity (such as size or length), but it does not validate or incorrectly validates that the quantity has the required properties.
CVSS v3 base score: 6.5
Published: 2022-02-04T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/framework/shape_inference.cc
Commit date: 2021-11-10 15:59:39 -0800
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

Diff: diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc
index 077837721ae4f..3316759aed5a7 100644
--- a/tensorflow/core/framework/shape_inference.cc
+++ b/tensorflow/core/framework/shape_inference.cc
@@ -12,10 +12,12 @@ 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.
 ==============================================================================*/
 #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"
 #include "tensorflow/core/framework/op_def.pb.h"
 #include "tensorflow/core/framework/partial_tensor_shape.h"
@@ -787,10 +789,23 @@ Status InferenceContext::InternalMakeShapeFromTensor(
     DimensionHandle shape_dim = Dim(tensor_shape, 0);
     if (!ValueKnown(shape_dim)) {
       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());
     return ReturnCreatedShape(dims, out);
   }

Fixed code: #include <cstdint>

    // 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);
    }