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/framework/tensor_slice.cc
tensorflow/core/framework/tensor_slice.h
tensorflow/core/framework/tensor_slice_test.cc
tensorflow/core/util/tensor_slice_reader.cc
tensorflow/core/util/tensor_slice_reader_test.cc
Commit message: Add BuildTensorSlice for building from unvalidated TensorSliceProtos.

This avoids several sources of crashes and undefined behavior when loading
invalid checkpoints.

PiperOrigin-RevId: 392785704
Change-Id: Icd9713c768b882f3b58b427eddac376060696833


Vulnerable code:
      TensorSlice ss_slice(tsp);

Fixed code:

#include <limits>

Status TensorSlice::BuildTensorSlice(const TensorSliceProto& proto,
                                     TensorSlice* output) {
  output->Clear();
  output->starts_.reserve(proto.extent_size());
  output->lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    int64_t l = GetExtentLength(e);
    if (e.start() != 0 || l != kFullExtent) {
      if (e.start() < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and positive length but got start = ",
            e.start(), ", length = ", l, ": extent = ", e.ShortDebugString());
      }
      // Calculating the extent end must not cause signed integer overflow.
      if (static_cast<uint64_t>(e.start()) + static_cast<uint64_t>(e.length()) >
          std::numeric_limits<int64_t>::max()) {
        return errors::InvalidArgument(
            "Extent end exceeds the maximum possible size: extent = ",
            e.ShortDebugString());
      }
    }
    output->starts_.push_back(e.start());
    output->lengths_.push_back(l);
  }

  return Status::OK();
}

  // This factory methods should be used instead of the constructor that takes a
  // `TensorSliceProto` if calling code cannot validate that the sizes specify a
  // valid `TensorSlice`.
  static Status BuildTensorSlice(const TensorSliceProto& proto,
                                 TensorSlice* output);

#include <limits>

// Testing `BuildTensorSlice` with valid and invalid input protos.
TEST(TensorSliceTest, BuildTensorSlice) {
  TensorSliceProto proto;
  TensorSlice({{0, -1}, {0, 10}, {14, 1}}).AsProto(&proto);
  TensorSlice s;

  // Successful building.
  {
    TF_ASSERT_OK(TensorSlice::BuildTensorSlice(proto, &s));
    EXPECT_EQ("-:0,10:14,1", s.DebugString());
  }

  // Failed building due to negative extent start.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(0)->set_start(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to negative extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to missing extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->clear_length();
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to extent end overflowing.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(
        std::numeric_limits<int64_t>::max());
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }
}

      TensorSlice ss_slice;
      status_ = TensorSlice::BuildTensorSlice(tsp, &ss_slice);
      if (!status_.ok()) return;
TEST(TensorSliceReaderTest, InvalidTensorSlice) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "invalid_slice_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_CHECK_OK(writer.Add("test", TensorShape({4, 5}),
                         TensorSlice::ParseOrDie("0,2:-"), data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.mutable_slice(0)->mutable_extent(0)->set_length(-10);
      }
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  // The negative exent length should cause loading to fail.
  EXPECT_FALSE(reader.status().ok());
}