CVE-2021-29551: Out-of-bounds Read (CWE-125)
value: TensorFlow is an end-to-end open source platform for machine learning. The implementation of `MatrixTriangularSolve`(https://github.com/tensorflow/tensorflow/blob/8cae746d8449c7dda5298327353d68613f16e798/tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h#L160-L240) fails to terminate kernel execution if one validation condition fails. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 5.5
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C
Files changed: tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h
Commit date: 2021-04-24 16:51:15 -0700
Commit message: Validate (and ensure validation sticks) inputs for `MatrixTriangularSolve`.

PiperOrigin-RevId: 370282444
Change-Id: Iaed61a0b0727cc42c830658b72eb69f785f48dc5

Diff: diff --git a/tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h b/tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h
index 99249f792b6ed..ce5392e62b9fa 100644
--- a/tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h
+++ b/tensorflow/core/kernels/linalg/matrix_triangular_solve_op_impl.h
@@ -160,10 +160,13 @@ class BaseMatrixTriangularSolveOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor& in0 = ctx->input(0);
     const Tensor& in1 = ctx->input(1);
 
     ValidateInputTensors(ctx, in0, in1);
+    if (!ctx->status().ok()) {
+      return;
+    }
 
     MatMulBCast bcast(in0.shape().dim_sizes(), in1.shape().dim_sizes());
     OP_REQUIRES(
         ctx, bcast.IsValid(),
         errors::InvalidArgument(
@@ -228,17 +231,26 @@ class MatrixTriangularSolveOp
   ~MatrixTriangularSolveOp() override {}
 
  private:
   void ValidateInputTensors(OpKernelContext* ctx, const Tensor& in0,
                             const Tensor& in1) override {
+    const auto in0_num_dims = in0.dims();
     OP_REQUIRES(
-        ctx, in0.dims() >= 2,
-        errors::InvalidArgument("In[0] ndims must be >= 2: ", in0.dims()));
+        ctx, in0_num_dims >= 2,
+        errors::InvalidArgument("In[0] ndims must be >= 2: ", in0_num_dims));
 
+    const auto in1_num_dims = in1.dims();
     OP_REQUIRES(
-        ctx, in1.dims() >= 2,
-        errors::InvalidArgument("In[0] ndims must be >= 2: ", in1.dims()));
+        ctx, in1_num_dims >= 2,
+        errors::InvalidArgument("In[1] ndims must be >= 2: ", in1_num_dims));
+
+    const auto in0_last_dim = in0.dim_size(in0_num_dims - 1);
+    const auto in0_prev_dim = in0.dim_size(in0_num_dims - 2);
+    OP_REQUIRES(ctx, in0_last_dim == in0_prev_dim,
+                errors::InvalidArgument(
+                    "In[0] matrices in the last dimensions must be square (",
+                    in0_last_dim, " =/= ", in0_prev_dim, ")"));
   }
 };
 
 #define REGISTER_BATCH_MATRIX_TRIANGULAR_SOLVE_CPU(TYPE)             \
   REGISTER_KERNEL_BUILDER(Name("MatrixTriangularSolve")              \

Vulnerable code:         ctx, in0.dims() >= 2,
        errors::InvalidArgument("In[0] ndims must be >= 2: ", in0.dims()));
        ctx, in1.dims() >= 2,
        errors::InvalidArgument("In[0] ndims must be >= 2: ", in1.dims()));
Fixed code:     if (!ctx->status().ok()) {
      return;
    }
    const auto in0_num_dims = in0.dims();
        ctx, in0_num_dims >= 2,
        errors::InvalidArgument("In[0] ndims must be >= 2: ", in0_num_dims));
    const auto in1_num_dims = in1.dims();
        ctx, in1_num_dims >= 2,
        errors::InvalidArgument("In[1] ndims must be >= 2: ", in1_num_dims));

    const auto in0_last_dim = in0.dim_size(in0_num_dims - 1);
    const auto in0_prev_dim = in0.dim_size(in0_num_dims - 2);
    OP_REQUIRES(ctx, in0_last_dim == in0_prev_dim,
                errors::InvalidArgument(
                    "In[0] matrices in the last dimensions must be square (",
                    in0_last_dim, " =/= ", in0_prev_dim, ")"));