CVE-2021-41206: Improper Validation of Integrity Check Value (CWE-354)
value: TensorFlow is an open source platform for machine learning. In affected versions several TensorFlow operations are missing validation for the shapes of the tensor arguments involved in the call. Depending on the API, this can result in undefined behavior and segfault or `CHECK`-fail related crashes but in some scenarios writes and reads from heap populated arrays are also possible. We have discovered these issues internally via tooling while working on improving/testing GPU op determinism. As such, we don't have reproducers and there will be multiple fixes for these issues. These fixes will be included in TensorFlow 2.7.0. We will also cherrypick these commits on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product does not validate or incorrectly validates the integrity check values or "checksums" of a message. This may prevent it from detecting if the data has been modified or corrupted in transmission.
CVSS v3 base score: 7.8
Published: 2021-11-05T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/linalg/matrix_solve_op.cc, tensorflow/python/kernel_tests/matrix_solve_op_test.py
Commit date: 2021-07-14 20:50:03 -0700
Commit message: Fix crash in MatrixSolve when inputs have different batch dimensions.

Before, the process would crash or certain elements would be silently ignored. Now an InvalidArgument is raised.

PiperOrigin-RevId: 384844020
Change-Id: Iba44417e383bdd0e1abc4012bfca83b2377dd335

Diff: diff --git a/tensorflow/core/kernels/linalg/matrix_solve_op.cc b/tensorflow/core/kernels/linalg/matrix_solve_op.cc
index 70f02bddf9b78..aeb0203b4a337 100644
--- a/tensorflow/core/kernels/linalg/matrix_solve_op.cc
+++ b/tensorflow/core/kernels/linalg/matrix_solve_op.cc
@@ -141,19 +141,26 @@ class MatrixSolveOpGpu : public AsyncOpKernel {
                           "Input and right-hand side must have same rank, got ",
                           ndims, " != ", rhs.dims()),
                       done);
     OP_REQUIRES_ASYNC(
         context, input.dim_size(ndims - 2) == n,
-        errors::InvalidArgument("Input matrices must be squares, got",
+        errors::InvalidArgument("Input matrices must be squares, got ",
                                 input.dim_size(ndims - 2), " != ", n),
         done);
     OP_REQUIRES_ASYNC(context, rhs.dim_size(ndims - 2) == n,
                       errors::InvalidArgument(
                           "Input matrix and right-hand side must have the "
-                          "same number of rows, got",
+                          "same number of rows, got ",
                           n, " != ", rhs.dim_size(ndims - 2)),
                       done);
+    for (int dim = 0; dim < ndims - 2; dim++) {
+      OP_REQUIRES_ASYNC(
+          context, input.dim_size(dim) == rhs.dim_size(dim),
+          errors::InvalidArgument(
+              "All input tensors must have the same outer dimensions."),
+          done);
+    }
 
     // Allocate output.
     Tensor* output;
     OP_REQUIRES_OK_ASYNC(
         context,
diff --git a/tensorflow/python/kernel_tests/matrix_solve_op_test.py b/tensorflow/python/kernel_tests/matrix_solve_op_test.py
index 0d149de2acb5e..1739b2272be81 100644
--- a/tensorflow/python/kernel_tests/matrix_solve_op_test.py
+++ b/tensorflow/python/kernel_tests/matrix_solve_op_test.py
@@ -110,10 +110,16 @@ class MatrixSolveOpTest(test.TestCase):
     matrix = constant_op.constant([[1., 0.], [0., 1.]])
     rhs = constant_op.constant([[1., 0.]])
     with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
       self.evaluate(linalg_ops.matrix_solve(matrix, rhs))
 
+    # The matrix and right-hand side should have the same batch dimensions
+    matrix = np.random.normal(size=(2, 6, 2, 2))
+    rhs = np.random.normal(size=(2, 3, 2, 2))
+    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+      self.evaluate(linalg_ops.matrix_solve(matrix, rhs))
+
   def testNotInvertible(self):
     # The input should be invertible.
     with self.assertRaisesOpError("Input matrix is not invertible."):
       # All rows of the matrix below add to zero
       matrix = constant_op.constant([[1., 0., -1.], [-1., 1., 0.],

Vulnerable code:         errors::InvalidArgument("Input matrices must be squares, got",
                          "same number of rows, got",
Fixed code:         errors::InvalidArgument("Input matrices must be squares, got ",
                          "same number of rows, got ",
    for (int dim = 0; dim < ndims - 2; dim++) {
      OP_REQUIRES_ASYNC(
          context, input.dim_size(dim) == rhs.dim_size(dim),
          errors::InvalidArgument(
              "All input tensors must have the same outer dimensions."),
          done);
    }
    # The matrix and right-hand side should have the same batch dimensions
    matrix = np.random.normal(size=(2, 6, 2, 2))
    rhs = np.random.normal(size=(2, 3, 2, 2))
    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
      self.evaluate(linalg_ops.matrix_solve(matrix, rhs))