CVE-2021-41203 (C, published 2021-11-05T21:15Z)
Description: value: TensorFlow is an open source platform for machine learning. In affected versions an attacker can trigger undefined behavior, integer overflows, segfaults and `CHECK`-fail crashes if they can change saved checkpoints from outside of TensorFlow. This is because the checkpoints loading infrastructure is missing validation for invalid file formats. The fixes will be included in TensorFlow 2.7.0. We will also cherrypick these commits on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
Severity: MEDIUM (CVSS v3 base score: 7.8)
Weakness: CWE-190 \u2014 Integer Overflow or Wraparound
The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/util/saved_tensor_slice_util.h
tensorflow/core/util/tensor_slice_reader.h
tensorflow/core/util/tensor_slice_reader_test.cc
Commit message: Avoid buffer overflow when loading tensors with insufficient data from checkpoints.

`CopyDataFromTensorSliceToTensorSlice` does not (and cannot conveniently)
provide any bounds checking on its own, so the size is instead checked prior
to passing unvalidated data to that function.

PiperOrigin-RevId: 392971286
Change-Id: If2073b36d4d5eedd386329f56729395fd7effee1


Vulnerable code:

Fixed code:
template <typename T>
int TensorProtoDataSize(const TensorProto& t);

  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {    \
    return t.FIELD##_val_size();                                  \
  }                                                               \
  template <>                                                     \
  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {      \
    return t.FIELD##_val_size() / 2;                                \
  }                                                                 \
  template <>                                                       \
template <>
inline int TensorProtoDataSize<qint32>(const TensorProto& t) {
  return t.int_val_size();
}

template <>
inline int TensorProtoDataSize<Eigen::half>(const TensorProto& t) {
  return t.half_val_size();
}

template <>
inline int TensorProtoDataSize<tstring>(const TensorProto& t) {
  return t.string_val_size();
}

    // Ensure the TensorSlice contains the expected amount of data.
    TensorShape shp_s;
    Status s = slice_s.SliceTensorShape(tss->shape(), &shp_s);
    if (!s.ok()) {
      VLOG(1) << "Failed to slice tensor " << name << ", slice "
              << slice_s.DebugString() << ": " << s;
      return false;
    }
    if (checkpoint::TensorProtoDataSize<T>(sts.data().data()) !=
        shp_s.num_elements()) {
      VLOG(1) << "Tensor " << name << ", slice " << slice_s.DebugString()
              << " had an unexpected amount of data: expected = "
              << shp_s.num_elements() << ", got = "
              << checkpoint::TensorProtoDataSize<T>(sts.data().data());
      return false;
    }
TEST(TensorSliceReaderTest, MissingTensorData) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "missing_data_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_ASSERT_OK(writer.Add("test", TensorShape({4, 5}),
                          TensorSlice::ParseOrDie("0,2:-"), data));
  TF_ASSERT_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [&](SavedTensorSlices sts) {
    if (sts.has_data()) {
      // Replace the data with only 4 elements.
      Fill(data, 4, sts.mutable_data()->mutable_data());
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  TF_ASSERT_OK(reader.status());

  // The tensor should be present, but loading it should fail due to the missing
  // data.
  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
  std::unique_ptr<Tensor> tensor;
  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
}