Commit message: Prevent a null pointer exception in TFLite

PiperOrigin-RevId: 370800206
Change-Id: Idd437ebce4ff224120d8eefc1c14c062173b71d6
CVE: CVE-2021-29590
CWE: CWE-125
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. The implementations of the `Minimum` and `Maximum` TFLite operators can be used to read data outside of bounds of heap allocated objects, if any of the two input tensor arguments are empty. This is because the broadcasting implementation(https://github.com/tensorflow/tensorflow/blob/0d45ea1ca641b21b73bcf9c00e0179cda284e7e7/tensorflow/lite/kernels/internal/reference/maximum_minimum.h#L52-L56) indexes in both tensors with the same index but does not validate that the index is within bounds. 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: maximum_minimum.cc

void TFLiteOperation(TfLiteContext* context, TfLiteNode* node,
                     const OpContext& op_context) {
  reference_ops::MaximumMinimumBroadcastSlow(
      GetTensorShape(op_context.input1),
      GetTensorData<data_type>(op_context.input1),
      GetTensorShape(op_context.input2),
      GetTensorData<data_type>(op_context.input2),
      GetTensorShape(op_context.output),
      GetTensorData<data_type>(op_context.output),
      op_type::template op<data_type>);
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  OpContext op_context(context, node);

  // If inputs have no element, shortcircuit.
  if (NumElements(op_context.input1) == 0 ||
      NumElements(op_context.input2) == 0) {
    return kTfLiteOk;
  }

  switch (op_context.output->type) {
    case kTfLiteFloat32:
      TFLiteOperation<kernel_type, float, OpType>(context, node, op_context);
      break;
    case kTfLiteUInt8:
      TFLiteOperation<kernel_type, uint8_t, OpType>(context, node, op_context);
      break;
    case kTfLiteInt8:
      TFLiteOperation<kernel_type, int8_t, OpType>(context, node, op_context);
      break;
    case kTfLiteInt32:
      TFLiteOperation<kernel_type, int32_t, OpType>(context, node, op_context);
      break;
    case kTfLiteInt64:
      TFLiteOperation<kernel_type, int64_t, OpType>(context, node, op_context);
      break;
    case kTfLiteInt16:
      TFLiteOperation<kernel_type, int16_t, OpType>(context, node, op_context);
      break;
    default:
      context->ReportError(context,
                           "Type %d is currently not supported by Maximum.",
                           op_context.output->type);
      return kTfLiteError;
  }
  return kTfLiteOk;
}

Project: tensorflow
File: maximum_minimum.cc

void TFLiteOperation<maximum_minimum::kGenericOptimized, int8, MaximumOp>(
    TfLiteContext* context, TfLiteNode* node, const OpContext& op_context) {
  tflite::ArithmeticParams op_params;
  const bool need_broadcast = optimized_ops::ProcessBroadcastShapes(
      GetTensorShape(op_context.input1), GetTensorShape(op_context.input2),
      &op_params);
  if (need_broadcast) {
    optimized_ops::BroadcastMaximumDispatch(
        op_params, GetTensorShape(op_context.input1),
        GetTensorData<int8>(op_context.input1),
        GetTensorShape(op_context.input2),
        GetTensorData<int8>(op_context.input2),
        GetTensorShape(op_context.output),
        GetTensorData<int8>(op_context.output), MaximumOp::template op<int8>);
    return;
  }
  reference_ops::MaximumMinimumBroadcastSlow(
      GetTensorShape(op_context.input1), GetTensorData<int8>(op_context.input1),
      GetTensorShape(op_context.input2), GetTensorData<int8>(op_context.input2),
      GetTensorShape(op_context.output), GetTensorData<int8>(op_context.output),
      MaximumOp::template op<int8>);
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MINIMUM_REF() {
  static TfLiteRegistration r = {
      nullptr, nullptr, maximum_minimum::Prepare,
      maximum_minimum::Eval<maximum_minimum::kReference,
                            maximum_minimum::MinimumOp>};
  return &r;
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MINIMUM_GENERIC_OPT() {
  static TfLiteRegistration r = {
      nullptr, nullptr, maximum_minimum::Prepare,
      maximum_minimum::Eval<maximum_minimum::kGenericOptimized,
                            maximum_minimum::MinimumOp>};
  return &r;
}

Project: tensorflow
File: maximum_minimum.cc

  static data_type op(data_type el1, data_type el2) {
    return el1 > el2 ? el1 : el2;
  }

Project: tensorflow
File: maximum_minimum.cc

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  OpContext op_context(context, node);
  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input1->type,
                          op_context.input2->type);
  op_context.output->type = op_context.input1->type;

  bool requires_broadcast =
      !HaveSameShapes(op_context.input1, op_context.input2);

  TfLiteIntArray* output_size = nullptr;
  if (requires_broadcast) {
    TF_LITE_ENSURE_OK(
        context, CalculateShapeForBroadcast(context, op_context.input1,
                                            op_context.input2, &output_size));
  } else {
    output_size = TfLiteIntArrayCopy(op_context.input1->dims);
  }

  return context->ResizeTensor(context, op_context.output, output_size);
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MAXIMUM() {
  return Register_MAXIMUM_GENERIC_OPT();
}

Project: tensorflow
File: maximum_minimum.cc

void TFLiteOperation<maximum_minimum::kGenericOptimized, int8, MinimumOp>(
    TfLiteContext* context, TfLiteNode* node, const OpContext& op_context) {
  tflite::ArithmeticParams op_params;
  const bool need_broadcast = optimized_ops::ProcessBroadcastShapes(
      GetTensorShape(op_context.input1), GetTensorShape(op_context.input2),
      &op_params);
  if (need_broadcast) {
    optimized_ops::BroadcastMinimumDispatch(
        op_params, GetTensorShape(op_context.input1),
        GetTensorData<int8>(op_context.input1),
        GetTensorShape(op_context.input2),
        GetTensorData<int8>(op_context.input2),
        GetTensorShape(op_context.output),
        GetTensorData<int8>(op_context.output), MinimumOp::template op<int8>);
    return;
  }
  reference_ops::MaximumMinimumBroadcastSlow(
      GetTensorShape(op_context.input1), GetTensorData<int8>(op_context.input1),
      GetTensorShape(op_context.input2), GetTensorData<int8>(op_context.input2),
      GetTensorShape(op_context.output), GetTensorData<int8>(op_context.output),
      MinimumOp::template op<int8>);
}

Project: tensorflow
File: maximum_minimum.cc

  static data_type op(data_type el1, data_type el2) {
    return el1 < el2 ? el1 : el2;
  }

Project: tensorflow
File: maximum_minimum.cc

  OpContext(TfLiteContext* context, TfLiteNode* node) {
    input1 = GetInput(context, node, kInputTensor1);
    input2 = GetInput(context, node, kInputTensor2);
    output = GetOutput(context, node, kOutputTensor);
  }

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MAXIMUM_REF() {
  static TfLiteRegistration r = {
      nullptr, nullptr, maximum_minimum::Prepare,
      maximum_minimum::Eval<maximum_minimum::kReference,
                            maximum_minimum::MaximumOp>};
  return &r;
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MAXIMUM_GENERIC_OPT() {
  static TfLiteRegistration r = {
      nullptr, nullptr, maximum_minimum::Prepare,
      maximum_minimum::Eval<maximum_minimum::kGenericOptimized,
                            maximum_minimum::MaximumOp>};
  return &r;
}

Project: tensorflow
File: maximum_minimum.cc

TfLiteRegistration* Register_MINIMUM() {
  return Register_MINIMUM_GENERIC_OPT();
}