Commit message: Prevent array write out-of-bounds.

If user passes an invalid axis, then we copy one too many dimensions to the output in the loop below these checks. Even if we didn't do that, there will be further issues with an invalid axis, so we check for that right now.

PiperOrigin-RevId: 371023299
Change-Id: I9eca37ffc2b29e8e48710f500701270ef0790224
CVE: CVE-2021-29603
CWE: CWE-787
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. A specially crafted TFLite model could trigger an OOB write on heap in the TFLite implementation of `ArgMin`/`ArgMax`(https://github.com/tensorflow/tensorflow/blob/102b211d892f3abc14f845a72047809b39cc65ab/tensorflow/lite/kernels/arg_min_max.cc#L52-L59). If `axis_value` is not a value between 0 and `NumDimensions(input)`, then the condition in the `if` is never true, so code writes past the last valid element of `output_dims->data`. 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: arg_min_max.cc

TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input,
                          const TfLiteTensor* axis, TfLiteTensor* output) {
  int axis_value;
  // Retrive all 8 bytes when axis type is kTfLiteInt64 to avoid data loss.
  if (axis->type == kTfLiteInt64) {
    axis_value = static_cast<int>(*GetTensorData<int64_t>(axis));
  } else {
    axis_value = *GetTensorData<int>(axis);
  }
  if (axis_value < 0) {
    axis_value += NumDimensions(input);
  }

  TF_LITE_ENSURE(context, axis_value >= 0);
  TF_LITE_ENSURE(context, axis_value < NumDimensions(input));

  // Copy the input dimensions to output except the axis dimension.
  TfLiteIntArray* output_dims = TfLiteIntArrayCreate(NumDimensions(input) - 1);
  int j = 0;
  for (int i = 0; i < NumDimensions(input); ++i) {
    if (i != axis_value) {
      output_dims->data[j] = SizeOfDimension(input, i);
      ++j;
    }
  }
  return context->ResizeTensor(context, output, output_dims);
}

Project: tensorflow
File: arg_min_max.cc

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node, bool is_arg_max) {
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* axis;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kAxis, &axis));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));
  if (IsDynamicTensor(output)) {
    TF_LITE_ENSURE_STATUS(ResizeOutput(context, input, axis, output));
  }

#define TF_LITE_ARG_MIN_MAX(data_type, axis_type, output_type) \
  optimized_ops::ArgMinMax(                                    \
      GetTensorShape(input), GetTensorData<data_type>(input),  \
      GetTensorData<axis_type>(axis), GetTensorShape(output),  \
      GetTensorData<output_type>(output), is_arg_max)
  if (axis->type == kTfLiteInt32) {
    switch (output->type) {
      case kTfLiteInt32: {
        switch (input->type) {
          case kTfLiteFloat32:
            TF_LITE_ARG_MIN_MAX(float, int32_t, int32_t);
            break;
          case kTfLiteUInt8:
            TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int32_t);
            break;
          case kTfLiteInt8:
            TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int32_t);
            break;
          case kTfLiteInt32:
            TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int32_t);
            break;
          default:
            context->ReportError(context,
                                 "Only float32, uint8, int8 and int32 are "
                                 "supported currently, got %s.",
                                 TfLiteTypeGetName(input->type));
            return kTfLiteError;
        }
      } break;
      case kTfLiteInt64: {
        switch (input->type) {
          case kTfLiteFloat32:
            TF_LITE_ARG_MIN_MAX(float, int32_t, int64_t);
            break;
          case kTfLiteUInt8:
            TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int64_t);
            break;
          case kTfLiteInt8:
            TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int64_t);
            break;
          case kTfLiteInt32:
            TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int64_t);
            break;
          default:
            context->ReportError(context,
                                 "Only float32, uint8, int8 and int32 are "
                                 "supported currently, got %s.",
                                 TfLiteTypeGetName(input->type));
            return kTfLiteError;
        }
      } break;
      default:
        context->ReportError(
            context, "Only int32 and int64 are supported currently, got %s.",
            TfLiteTypeGetName(output->type));
        return kTfLiteError;
    }
  } else {
    switch (output->type) {
      case kTfLiteInt32: {
        switch (input->type) {
          case kTfLiteFloat32:
            TF_LITE_ARG_MIN_MAX(float, int64_t, int32_t);
            break;
          case kTfLiteUInt8:
            TF_LITE_ARG_MIN_MAX(uint8_t, int64_t, int32_t);
            break;
          case kTfLiteInt8:
            TF_LITE_ARG_MIN_MAX(int8_t, int64_t, int32_t);
            break;
          case kTfLiteInt32:
            TF_LITE_ARG_MIN_MAX(int32_t, int64_t, int32_t);
            break;
          default:
            context->ReportError(context,
                                 "Only float32, uint8, int8 and int32 are "
                                 "supported currently, got %s.",
                                 TfLiteTypeGetName(input->type));
            return kTfLiteError;
        }
      } break;
      case kTfLiteInt64: {
        switch (input->type) {
          case kTfLiteFloat32:
            TF_LITE_ARG_MIN_MAX(float, int64_t, int64_t);
            break;
          case kTfLiteUInt8:
            TF_LITE_ARG_MIN_MAX(uint8_t, int64_t, int64_t);
            break;
          case kTfLiteInt8:
            TF_LITE_ARG_MIN_MAX(int8_t, int64_t, int64_t);
            break;
          case kTfLiteInt32:
            TF_LITE_ARG_MIN_MAX(int32_t, int64_t, int64_t);
            break;
          default:
            context->ReportError(context,
                                 "Only float32, uint8, int8 and int32 are "
                                 "supported currently, got %s.",
                                 TfLiteTypeGetName(input->type));
            return kTfLiteError;
        }
      } break;
      default:
        context->ReportError(
            context, "Only int32 and int64 are supported currently, got %s.",
            TfLiteTypeGetName(output->type));
        return kTfLiteError;
    }
  }
#undef TF_LITE_ARG_MIN_MAX

  return kTfLiteOk;
}