CVE-2021-29535: Out-of-bounds Write (CWE-787)
value: TensorFlow is an end-to-end open source platform for machine learning. An attacker can cause a heap buffer overflow in `QuantizedMul` by passing in invalid thresholds for the quantization. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/87cf4d3ea9949051e50ca3f071fc909538a51cd0/tensorflow/core/kernels/quantized_mul_op.cc#L287-L290) assumes that the 4 arguments are always valid scalars and tries to access the numeric value directly. However, if any of these tensors is empty, then `.flat<T>()` is an empty buffer and accessing the element at position 0 results in overflow. 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 writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/quantized_mul_op.cc
Commit date: 2021-04-21 16:20:08 -0700
Commit message: Validate inputs to `QuantizedMul`

PiperOrigin-RevId: 369756982
Change-Id: I00d960cc3b9316fd7a86bd37a44e341c96e17624

Diff: diff --git a/tensorflow/core/kernels/quantized_mul_op.cc b/tensorflow/core/kernels/quantized_mul_op.cc
index fb56f68bf14db..22cff8939449a 100644
--- a/tensorflow/core/kernels/quantized_mul_op.cc
+++ b/tensorflow/core/kernels/quantized_mul_op.cc
@@ -282,14 +282,26 @@ class QuantizedMulOp : public OpKernel {
   explicit QuantizedMulOp(OpKernelConstruction* context) : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     const Tensor& x = context->input(0);
     const Tensor& y = context->input(1);
-    const float min_x = context->input(2).flat<float>()(0);
-    const float max_x = context->input(3).flat<float>()(0);
-    const float min_y = context->input(4).flat<float>()(0);
-    const float max_y = context->input(5).flat<float>()(0);
+    auto& min_x_tensor = context->input(2);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_x_tensor.shape()),
+                errors::InvalidArgument("min_x must be a scalar"));
+    const float min_x = min_x_tensor.flat<float>()(0);
+    auto& max_x_tensor = context->input(3);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_x_tensor.shape()),
+                errors::InvalidArgument("max_x must be a scalar"));
+    const float max_x = max_x_tensor.flat<float>()(0);
+    auto& min_y_tensor = context->input(4);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_y_tensor.shape()),
+                errors::InvalidArgument("min_y must be a scalar"));
+    const float min_y = min_y_tensor.flat<float>()(0);
+    auto& max_y_tensor = context->input(5);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_y_tensor.shape()),
+                errors::InvalidArgument("max_y must be a scalar"));
+    const float max_y = max_y_tensor.flat<float>()(0);
 
     BCast bcast(BCast::FromShape(x.shape()), BCast::FromShape(y.shape()));
     if (!bcast.IsValid()) {
       context->SetStatus(errors::InvalidArgument(
           "Incompatible shapes: ", x.shape().DebugString(), " vs. ",

Vulnerable code:     const float min_x = context->input(2).flat<float>()(0);
    const float max_x = context->input(3).flat<float>()(0);
    const float min_y = context->input(4).flat<float>()(0);
    const float max_y = context->input(5).flat<float>()(0);
Fixed code:     auto& min_x_tensor = context->input(2);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_x_tensor.shape()),
                errors::InvalidArgument("min_x must be a scalar"));
    const float min_x = min_x_tensor.flat<float>()(0);
    auto& max_x_tensor = context->input(3);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_x_tensor.shape()),
                errors::InvalidArgument("max_x must be a scalar"));
    const float max_x = max_x_tensor.flat<float>()(0);
    auto& min_y_tensor = context->input(4);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_y_tensor.shape()),
                errors::InvalidArgument("min_y must be a scalar"));
    const float min_y = min_y_tensor.flat<float>()(0);
    auto& max_y_tensor = context->input(5);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_y_tensor.shape()),
                errors::InvalidArgument("max_y must be a scalar"));
    const float max_y = max_y_tensor.flat<float>()(0);