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
CVE: CVE-2021-41203
CWE: CWE-345
Vulnerability description: 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.

Project: tensorflow
File: saved_tensor_slice_util.h

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

Project: tensorflow
File: saved_tensor_slice_util.h

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

Project: tensorflow
File: tensor_slice_reader.h

bool TensorSliceReader::CopySliceData(const string& name,
                                      const TensorSlice& slice, T* data) const {
  std::vector<std::pair<TensorSlice, string>> details;
  const TensorSliceSet* tss;
  {
    mutex_lock l(mu_);
    tss = FindTensorSlice(name, slice, &details);
    if (!tss && !all_shards_loaded_) {
      VLOG(1) << "Did not find slice in preferred shard, loading all shards."
              << name << ": " << slice.DebugString();
      LoadAllShards();
      tss = FindTensorSlice(name, slice, &details);
    }
    if (!tss) {
      // No such tensor
      return false;
    }
  }
  // We have the data -- copy it over.
  string value;
  for (const auto& x : details) {
    const TensorSlice& slice_s = x.first;
    const string& fname = x.second;
    int idx = gtl::FindWithDefault(fname_to_index_, fname, -1);
    CHECK_GE(idx, 0) << "Failed to find the index for filename " << fname;
    // We read a record in the corresponding sstable
    const string key = EncodeTensorNameSlice(name, slice_s);
    if (!sss_[idx]->Get(key, &value)) {
      VLOG(1) << "Failed to seek to the record for tensor " << name
              << ", slice " << slice_s.DebugString()
              << ": computed key = " << key;
      return false;
    }
    SavedTensorSlices sts;
    if (!ParseProtoUnlimited(&sts, value)) {
      VLOG(1) << "Failed to parse the record for tensor " << name << ", slice "
              << slice_s.DebugString() << ": computed key = " << key;
      return false;
    }
    // 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;
    }
    CopyDataFromTensorSliceToTensorSlice(
        tss->shape(), slice_s, slice,
        checkpoint::TensorProtoData<T>(sts.data().data()), data);
  }
  return true;
}

Project: tensorflow
File: tensor_slice_reader.h

  const Status status() const { return status_; }

Project: tensorflow
File: saved_tensor_slice_util.h

inline protobuf::RepeatedPtrField<string>* MutableTensorProtoData<tstring>(
    TensorProto* t) {
  static_assert(SaveTypeTraits<tstring>::supported,
                "Specified type tstring not supported for Save");
  return t->mutable_string_val();
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline const int32* TensorProtoData<qint32>(const TensorProto& t) {
  static_assert(SaveTypeTraits<qint32>::supported,
                "Specified type qint32 not supported for Restore");
  return reinterpret_cast<const int32*>(t.int_val().data());
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline const string* const* TensorProtoData<tstring>(const TensorProto& t) {
  static_assert(SaveTypeTraits<tstring>::supported,
                "Specified type tstring not supported for Restore");
  return t.string_val().data();
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline int TensorProtoDataSize<qint32>(const TensorProto& t) {
  return t.int_val_size();
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline protobuf::RepeatedField<int32>* MutableTensorProtoData<Eigen::half>(
    TensorProto* t) {
  return t->mutable_half_val();
}

Project: tensorflow
File: tensor_slice_reader.h

  const string& filepattern() const { return filepattern_; }

Project: tensorflow
File: tensor_slice_reader.h

  const std::unordered_map<string, TensorSliceSet*>& Tensors() const {
    return tensors_;
  }

Project: tensorflow
File: tensor_slice_reader.h

  int num_files() const { return sss_.size(); }

Project: tensorflow
File: tensor_slice_reader_test.cc

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());
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline void Fill(const Eigen::half* data, size_t n, TensorProto* t) {
  typename protobuf::RepeatedField<int32>* val = t->mutable_half_val();
  val->Resize(n, 0);
  for (size_t i = 0; i < n; ++i) {
    val->Set(i, Eigen::numext::bit_cast<uint16>(data[i]));
  }
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline void Fill(const qint32* data, size_t n, TensorProto* t) {
  const int32* p = reinterpret_cast<const int32*>(data);
  typename protobuf::RepeatedField<int32> copy(p, p + n);
  t->mutable_int_val()->Swap(&copy);
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline void Fill(const tstring* data, size_t n, TensorProto* t) {
  typename protobuf::RepeatedPtrField<string> copy(data, data + n);
  t->mutable_string_val()->Swap(&copy);
}

Project: tensorflow
File: saved_tensor_slice_util.h

inline const int* TensorProtoData<Eigen::half>(const TensorProto& t) {
  return t.half_val().data();
}