Commit message: Prevent memory overflow in ParseAttrValue from nested tensors.

PiperOrigin-RevId: 370108442
Change-Id: I84d64a5e8895a6aeffbf4749841b4c54d51b5889
CVE: CVE-2021-29615
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. The implementation of `ParseAttrValue`(https://github.com/tensorflow/tensorflow/blob/c22d88d6ff33031aa113e48aa3fc9aa74ed79595/tensorflow/core/framework/attr_value_util.cc#L397-L453) can be tricked into stack overflow due to recursion by giving in a specially crafted input. 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

bool AreAttrValuesEqual(const AttrValue& a, const AttrValue& b,
                        TensorProtosEquality tensor_equality) {
  if (a.type() != b.type()) {
    return false;
  } else if (a.type() != DT_INVALID && b.type() != DT_INVALID) {
    return a.type() == b.type();
  }

  if (a.has_tensor() != b.has_tensor()) {
    return false;
  } else if (a.has_tensor() && b.has_tensor()) {
    return tensor_equality(a.tensor(), b.tensor());
  }

  // `func` field contains a nested AttrValue. Compare such AttrValues
  // recursively.
  if (a.has_func() != b.has_func()) {
    return false;
  } else if (a.has_func() && b.has_func()) {
    const NameAttrList& af = a.func();
    const NameAttrList& bf = b.func();
    if (af.name() != bf.name()) return false;
    std::unordered_map<string, AttrValue> am(af.attr().begin(),
                                             af.attr().end());
    for (const auto& bm_pair : bf.attr()) {
      const auto& iter = am.find(bm_pair.first);
      if (iter == am.end()) return false;
      if (!AreAttrValuesEqual(iter->second, bm_pair.second, tensor_equality))
        return false;
      am.erase(iter);
    }
    if (!am.empty()) return false;
    return true;
  }

  // All other fields in AttrValue have deterministic representations.
  // It is safe to compare their serialized strings.
  return AreSerializedProtosEqual(a, b);
}

Project: tensorflow

bool FastAreAttrValuesEqual(const AttrValue& a, const AttrValue& b) {
  return AreAttrValuesEqual(a, b, FastAreTensorProtosEqual);
}

Project: tensorflow

int64 TensorByteSize(const TensorProto& t) {
  // num_elements returns -1 if shape is not fully defined.
  int64 num_elems = TensorShape(t.tensor_shape()).num_elements();
  return num_elems < 0 ? -1 : num_elems * DataTypeSize(t.dtype());
}

Project: tensorflow

bool AreTensorProtosEqual(const TensorProto& lhs, const TensorProto& rhs) {
  Tensor lhs_t(lhs.dtype());
  bool success = lhs_t.FromProto(lhs);
  DCHECK(success);

  Tensor rhs_t(rhs.dtype());
  success = rhs_t.FromProto(rhs);
  DCHECK(success);

  TensorProto lhs_tp;
  lhs_t.AsProtoTensorContent(&lhs_tp);

  TensorProto rhs_tp;
  rhs_t.AsProtoTensorContent(&rhs_tp);

  return AreSerializedProtosEqual(lhs_tp, rhs_tp);
}

Project: tensorflow

bool AreAttrValuesEqual(const AttrValue& a, const AttrValue& b) {
  return AreAttrValuesEqual(a, b, AreTensorProtosEqual);
}

Project: tensorflow

bool FastAreTensorProtosEqual(const TensorProto& lhs, const TensorProto& rhs) {
  // A small TensorProto can expand into a giant Tensor.  So we avoid
  // conversion to an actual Tensor if we can quickly rule out equality
  // by comparing the Tensor size since different sized Tensors are definitely
  // different.
  const int64 lhs_tensor_bytes = TensorByteSize(lhs);
  const int64 rhs_tensor_bytes = TensorByteSize(rhs);
  if (lhs_tensor_bytes != rhs_tensor_bytes) {
    return false;
  }

  // If the tensor is very large, we'll only compare the proto representation
  // (even though this may miss some equivalent tensors whose actual tensor
  // values are the same but which are described by different TensorProtos).
  if (lhs_tensor_bytes > kMaxAttrValueTensorByteSize) {
    return AreSerializedProtosEqual(lhs, rhs);
  }

  // If the TensorProto representation expands into a much bigger Tensor,
  // we have a fast-path that first compares the protos.
  const int64 lhs_proto_bytes = lhs.ByteSizeLong();
  const bool large_expansion =
      (lhs_proto_bytes < 512 && lhs_tensor_bytes > 4096);
  if (large_expansion && AreSerializedProtosEqual(lhs, rhs)) {
    return true;
  }

  // Fall back to the general code in AreTensorProtosEqual.
  return AreTensorProtosEqual(lhs, rhs);
}