CVE-2022-35973: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. If `QuantizedMatMul` is given nonscalar input for: `min_a`, `max_a`, `min_b`, or `max_b` It gives a segfault that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit aca766ac7693bf29ed0df55ad6bfcc78f35e7f48. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
CVSS v3 base score: 7.5
Published: 2022-09-16T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/quantized_matmul_op.cc, tensorflow/core/kernels/quantized_matmul_op_test.cc
Commit date: 2022-08-04 11:43:20 -0700
Commit message: Fix tf.raw_ops. QuantizedMatMul vulnerability from non scalar min/max a/b arguments.

PiperOrigin-RevId: 465364257

Diff: diff --git a/tensorflow/core/kernels/quantized_matmul_op.cc b/tensorflow/core/kernels/quantized_matmul_op.cc
index 9d3b5279e4bb8..ae65dc3b5e38c 100644
--- a/tensorflow/core/kernels/quantized_matmul_op.cc
+++ b/tensorflow/core/kernels/quantized_matmul_op.cc
@@ -18,15 +18,18 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #define GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK
 #include "public/gemmlowp.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/meta_support.h"
 #include "tensorflow/core/kernels/quantization_utils.h"
 #include "tensorflow/core/kernels/reference_gemm.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
 
 // We have to break this out as a separate function because there are multiple
 // combinations of transpose attributes we need to support, and they have to be
@@ -73,13 +76,25 @@ class QuantizedMatMulOp : public OpKernel {
   }
 
   void Compute(OpKernelContext* context) override {
     const Tensor& a = context->input(0);
     const Tensor& b = context->input(1);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(context->input(2).shape()),
+                errors::InvalidArgument("min_a must be a scalar, but got shape",
+                                        context->input(2).shape()));
     const float min_a = context->input(2).flat<float>()(0);
+    OP_REQUIRES(context, context->input(3).NumElements() == 1,
+                errors::InvalidArgument("max_a must be a scalar, but got shape",
+                                        context->input(3).shape()));
     const float max_a = context->input(3).flat<float>()(0);
+    OP_REQUIRES(context, context->input(4).NumElements() == 1,
+                errors::InvalidArgument("min_b must be a scalar, but got shape",
+                                        context->input(4).shape()));
     const float min_b = context->input(4).flat<float>()(0);
+    OP_REQUIRES(context, context->input(5).NumElements() == 1,
+                errors::InvalidArgument("max_b must be a scalar, but got shape",
+                                        context->input(5).shape()));
     const float max_b = context->input(5).flat<float>()(0);
 
     // Make sure that we have valid quantization ranges for the input buffers.
     // If the difference between the min and max is negative or zero, it makes
     // it hard to do meaningful intermediate operations on the values.
diff --git a/tensorflow/core/kernels/quantized_matmul_op_test.cc b/tensorflow/core/kernels/quantized_matmul_op_test.cc
index c9f05dbc10bb8..f562a2ebcb744 100644
--- a/tensorflow/core/kernels/quantized_matmul_op_test.cc
+++ b/tensorflow/core/kernels/quantized_matmul_op_test.cc
@@ -60,14 +60,14 @@ TEST_F(QuantizedMatMulTest, Small_NoParams) {
   // |  7 |  8 |  9 | 10 |
   // | 11 | 12 | 13 | 14 |
   // | 15 | 16 | 17 | 18 |
   AddInputFromArray<quint8>(TensorShape({3, 4}),
                             {7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18});
-  AddInputFromArray<float>(TensorShape({1}), {0});
-  AddInputFromArray<float>(TensorShape({1}), {255.0f});
-  AddInputFromArray<float>(TensorShape({1}), {0});
-  AddInputFromArray<float>(TensorShape({1}), {255.0f});
+  AddInputFromArray<float>(TensorShape({}), {0});
+  AddInputFromArray<float>(TensorShape({}), {255.0f});
+  AddInputFromArray<float>(TensorShape({}), {0});
+  AddInputFromArray<float>(TensorShape({}), {255.0f});
 
   TF_ASSERT_OK(RunOpKernel());
   // Here are the results we expect, from hand calculations:
   // (1 * 7) + (2 * 11) + (3 * 15) = 74
   // (1 * 8) + (2 * 12) + (3 * 16) = 80
@@ -116,14 +116,14 @@ TEST_F(QuantizedMatMulTest, VerySmall_WithParams) {
   AddInputFromArray<quint8>(TensorShape({a_rows, a_cols}), {11});
 
   // The B matrix is:
   // |   1 |
   AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
-  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
-  AddInputFromArray<float>(TensorShape({1}), {243.0f});
-  AddInputFromArray<float>(TensorShape({1}), {1.0f});
-  AddInputFromArray<float>(TensorShape({1}), {256.0f});
+  AddInputFromArray<float>(TensorShape({}), {-12.0f});
+  AddInputFromArray<float>(TensorShape({}), {243.0f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
   TF_ASSERT_OK(RunOpKernel());
   // We're requesting C = A.transposed() * B,
   // so we expect to get these results:
   // 1*-1 = -1
   // | -1 |
@@ -160,16 +160,54 @@ TEST_F(QuantizedMatMulTest, VerySmall_BadRange) {
   AddInputFromArray<quint8>(TensorShape({a_rows, a_cols}), {11});
 
   // The B matrix is:
   // |   1 |
   AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
-  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
-  AddInputFromArray<float>(TensorShape({1}), {243.0f});
+  AddInputFromArray<float>(TensorShape({}), {-12.0f});
+  AddInputFromArray<float>(TensorShape({}), {243.0f});
   // Here we set the range so that the min and max are equal, so we expect to
   // see an error when we run.
-  AddInputFromArray<float>(TensorShape({1}), {1.0f});
-  AddInputFromArray<float>(TensorShape({1}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
+  EXPECT_EQ(::tensorflow::error::INVALID_ARGUMENT, RunOpKernel().code());
+}
+
+// This test multiplies two 1x1 8bit matrices, but sets invalid quantized min
+// and max values, so we expect to get an error
+TEST_F(QuantizedMatMulTest, VerySmall_BadMinMax) {
+  // These parameters reflect a typical production usage of eight-bit matmuls
+  // in an Inception-style network.
+  const bool transpose_a = true;
+  const int a_rows = 1;
+  const int a_cols = 1;
+  const int b_rows = 1;
+  const int b_cols = 1;
+  const bool transpose_b = false;
+  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op", "QuantizedMatMul")
+                   .Input(FakeInput(DT_QUINT8))
+                   .Input(FakeInput(DT_QUINT8))
+                   .Input(FakeInput(DT_FLOAT))
+                   .Input(FakeInput(DT_FLOAT))
+                   .Input(FakeInput(DT_FLOAT))
+                   .Input(FakeInput(DT_FLOAT))
+                   .Attr("Toutput", DataTypeToEnum<qint32>::v())
+                   .Attr("transpose_a", transpose_a)
+                   .Attr("transpose_b", transpose_b)
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  // The A matrix is:
+  // |  -1 |
+  AddInputFromArray<quint8>(TensorShape({a_rows, a_cols}), {11});
+  // The B matrix is:
+  // |   1 |
+  AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
+  // Here we set the error of a non scalar min_a value, so we expect to see an
+  // error when we run.
+  AddInputFromArray<float>(TensorShape({1}), {2});
+  AddInputFromArray<float>(TensorShape({}), {243.0f});
+  AddInputFromArray<float>(TensorShape({}), {1.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
   EXPECT_EQ(::tensorflow::error::INVALID_ARGUMENT, RunOpKernel().code());
 }
 
 // This test multiplies a couple of small 8-bit matrices, and compares the
 // results with hand-calculated expectations. It uses shifts and offsets to
@@ -231,14 +269,14 @@ TEST_F(QuantizedMatMulTest, Small_WithParams) {
                                                                2,
                                                                5,
                                                                3,
                                                                6,
                                                            });
-  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
-  AddInputFromArray<float>(TensorShape({1}), {243.0f});
-  AddInputFromArray<float>(TensorShape({1}), {0});
-  AddInputFromArray<float>(TensorShape({1}), {255.0f});
+  AddInputFromArray<float>(TensorShape({}), {-12.0f});
+  AddInputFromArray<float>(TensorShape({}), {243.0f});
+  AddInputFromArray<float>(TensorShape({}), {0});
+  AddInputFromArray<float>(TensorShape({}), {255.0f});
   TF_ASSERT_OK(RunOpKernel());
   // We're requesting C = A.transposed() * B,
   // so we expect to get these results:
   // 1*-1 + 2*-5 + 3*-9 = -38
   // 4*-1 + 5*-5 + 6*-9 = -83
@@ -324,14 +362,14 @@ TEST_F(QuantizedMatMulTest, Medium_WithParams) {
        0.012448,  0.0709718,  -0.392394,  -0.434215});
   Tensor b_quantized = FloatTensorToQuantized<quint8>(b_float, b_min, b_max);
 
   AddInputFromArray<quint8>(a_quantized.shape(), a_quantized.flat<quint8>());
   AddInputFromArray<quint8>(b_quantized.shape(), b_quantized.flat<quint8>());
-  AddInputFromArray<float>(TensorShape({1}), {a_min});
-  AddInputFromArray<float>(TensorShape({1}), {a_max});
-  AddInputFromArray<float>(TensorShape({1}), {b_min});
-  AddInputFromArray<float>(TensorShape({1}), {b_max});
+  AddInputFromArray<float>(TensorShape({}), {a_min});
+  AddInputFromArray<float>(TensorShape({}), {a_max});
+  AddInputFromArray<float>(TensorShape({}), {b_min});
+  AddInputFromArray<float>(TensorShape({}), {b_max});
   TF_ASSERT_OK(RunOpKernel());
 
   Tensor expected_float(DT_FLOAT, {a_cols, b_cols});
   test::FillValues<float>(
       &expected_float,

Vulnerable code:   AddInputFromArray<float>(TensorShape({1}), {0});
  AddInputFromArray<float>(TensorShape({1}), {255.0f});
  AddInputFromArray<float>(TensorShape({1}), {0});
  AddInputFromArray<float>(TensorShape({1}), {255.0f});
  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
  AddInputFromArray<float>(TensorShape({1}), {243.0f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
  AddInputFromArray<float>(TensorShape({1}), {243.0f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {-12.0f});
  AddInputFromArray<float>(TensorShape({1}), {243.0f});
  AddInputFromArray<float>(TensorShape({1}), {0});
  AddInputFromArray<float>(TensorShape({1}), {255.0f});
  AddInputFromArray<float>(TensorShape({1}), {a_min});
  AddInputFromArray<float>(TensorShape({1}), {a_max});
  AddInputFromArray<float>(TensorShape({1}), {b_min});
  AddInputFromArray<float>(TensorShape({1}), {b_max});
Fixed code: #include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/errors.h"
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(context->input(2).shape()),
                errors::InvalidArgument("min_a must be a scalar, but got shape",
                                        context->input(2).shape()));
    OP_REQUIRES(context, context->input(3).NumElements() == 1,
                errors::InvalidArgument("max_a must be a scalar, but got shape",
                                        context->input(3).shape()));
    OP_REQUIRES(context, context->input(4).NumElements() == 1,
                errors::InvalidArgument("min_b must be a scalar, but got shape",
                                        context->input(4).shape()));
    OP_REQUIRES(context, context->input(5).NumElements() == 1,
                errors::InvalidArgument("max_b must be a scalar, but got shape",
                                        context->input(5).shape()));
  AddInputFromArray<float>(TensorShape({}), {0});
  AddInputFromArray<float>(TensorShape({}), {255.0f});
  AddInputFromArray<float>(TensorShape({}), {0});
  AddInputFromArray<float>(TensorShape({}), {255.0f});
  AddInputFromArray<float>(TensorShape({}), {-12.0f});
  AddInputFromArray<float>(TensorShape({}), {243.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-12.0f});
  AddInputFromArray<float>(TensorShape({}), {243.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  EXPECT_EQ(::tensorflow::error::INVALID_ARGUMENT, RunOpKernel().code());
}

// This test multiplies two 1x1 8bit matrices, but sets invalid quantized min
// and max values, so we expect to get an error
TEST_F(QuantizedMatMulTest, VerySmall_BadMinMax) {
  // These parameters reflect a typical production usage of eight-bit matmuls
  // in an Inception-style network.
  const bool transpose_a = true;
  const int a_rows = 1;
  const int a_cols = 1;
  const int b_rows = 1;
  const int b_cols = 1;
  const bool transpose_b = false;
  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op", "QuantizedMatMul")
                   .Input(FakeInput(DT_QUINT8))
                   .Input(FakeInput(DT_QUINT8))
                   .Input(FakeInput(DT_FLOAT))
                   .Input(FakeInput(DT_FLOAT))
                   .Input(FakeInput(DT_FLOAT))
                   .Input(FakeInput(DT_FLOAT))
                   .Attr("Toutput", DataTypeToEnum<qint32>::v())
                   .Attr("transpose_a", transpose_a)
                   .Attr("transpose_b", transpose_b)
                   .Finalize(node_def()));
  TF_ASSERT_OK(InitOp());
  // The A matrix is:
  // |  -1 |
  AddInputFromArray<quint8>(TensorShape({a_rows, a_cols}), {11});
  // The B matrix is:
  // |   1 |
  AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
  // Here we set the error of a non scalar min_a value, so we expect to see an
  // error when we run.
  AddInputFromArray<float>(TensorShape({1}), {2});
  AddInputFromArray<float>(TensorShape({}), {243.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-12.0f});
  AddInputFromArray<float>(TensorShape({}), {243.0f});
  AddInputFromArray<float>(TensorShape({}), {0});
  AddInputFromArray<float>(TensorShape({}), {255.0f});
  AddInputFromArray<float>(TensorShape({}), {a_min});
  AddInputFromArray<float>(TensorShape({}), {a_max});
  AddInputFromArray<float>(TensorShape({}), {b_min});
  AddInputFromArray<float>(TensorShape({}), {b_max});