Commit message: Fix `tf.io.decode_raw` bugs and update documentation.

Fixes cases where specifying `fixed_length` resulted in data loss and even segfault and corruption of the Python interpreter. The fix is subtle but needed due to pointer arithmetic rules.

Makes sure that `fixed_length` does not change the output when present but not needed.

Eliminates needless copy and cast in the main codepath.

PiperOrigin-RevId: 371322725
Change-Id: I514ef67a2961c86422f69d05122d31615e87896c
CVE: CVE-2021-29614
CWE: CWE-787
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. The implementation of `tf.io.decode_raw` produces incorrect results and crashes the Python interpreter when combining `fixed_length` and wider datatypes. The implementation of the padded version(https://github.com/tensorflow/tensorflow/blob/1d8903e5b167ed0432077a3db6e462daf781d1fe/tensorflow/core/kernels/decode_padded_raw_op.cc) is buggy due to a confusion about pointer arithmetic rules. First, the code computes(https://github.com/tensorflow/tensorflow/blob/1d8903e5b167ed0432077a3db6e462daf781d1fe/tensorflow/core/kernels/decode_padded_raw_op.cc#L61) the width of each output element by dividing the `fixed_length` value to the size of the type argument. The `fixed_length` argument is also used to determine the size needed for the output tensor(https://github.com/tensorflow/tensorflow/blob/1d8903e5b167ed0432077a3db6e462daf781d1fe/tensorflow/core/kernels/decode_padded_raw_op.cc#L63-L79). This is followed by reencoding code(https://github.com/tensorflow/tensorflow/blob/1d8903e5b167ed0432077a3db6e462daf781d1fe/tensorflow/core/kernels/decode_padded_raw_op.cc#L85-L94). The erroneous code is the last line above: it is moving the `out_data` pointer by `fixed_length * sizeof(T)` bytes whereas it only copied at most `fixed_length` bytes from the input. This results in parts of the input not being decoded into the output. Furthermore, because the pointer advance is far wider than desired, this quickly leads to writing to outside the bounds of the backing data. This OOB write leads to interpreter crash in the reproducer mentioned here, but more severe attacks can be mounted too, given that this gadget allows writing to periodically placed locations in memory. 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: decode_padded_raw_op.cc

  explicit DecodePaddedRawOp(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("out_type", &out_type_));

    const bool host_is_little_endian = port::kLittleEndian;
    bool data_is_little_endian;
    OP_REQUIRES_OK(context,
                   context->GetAttr("little_endian", &data_is_little_endian));
    convert_data_endianness_ = host_is_little_endian != data_is_little_endian;
  }

Project: tensorflow
File: decode_padded_raw_op.cc

  void Compute(OpKernelContext* context) override {
    const auto& input = context->input(0);
    auto flat_in = input.flat<tstring>();

    int fixed_length;
    const auto& length_input = context->input(1);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(length_input.shape()),
                errors::InvalidArgument("k must be scalar, got shape ",
                                        length_input.shape().DebugString()));
    fixed_length = length_input.scalar<int32>()();

    OP_REQUIRES(
        context, fixed_length % sizeof(T) == 0,
        errors::InvalidArgument(
            "fixed_length (", fixed_length,
            ") must be a multiple of the size of out_type (", sizeof(T), ")"));

    OP_REQUIRES(context, fixed_length > 0,
                errors::InvalidArgument("fixed_length (", fixed_length,
                                        ") must be greater than zero."));

    int width = fixed_length / sizeof(T);

    TensorShape out_shape = input.shape();
    out_shape.AddDim(width);
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(
        context, context->allocate_output("output", out_shape, &output_tensor));

    if (flat_in.size() == 0) {  // Empty input
      return;
    }

    auto out = output_tensor->flat_inner_dims<T>();
    T* out_data = out.data();

    // Forcibly clear memory - we're going to copy variable length strings in,
    // and need to ensure that if we don't write to byte N when we copy, that
    // we're not getting random data.
    memset(out_data, 0, fixed_length * flat_in.size());

    // If the data is already in the host's byte order, or if the width of the
    // output type is a single byte (meaning the ordering doesn't matter), we
    // can copy the memory directly.
    if (!convert_data_endianness_ || sizeof(T) == 1) {
      for (int64 i = 0; i < flat_in.size(); ++i) {
        const auto to_copy =
            std::min(flat_in(i).size(), static_cast<size_t>(fixed_length));
        memcpy(out_data, flat_in(i).data(), to_copy);
        // Note: increase out_data by width since it's already of type T* so
        // each shift amount is implicitly multiplied by sizeof(T) according to
        // pointer arithmetic rules.
        out_data += width;
      }
    } else {
      // Otherwise, the data is not in the host's byte order, and rather than a
      // direct copy, we need to reverse the byte ordering of each element.
      for (int64 i = 0; i < flat_in.size(); ++i) {
        const char* in_data_bytes =
            reinterpret_cast<const char*>(flat_in(i).data());
        char* out_data_bytes = reinterpret_cast<char*>(out_data);
        const char* p_in = in_data_bytes;
        char* p_out = out_data_bytes;
        for (; p_in < in_data_bytes + fixed_length;
             p_in += sizeof(T), p_out += sizeof(T)) {
          std::reverse_copy(p_in, p_in + sizeof(T), p_out);
        }
        // Note: increase out_data by width since it's already of type T* so
        // each shift amount is implicitly multiplied by sizeof(T) according to
        // pointer arithmetic rules.
        out_data += width;
      }
    }
  }