Commit message: A negative size in one of the split sizes allowed the computed size of another
to exceed the total dimension, leading to a segfault and security vulnerability.
Adding a check for negative sizes prevents this.

PiperOrigin-RevId: 401035665
Change-Id: I79bbe329787dac82aa4bf60397a9129b716aedab
CVE: CVE-2021-41222
CWE: CWE-703
Vulnerability description: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `SplitV` can trigger a segfault is an attacker supplies negative arguments. This occurs whenever `size_splits` contains more than one value and at least one value is negative. 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: split_v_op.cc

  explicit SplitVOpCPU(OpKernelConstruction* c) : Base(c) {}

Project: tensorflow
File: split_v_op.cc

  std::tuple<IndexType, IndexType, IndexType> SetDims(
      const TensorShape& input_shape, const int32_t split_dim) const {
    static_assert(std::is_integral<IndexType>::value,
                  "IndexType must be an integer type");
    int32_t prefix_dim_size = 1;
    for (int i = 0; i < split_dim; ++i) {
      prefix_dim_size *= input_shape.dim_size(i);
    }

    // Caller must ensure that dim_size and suffix_dim_size are <
    // std::numeric_limits<IndexType>::max()
    IndexType split_dim_size =
        static_cast<IndexType>(input_shape.dim_size(split_dim));

    IndexType suffix_dim_size = 1;
    for (int i = split_dim + 1; i < input_shape.dims(); ++i) {
      suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
    }
    return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
  }

Project: tensorflow
File: split_v_op.cc

  explicit SplitVOpGPU(OpKernelConstruction* c) : Base(c) {}

Project: tensorflow
File: split_v_op.cc

  void ParallelSplitByInputData(OpKernelContext* context,
                                const InputReshapedType& input_reshaped,
                                const TensorShape& input_shape,
                                const std::vector<Tlen>& split_sizes_vec,
                                const int32_t split_dim) const {
    const T* p_data = input_reshaped.data();
    const uint32 elem_pkg = input_reshaped.dimensions().rank() == 3
                                ? input_reshaped.dimension(2)
                                : 1;
    const uint32 line_elem_num =
        (input_reshaped.dimensions().rank() >= 2 ? input_reshaped.dimension(1)
                                                 : 1) *
        elem_pkg;
    const uint32 line_num = input_reshaped.dimension(0);

    // Prepare the output matrix.
    std::vector<T*> outputs(split_sizes_vec.size());
    for (uint64 i = 0; i < split_sizes_vec.size(); ++i) {
      TensorShape output_shape(input_shape);
      output_shape.set_dim(split_dim, split_sizes_vec[i]);
      Tensor* result = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(i, output_shape, &result));
      outputs[i] = static_cast<T*>(&result->flat<T>()(0));
    }

    auto sub_split_func = [&split_sizes_vec, &p_data, elem_pkg, &outputs,
                           line_elem_num](int32_t start_part,
                                          int32_t end_part) {
      int start = start_part * line_elem_num;
      int end = end_part * line_elem_num;
      uint32 times = 0;
      for (int32_t i = start; i < end;) {
        for (uint32 j = 0; j < split_sizes_vec.size(); ++j) {
          const auto copy_elem_num = split_sizes_vec[j] * elem_pkg;
          std::copy_n(p_data + i, copy_elem_num,
                      &(outputs[j][(start_part + times) * copy_elem_num]));
          i += copy_elem_num;
        }
        ++times;
      }
    };

    uint32 part_size =
        context->device()->tensorflow_cpu_worker_threads()->num_threads;
    Shard(part_size,
          context->device()->tensorflow_cpu_worker_threads()->workers, line_num,
          line_num, sub_split_func);
  }

Project: tensorflow
File: split_v_op.cc

  void ComputeEasyCases(OpKernelContext* context, bool* done,
                        std::vector<Tlen>* split_sizes_vec) {
    const int32_t num_split = context->num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const Tensor& split_tensor = context->input(1);
    const Tensor& split_dim_tensor = context->input(2);

    OP_REQUIRES(context, split_dim_tensor.NumElements() == 1,
                errors::InvalidArgument("split_dim_tensor must have "
                                        "exactly one element."));

    const int32_t split_dim_orig = split_dim_tensor.flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;

    OP_REQUIRES(
        context,
        split_tensor.dims() == 1 && split_tensor.NumElements() == num_split,
        errors::InvalidArgument("size of the split_tensor must be 1-D and have "
                                "the same elements as outputs got ",
                                split_tensor.dims(), " -D and ",
                                split_tensor.NumElements(), " elements"));

    auto split_sizes_d = split_tensor.vec<Tlen>();

    split_sizes_vec->resize(split_sizes_d.size());

    std::copy(split_sizes_d.data(), split_sizes_d.data() + split_sizes_d.size(),
              split_sizes_vec->begin());

    OP_REQUIRES(
        context, num_split > 0,
        errors::InvalidArgument(
            "Number of ways to split should be > 0, but got ", num_split));

    OP_REQUIRES(
        context, 0 <= split_dim && split_dim < input.dims(),
        errors::InvalidArgument("-input rank(-", input.dims(),
                                ") <= split_dim < input rank (", input.dims(),
                                "), but got ", split_dim_orig));

    Tlen input_size_split_dim = input_shape.dim_size(split_dim);

    // Special case 1: num_split == 1. Nothing to do.
    if (num_split == 1) {
      context->set_output(0, context->input(0));
      OP_REQUIRES(
          context, (*split_sizes_vec)[0] == input_size_split_dim,
          errors::InvalidArgument("If there is only one output, it must have "
                                  "the same size as the input. Input size: ",
                                  input_size_split_dim,
                                  " output size: ", (*split_sizes_vec)[0]));
      *done = true;
      return;
    }

    // Determine sizes of output, in case of a -1 input value
    int neg_one_dim = -1;
    Tlen determined_size = 0;
    for (int d = 0; d < split_sizes_vec->size(); ++d) {
      Tlen size = (*split_sizes_vec)[d];

      if (size == -1) {
        OP_REQUIRES(context, neg_one_dim == -1,
                    errors::InvalidArgument("There can only be one -1 in the "
                                            "input."));
        neg_one_dim = d;
      } else {
        determined_size += size;
      }
    }

    OP_REQUIRES(
        context,
        (neg_one_dim == -1 && determined_size == input_size_split_dim) ||
            (neg_one_dim >= 0 && determined_size <= input_size_split_dim),
        errors::InvalidArgument("Determined shape must either match "
                                "input shape along split_dim exactly if "
                                "fully specified, or be less than the size of "
                                "the input along split_dim if not fully "
                                "specified.  Got: ",
                                determined_size));

    if (neg_one_dim >= 0) {
      (*split_sizes_vec)[neg_one_dim] = input_size_split_dim - determined_size;
    }

    for (int i = 0; i < split_sizes_vec->size(); ++i) {
      const Tlen& split_size = (*split_sizes_vec)[i];
      OP_REQUIRES(context, split_size >= Tlen(0),
                  errors::InvalidArgument("Split size at index ", i,
                                          " must be >= 0. Got: ", split_size));
    }

    // Special case 2: split along the 1st dimension. The requirements are that
    // either we are splitting the outer dimension of two or more such that
    // every outer subpart is aligned or that the split sizes mean that they are
    // always aligned. In these cases, we can share the underlying buffer.
    //
    // Apply this optimization conservatively: if input is aligned,
    // the resulting tensors must be aligned. It's conservative
    // because if the immediate consumer of the resulting tensors are
    // not using eigen for computation, its perfectly fine to avoid
    // the copying.
    if (SplitHasAlignedOutputsInFirstDimension(
            input_shape, split_dim, absl::MakeConstSpan(*split_sizes_vec))) {
      Tlen start = 0;
      for (int i = 0; i < num_split; ++i) {
        context->set_output(i,
                            input.Slice(start, start + (*split_sizes_vec)[i]));
        start += (*split_sizes_vec)[i];
      }
      *done = true;
      return;
    }
  }

Project: tensorflow
File: split_v_op.cc

  void operator()(OpKernelContext* context,
                  const InputReshapedType& input_reshaped,
                  const std::vector<int64_t>& split_start_points,
                  const TensorShape& input_shape, int32_t split_dim,
                  Eigen::DenseIndex prefix_dim_size,
                  Eigen::DenseIndex split_dim_size,
                  Eigen::DenseIndex suffix_dim_size,
                  std::vector<Tlen>& split_sizes_vec,
                  const MakeSizesType& make_sizes,
                  const ReshapeResultType& reshape_result) const {
    Eigen::DSizes<Eigen::DenseIndex, NDims> indices;
    for (int i = 0; i < NDims; ++i) {
      indices[i] = 0;
    }
    const auto num_threads =
        context->device()->tensorflow_cpu_worker_threads()->num_threads;
    // TODO(jewillco): Tune heuristic further.
    const auto input_element_count = input_shape.num_elements();
    const int num_split = split_start_points.size();
    const bool use_parallelism_between_outputs =
        (num_split >= kMinimumSplitNum &&
         input_element_count >= std::min(num_threads, num_split) * 4096 &&
         input_element_count < num_split * 180 * 1024);

    auto range_output_func = [&indices, context, &input_shape, split_dim,
                              &split_sizes_vec, &split_start_points,
                              use_parallelism_between_outputs, &input_reshaped,
                              &make_sizes,
                              &reshape_result](int64_t start, int64_t limit) {
      for (int64_t i = start; i < limit; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));

        const auto sizes = make_sizes(split_sizes_vec[i]);

        if (sizes.TotalSize() > 0) {
          auto result_shaped = reshape_result(result, split_sizes_vec[i]);

          auto current_indices = indices;
          current_indices[NDims - 2] = split_start_points[i];
          if (use_parallelism_between_outputs) {
            // Use sequential implementation for single output.
            result_shaped = input_reshaped.slice(current_indices, sizes);
          } else {
            // This implementation may be parallel internally.
            functor::Split<CPUDevice, T, NDims>()(
                context->eigen_device<CPUDevice>(), result_shaped,
                input_reshaped, current_indices, sizes);
          }
        }
      }
    };

    // 1. Parallel performance is not as good as serial when the amount of data
    // is too small (<kMinimumInputSize);
    // 2. There is sufficient data on the 0th dimension to ensure parallelism;
    // 3. This method only supports non-zero split.
    if ((input_element_count >= kMinimumInputSize) &&
        input_reshaped.dimension(0) > kMinimumDim0Size && split_dim) {
      // Each thread processes the same amount of data, and then copies data
      // to all output tensors .
      ParallelSplitByInputData(context, input_reshaped, input_shape,
                               split_sizes_vec, split_dim);
    } else if (use_parallelism_between_outputs) {
      // A thread maps a output tensor, this thread will traverse all the data,
      // and then put specified data to mapped output tensor. Run in parallel,
      // disabling parallelism in functor.
      Shard(num_split,
            context->device()->tensorflow_cpu_worker_threads()->workers,
            num_split, input_element_count / num_split, range_output_func);
    } else {
      // Run sequentially, but allow internal parallelism in functor.
      range_output_func(0, num_split);
    }
  }

Project: tensorflow
File: split_v_op.cc

  explicit SplitVOpBase(OpKernelConstruction* c) : OpKernel(c) {}

Project: tensorflow
File: split_v_op.cc

  void Compute(OpKernelContext* context) override {
    bool done = false;
    std::vector<Tlen> split_sizes_vec;
    Base::ComputeEasyCases(context, &done, &split_sizes_vec);
    if (!context->status().ok() || done) {
      return;
    }
    const int32_t num_split = Base::num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32_t split_dim_orig = context->input(2).flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;
    OP_REQUIRES(
        context,
        FastBoundsCheck(input.NumElements(), std::numeric_limits<int32>::max()),
        errors::InvalidArgument("Split on GPU requires input size "
                                "< max int32"));

    int32_t prefix_dim_size;
    int32_t split_dim_size;
    int32_t suffix_dim_size;
    std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
        Base::template SetDims<int32>(input_shape, split_dim);

    // use the same approach as concat (see documentation there)
    // reshape to 2D

    if (num_split > 16) {
      GpuDeviceArrayOnHost<T*> ptrs(context, num_split);
      OP_REQUIRES_OK(context, ptrs.Init());

      GpuDeviceArrayOnHost<Tlen> offsets(context, num_split + 1);
      OP_REQUIRES_OK(context, offsets.Init());

      Tlen offset = 0;
      int entry = split_sizes_vec[0];
      bool fixed_size =
          std::all_of(split_sizes_vec.begin(), split_sizes_vec.end(),
                      [&entry](int n) { return n == entry; });

      for (int i = 0; i < num_split; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));
        ptrs.Set(i, result->flat<T>().data());
        offsets.Set(i, offset);
        offset += split_sizes_vec[i] * suffix_dim_size;
      }
      offsets.Set(num_split, offset);
      OP_REQUIRES_OK(context, ptrs.Finalize());
      OP_REQUIRES_OK(context, offsets.Finalize());

      if (input.NumElements() > 0) {
        SplitVOpGPULaunch<T, Tlen>().Run(
            context->eigen_device<GPUDevice>(), fixed_size,
            input.flat<T>().data(), prefix_dim_size,
            input.NumElements() / prefix_dim_size, offsets.data(), ptrs.data());
        OP_REQUIRES(
            context, context->op_device_context()->stream()->ok(),
            errors::Internal("Launch of gpu kernel for SplitVOp failed"));
      }
    } else {
      Eigen::DenseIndex prefix_dim_size;
      Eigen::DenseIndex split_dim_size;
      Eigen::DenseIndex suffix_dim_size;

      std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
          Base::template SetDims<Eigen::DenseIndex>(input_shape, split_dim);
      auto input_reshaped = input.shaped<T, 2>(
          {prefix_dim_size, split_dim_size * suffix_dim_size});

      Eigen::DSizes<Eigen::DenseIndex, 2> indices{0, 0};

      for (int i = 0; i < num_split; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));

        Eigen::DSizes<Eigen::DenseIndex, 2> sizes{
            prefix_dim_size, split_sizes_vec[i] * suffix_dim_size};

        if (sizes.TotalSize() > 0) {
          auto result_shaped = result->shaped<T, 2>(
              {prefix_dim_size, split_sizes_vec[i] * suffix_dim_size});

          functor::SplitCustom<GPUDevice, T>()(
              context->eigen_device<GPUDevice>(), result_shaped, input_reshaped,
              indices, sizes);
        }
        indices[1] += split_sizes_vec[i] * suffix_dim_size;
      }
    }
  }

Project: tensorflow
File: split_v_op.cc

  void Compute(OpKernelContext* context) override {
    bool done = false;
    std::vector<Tlen> split_sizes_vec;
    Base::ComputeEasyCases(context, &done, &split_sizes_vec);
    if (!context->status().ok() || done) {
      return;
    }
    const int32_t num_split = Base::num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32_t split_dim_orig = context->input(2).flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;

    // Android also uses int32 indexing, so check here also.
    OP_REQUIRES(
        context,
        FastBoundsCheck(input.NumElements(),
                        std::numeric_limits<Eigen::DenseIndex>::max()),
        errors::InvalidArgument("Split requires input size < ",
                                std::numeric_limits<Eigen::DenseIndex>::max()));

    Eigen::DenseIndex prefix_dim_size;
    Eigen::DenseIndex split_dim_size;
    Eigen::DenseIndex suffix_dim_size;

    std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
        Base::template SetDims<Eigen::DenseIndex>(input_shape, split_dim);
    std::vector<int64_t> split_start_points(num_split);
    for (int i = 0; i < num_split; ++i) {
      if (i == 0) {
        split_start_points[i] = 0;
      } else {
        split_start_points[i] =
            split_start_points[i - 1] + split_sizes_vec[i - 1];
      }
    }

    if (prefix_dim_size == 1) {
      auto input_reshaped =
          input.shaped<T, 2>({split_dim_size, suffix_dim_size});
      auto make_sizes = [&](Eigen::DenseIndex split_size) {
        return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size};
      };
      auto reshape_result = [&](Tensor* result, Tlen split_size) {
        return result->shaped<T, 2>({split_size, suffix_dim_size});
      };
      SplitVOpCPUImpl<T, Tlen, decltype(input_reshaped), 2>{}(
          context, input_reshaped, split_start_points, input_shape, split_dim,
          prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec,
          make_sizes, reshape_result);
    } else {
      auto input_reshaped = input.shaped<T, 3>(
          {prefix_dim_size, split_dim_size, suffix_dim_size});
      auto make_sizes = [&](Eigen::DenseIndex split_size) {
        return Eigen::DSizes<Eigen::DenseIndex, 3>{prefix_dim_size, split_size,
                                                   suffix_dim_size};
      };
      auto reshape_result = [&](Tensor* result, Tlen split_size) {
        return result->shaped<T, 3>(
            {prefix_dim_size, split_size, suffix_dim_size});
      };
      SplitVOpCPUImpl<T, Tlen, decltype(input_reshaped), 3>{}(
          context, input_reshaped, split_start_points, input_shape, split_dim,
          prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec,
          make_sizes, reshape_result);
    }
  }

Project: tensorflow
File: split_v_op.cc

  static bool SplitHasAlignedOutputsInFirstDimension(
      const TensorShape& input_shape, int32_t split_dim,
      absl::Span<const Tlen> split_sizes) {
    if (split_dim != 0) {
      return false;
    }
    Tlen start = 0;
    for (const Tlen split_size : split_sizes) {
      if (!IsDim0SliceAligned<T>(input_shape, start, start + split_size)) {
        return false;
      }
      start += split_size;
    }
    return true;
  }