Commit message: Validate inputs to `QuantizedMul`

PiperOrigin-RevId: 369756982
Change-Id: I00d960cc3b9316fd7a86bd37a44e341c96e17624
CVE: CVE-2021-29535
CWE: CWE-787
Vulnerability description: 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.

Project: tensorflow
File: quantized_mul_op.cc

  void Compute(OpKernelContext* context) override {
    const Tensor& x = context->input(0);
    const Tensor& y = context->input(1);
    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. ",
          y.shape().DebugString()));
      return;
    }
    Tensor* z;
    OP_REQUIRES_OK(context, context->allocate_output(
                                0, BCast::ToShape(bcast.output_shape()), &z));

    // 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.
    OP_REQUIRES(context, (max_x > min_x),
                errors::InvalidArgument("max_x must be larger than min_a."));
    OP_REQUIRES(context, (max_y > min_y),
                errors::InvalidArgument("max_x must be larger than min_b."));
    const int32 offset_x = FloatToQuantizedUnclamped<T>(0.0f, min_x, max_x);
    const int32 offset_y = FloatToQuantizedUnclamped<T>(0.0f, min_y, max_y);
    const T* x_data = x.flat<T>().data();
    const T* y_data = y.flat<T>().data();
    Toutput* z_data = z->flat<Toutput>().data();

    const int ndims = bcast.x_reshape().size();
    if (ndims <= 1) {
      if (x.NumElements() == 1) {
        ScalarMultiply<T, Toutput>(context, y_data, offset_y, y.NumElements(),
                                   x_data[0], offset_x, z_data);
      } else if (y.NumElements() == 1) {
        ScalarMultiply<T, Toutput>(context, x_data, offset_x, x.NumElements(),
                                   y_data[0], offset_y, z_data);
      } else {
        VectorMultiply<T, Toutput>(context, x_data, offset_x, y_data, offset_y,
                                   x.NumElements(), z_data);
      }
    } else if (ndims == 2) {
      const T* vector_data;
      int64 vector_num_elements;
      int32 vector_offset;
      const T* tensor_data;
      int64 tensor_num_elements;
      int32 tensor_offset;
      if (x.NumElements() < y.NumElements()) {
        vector_data = x_data;
        vector_num_elements = x.NumElements();
        vector_offset = offset_x;
        tensor_data = y_data;
        tensor_num_elements = y.NumElements();
        tensor_offset = offset_y;
      } else {
        vector_data = y_data;
        vector_num_elements = y.NumElements();
        vector_offset = offset_y;
        tensor_data = x_data;
        tensor_num_elements = x.NumElements();
        tensor_offset = offset_x;
      }
      if (vector_num_elements == 0) {
        context->SetStatus(
            errors::InvalidArgument("vector must have at least 1 element"));
        return;
      }
      VectorTensorMultiply<T, Toutput>(
          vector_data, vector_offset, vector_num_elements, tensor_data,
          tensor_offset, tensor_num_elements, z_data);
    } else {
      LOG(INFO) << "ndims=" << ndims;
      LOG(INFO) << "bcast.x_reshape()="
                << TensorShape(bcast.x_reshape()).DebugString();
      LOG(INFO) << "bcast.y_reshape()="
                << TensorShape(bcast.y_reshape()).DebugString();
      LOG(INFO) << "bcast.x_bcast()="
                << TensorShape(bcast.x_bcast()).DebugString();
      LOG(INFO) << "bcast.y_bcast()="
                << TensorShape(bcast.y_bcast()).DebugString();

      context->SetStatus(errors::Unimplemented(
          "Broadcast between ", context->input(0).shape().DebugString(),
          " and ", context->input(1).shape().DebugString(),
          " is not supported yet."));
      return;
    }

    float min_z_value;
    float max_z_value;
    QuantizationRangeForMultiplication<T, T, Toutput>(
        min_x, max_x, min_y, max_y, &min_z_value, &max_z_value);
    Tensor* z_min = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(1, {}, &z_min));
    z_min->flat<float>()(0) = min_z_value;

    Tensor* z_max = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(2, {}, &z_max));
    z_max->flat<float>()(0) = max_z_value;
  }