Commit message: Adding more validation checks to _ParallelConcatUpdate to avoid NPE.

PiperOrigin-RevId: 402569467
Change-Id: I2db122dab68be2a5e4e8dd3375f5a70c4d2307ec
CVE: CVE-2021-41207
CWE: CWE-369
Vulnerability description: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `ParallelConcat` misses some input validation and can produce a division by 0. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.

Project: tensorflow
File: inplace_ops.cc

Status DoCopy(const CPUDevice& device, const Tensor& x, Tensor* y) {
  CHECK_EQ(x.dtype(), y->dtype());
  switch (x.dtype()) {
#define CASE(type)                                   \
  case DataTypeToEnum<type>::value:                  \
    y->flat<type>().device(device) = x.flat<type>(); \
    break;

    TF_CALL_NUMBER_TYPES(CASE);
    TF_CALL_bool(CASE);
    TF_CALL_tstring(CASE);
#undef CASE
    default:
      return errors::InvalidArgument("Unsupported data type: ",
                                     DataTypeString(x.dtype()));
  }
  return Status::OK();
}

Project: tensorflow
File: inplace_ops.cc

  void Compute(OpKernelContext*) override {}

Project: tensorflow
File: inplace_ops.cc

void DoInplaceStringUpdateOp(const CPUDevice& d, const Tensor& i,
                             const Tensor& v, Tensor* y) {
  auto Ti = i.flat<int32>();
  auto Tv = v.flat_outer_dims<tstring>();
  auto Ty = y->flat_outer_dims<tstring>();
  auto nrows = Ty.dimension(0);
  for (int64_t j = 0; j < Ti.size(); ++j) {
    auto r = (Ti(j) % nrows + nrows) % nrows;  // Guard index range.
    Ty.template chip<0>(r).device(d) = Tv.template chip<0>(j);
  }
}

Project: tensorflow
File: inplace_ops.cc

  void Compute(OpKernelContext* ctx) override {
    Tensor* out = nullptr;
    // We do not know whether the output will be used on GPU. Setting it to be
    // gpu-compatible for now.
    AllocatorAttributes attr;
    attr.set_gpu_compatible(true);
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, shape_, &out, attr));
  }

Project: tensorflow
File: inplace_ops.cc

  void Compute(OpKernelContext* ctx) override {
    auto x = ctx->input(0);
    auto i = ctx->input(1);
    auto v = ctx->input(2);

    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(i.shape()),
                errors::InvalidArgument("i must be a vector. ",
                                        i.shape().DebugString()));
    OP_REQUIRES(ctx, x.dims() == v.dims(),
                errors::InvalidArgument(
                    "x and v shape doesn't match (ranks differ): ",
                    x.shape().DebugString(), " vs. ", v.shape().DebugString()));
    for (int i = 1; i < x.dims(); ++i) {
      OP_REQUIRES(
          ctx, x.dim_size(i) == v.dim_size(i),
          errors::InvalidArgument("x and v shape doesn't match at index ", i,
                                  " : ", x.shape().DebugString(), " vs. ",
                                  v.shape().DebugString()));
    }
    OP_REQUIRES(ctx, i.dim_size(0) == v.dim_size(0),
                errors::InvalidArgument(
                    "i and x shape doesn't match at index 0: ",
                    i.shape().DebugString(), " vs. ", v.shape().DebugString()));

    Tensor y = x;  // This creates an alias intentionally.
    // Skip processing if tensors are empty.
    if (x.NumElements() > 0 && v.NumElements() > 0) {
      OP_REQUIRES_OK(ctx, DoCompute(ctx, i, v, &y));
    }
    ctx->set_output(0, y);
  }

Project: tensorflow

  void Compute(OpKernelContext* ctx) override {
    auto value = ctx->input(0);
    // Value should be at least rank 1. Also the 0th dimension should be
    // at least loc_.
    OP_REQUIRES(ctx, value.dims() >= 1,
                errors::InvalidArgument("value should be at least rank 1."));
    OP_REQUIRES(
        ctx, value.dim_size(0) > loc_,
        errors::InvalidArgument("0th dimension of value = ", value.dim_size(0),
                                " is less than loc_=", loc_));

    auto update = ctx->input(1);

    OP_REQUIRES(
        ctx, value.dims() == update.dims(),
        errors::InvalidArgument("value and update shape doesn't match: ",
                                value.shape().DebugString(), " vs. ",
                                update.shape().DebugString()));
    for (int i = 1; i < value.dims(); ++i) {
      OP_REQUIRES(
          ctx, value.dim_size(i) == update.dim_size(i),
          errors::InvalidArgument("value and update shape doesn't match ",
                                  value.shape().DebugString(), " vs. ",
                                  update.shape().DebugString()));
    }
    OP_REQUIRES(ctx, 1 == update.dim_size(0),
                errors::InvalidArgument("update shape doesn't match: ",
                                        update.shape().DebugString()));

    Tensor output = value;  // This creates an alias intentionally.
    const auto& d = ctx->eigen_device<Device>();
    OP_REQUIRES_OK(
        ctx, ::tensorflow::functor::DoParallelConcat(d, update, loc_, &output));
    ctx->set_output(0, output);
  }

Project: tensorflow
File: inplace_ops.cc

  void Compute(OpKernelContext* ctx) override {
    const Tensor& shape = ctx->input(0);
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsVector(shape.shape()),
        errors::InvalidArgument("shape must be a vector of int32, got shape ",
                                shape.shape().DebugString()));
    auto dims = shape.flat<int32>();
    TensorShape out_shape;
    OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(
                            reinterpret_cast<const int32*>(dims.data()),
                            dims.size(), &out_shape));
    Tensor* out = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out));

    if (init_) {
      functor::SetZeroFunctor<Device, T>()(ctx->eigen_device<Device>(),
                                           out->flat<T>());
    }
  }

Project: tensorflow
File: inplace_ops.cc

  Status DoCompute(OpKernelContext* ctx, const Tensor& x, Tensor* y) override {
    const auto& d = ctx->eigen_device<Device>();
    return ::tensorflow::functor::DoCopy(d, x, y);
  }

Project: tensorflow
File: inplace_ops.cc

  Status DoCompute(OpKernelContext* ctx, const Tensor& i, const Tensor& v,
                   Tensor* y) override {
    const auto& d = ctx->eigen_device<Device>();
    return ::tensorflow::functor::DoInplace(d, op, i, v, y);
  }

Project: tensorflow
File: inplace_ops.cc

  explicit CopyOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {}

Project: tensorflow
File: inplace_ops.cc

  explicit ParallelConcatStart(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &shape_));
  }

Project: tensorflow
File: inplace_ops.cc

  explicit FailureKernel(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx,
                   errors::Internal("Found instance of parallel_stack which "
                                    "could not be properly replaced."));
  }

Project: tensorflow
File: inplace_ops.cc

void DoInplaceOp(const CPUDevice& d, InplaceOpType op, const Tensor& i,
                 const Tensor& v, Tensor* y) {
  auto Ti = i.flat<int32>();
  auto Tv = v.flat_outer_dims<T>();
  auto Ty = y->flat_outer_dims<T>();
  auto nrows = Ty.dimension(0);
  for (int64_t j = 0; j < Ti.size(); ++j) {
    auto r = (Ti(j) % nrows + nrows) % nrows;  // Guard index range.
    switch (op) {
      case I_UPDATE:
        Ty.template chip<0>(r).device(d) = Tv.template chip<0>(j);
        break;
      case I_ADD:
        Ty.template chip<0>(r).device(d) += Tv.template chip<0>(j);
        break;
      case I_SUB:
        Ty.template chip<0>(r).device(d) -= Tv.template chip<0>(j);
        break;
    }
  }
}

Project: tensorflow
File: inplace_ops.cc

  explicit EmptyOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("init", &init_));
  }

Project: tensorflow
File: inplace_ops.cc

Status DoParallelConcat(const CPUDevice& d, const Tensor& value, int32_t loc,
                        Tensor* output) {
  CHECK_EQ(value.dtype(), output->dtype());
  switch (value.dtype()) {
#define CASE(type)                  \
  case DataTypeToEnum<type>::value: \
    return DoParallelConcatUpdate<CPUDevice, type>(d, value, loc, output);
    TF_CALL_POD_TYPES(CASE);
    TF_CALL_tstring(CASE);
    TF_CALL_variant(CASE);
#undef CASE
    default:
      return errors::InvalidArgument("Unsupported data type: ",
                                     DataTypeString(value.dtype()));
  }
}

Project: tensorflow
File: inplace_ops.cc

  explicit CopyOp(OpKernelConstruction* ctx) : CopyOpBase(ctx) {}

Project: tensorflow
File: inplace_ops.cc

  explicit InplaceOp(OpKernelConstruction* ctx) : InplaceOpBase(ctx) {}

Project: tensorflow
File: inplace_ops.cc

Status DoParallelConcatUpdate(const Device& d, const Tensor& value, int32_t loc,
                              Tensor* output) {
  auto Tvalue = value.shaped<T, 2>({1, value.NumElements()});
  auto Toutput = output->flat_outer_dims<T>();
  auto nrows = Toutput.dimension(0);
  auto r = (loc % nrows + nrows) % nrows;  // Guard index range.
  Toutput.template chip<0>(r).device(d) = Tvalue.template chip<0>(0);
  return Status::OK();
}

Project: tensorflow
File: inplace_ops.cc

Status DoInplace(const CPUDevice& device, InplaceOpType op, const Tensor& i,
                 const Tensor& v, Tensor* y) {
  CHECK_EQ(v.dtype(), y->dtype());
  if (op == I_UPDATE) {
    if (v.dtype() == DT_STRING) {
      DoInplaceStringUpdateOp(device, i, v, y);
      return Status::OK();
    } else if (v.dtype() == DT_BOOL) {
      DoInplaceOp<bool>(device, op, i, v, y);
      return Status::OK();
    }
  }
  switch (v.dtype()) {
#define CASE(type)                          \
  case DataTypeToEnum<type>::value:         \
    DoInplaceOp<type>(device, op, i, v, y); \
    break;
    TF_CALL_NUMBER_TYPES(CASE);
#undef CASE
    default:
      return errors::InvalidArgument("Unsupported data type: ",
                                     DataTypeString(v.dtype()));
  }
  return Status::OK();
}

Project: tensorflow
File: inplace_ops.cc

  explicit ParallelConcatUpdate(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("loc", &loc_));
  }

Project: tensorflow
File: inplace_ops.cc

  explicit InplaceOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {}

Project: tensorflow
File: inplace_ops.cc

  void Compute(OpKernelContext* ctx) override {
    auto x = ctx->input(0);
    Tensor* y;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));
    OP_REQUIRES_OK(ctx, DoCompute(ctx, x, y));
  }