Commit message: Prevent division by 0

PiperOrigin-RevId: 370984990
Change-Id: Ib324955bbeb1cbd97c82fd5d61a00a2697c9a2de
CVE: CVE-2021-29597
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. The implementation of the `SpaceToBatchNd` TFLite operator is [vulnerable to a division by zero error](https://github.com/tensorflow/tensorflow/blob/412c7d9bb8f8a762c5b266c9e73bfa165f29aac8/tensorflow/lite/kernels/space_to_batch_nd.cc#L82-L83). An attacker can craft a model such that one dimension of the `block` input is 0. Hence, the corresponding value in `block_shape` is 0. 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: space_to_batch_nd.cc

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

  SpaceToBatchNDContext op_context(context, node);
  TF_LITE_ENSURE(context,
                 NumDimensions(op_context.input) >= kInputMinDimensionNum);
  TF_LITE_ENSURE(context,
                 NumDimensions(op_context.input) <= kInputMaxDimensionNum);
  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type,
                          op_context.output->type);

  if (!IsConstantTensor(op_context.block_shape) ||
      !IsConstantTensor(op_context.paddings)) {
    SetTensorToDynamic(op_context.output);
    return kTfLiteOk;
  }
  return ResizeOutputTensor(context, &op_context);
}

Project: tensorflow
File: space_to_batch_nd.cc

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

Project: tensorflow
File: space_to_batch_nd.cc

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

Project: tensorflow
File: space_to_batch_nd.cc

TfLiteRegistration* Register_SPACE_TO_BATCH_ND() {
  // return Register_SPACE_TO_BATCH_ND_REF();
  return Register_SPACE_TO_BATCH_ND_GENERIC_OPT();
}

Project: tensorflow
File: space_to_batch_nd.cc

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

  // Resize the output tensor if the output tensor is dynamic.
  if (IsDynamicTensor(op_context.output)) {
    TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
  }

#define TF_LITE_SPACE_TO_BATCH_ND(type, scalar, pad_value)             \
  tflite::SpaceToBatchParams op_params;                                \
  op_params.output_offset = pad_value;                                 \
  type::SpaceToBatchND(op_params, GetTensorShape(op_context.input),    \
                       GetTensorData<scalar>(op_context.input),        \
                       GetTensorShape(op_context.block_shape),         \
                       GetTensorData<int32_t>(op_context.block_shape), \
                       GetTensorShape(op_context.paddings),            \
                       GetTensorData<int32_t>(op_context.paddings),    \
                       GetTensorShape(op_context.output),              \
                       GetTensorData<scalar>(op_context.output))
  switch (op_context.input->type) {  // Already know in/out types are same.
    case kTfLiteFloat32:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, float, 0);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, float, 0);
      }
      break;
    case kTfLiteUInt8:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, uint8_t,
                                  op_context.output->params.zero_point);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, uint8_t,
                                  op_context.output->params.zero_point);
      }
      break;
    case kTfLiteInt8:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int8_t,
                                  op_context.output->params.zero_point);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int8_t,
                                  op_context.output->params.zero_point);
      }
      break;
    case kTfLiteInt32:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int32_t, 0);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int32_t, 0);
      }
      break;
    case kTfLiteInt64:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int64_t, 0);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int64_t, 0);
      }
      break;
    default:
      context->ReportError(
          context, "Type %d is currently not supported by SpaceToBatch.",
          op_context.input->type);
      return kTfLiteError;
  }
#undef TF_LITE_SPACE_TO_BATCH_ND
  return kTfLiteOk;
}

Project: tensorflow
File: space_to_batch_nd.cc

  SpaceToBatchNDContext(TfLiteContext* context, TfLiteNode* node) {
    input = GetInput(context, node, 0);
    block_shape = GetInput(context, node, 1);
    paddings = GetInput(context, node, 2);
    output = GetOutput(context, node, 0);
  }

Project: tensorflow
File: space_to_batch_nd.cc

TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
                                SpaceToBatchNDContext* op_context) {
  TfLiteIntArray* input_size = op_context->input->dims;
  const int32* block_shape = GetTensorData<int32>(op_context->block_shape);
  const int32* paddings_data = GetTensorData<int32>(op_context->paddings);

  int spatial_dims_num = input_size->size - 2;
  // Block_shape should be a 1D tensor with dimension [spatial_dims_num].
  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), 1);
  TF_LITE_ENSURE_EQ(context, op_context->block_shape->dims->data[0],
                    spatial_dims_num);
  // Paddings should be a 2D tensor with dimension [spatial_dims_num, 2].
  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->paddings), 2);
  TF_LITE_ENSURE_EQ(context, op_context->paddings->dims->data[0],
                    spatial_dims_num);
  TF_LITE_ENSURE_EQ(context, op_context->paddings->dims->data[1], 2);

  TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size);

  // Ensures the input height and width (with padding) is a multiple of block
  // shape height and width.
  int output_batch_size = input_size->data[0];
  for (int dim = 0; dim < spatial_dims_num; ++dim) {
    int final_dim_size = (input_size->data[dim + 1] + paddings_data[dim * 2] +
                          paddings_data[dim * 2 + 1]);
    TF_LITE_ENSURE(context, block_shape[dim] != 0);
    TF_LITE_ENSURE_EQ(context, final_dim_size % block_shape[dim], 0);
    output_size->data[dim + 1] = final_dim_size / block_shape[dim];
    output_batch_size *= block_shape[dim];
  }

  output_size->data[0] = output_batch_size;
  output_size->data[input_size->size - 1] =
      input_size->data[input_size->size - 1];

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