CVE-2022-23564 (C++, published 2022-02-04T23:15Z)
Description: value: Tensorflow is an Open Source Machine Learning Framework. When decoding a resource handle tensor from protobuf, a TensorFlow process can encounter cases where a `CHECK` assertion is invalidated based on user controlled arguments. This allows attackers to cause denial of services in TensorFlow processes. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
Severity: MEDIUM (CVSS v3 base score: 6.5)
Weakness: CWE-617 \u2014 Reachable Assertion
The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/framework/BUILD
tensorflow/core/framework/resource_handle.cc
tensorflow/core/framework/resource_handle.h
tensorflow/core/framework/tensor.cc
Commit message: Prevent `CHECK`-fail when decoding resource handles from proto

In certain scenarios, the proto might contain tensors that have too many elements (overflow). This is a `CHECK`-fail in general, but we should prevent this, given how many CVEs caused by that we have received this year (a large fraction of 200).

PiperOrigin-RevId: 408049766
Change-Id: I2ac20b247aa8ed9110846fbdb7a0a9401f2c168c


Vulnerable code:
  FromProto(proto);
void ResourceHandle::FromProto(const ResourceHandleProto& proto) {
    PartialTensorShape shape(dtype_and_shape.shape());
  const bool status = proto.ParseFromString(s);
  if (status) FromProto(proto);
  return status;
    ps[i].FromProto(proto);
  void FromProto(const ResourceHandleProto& proto);

Fixed code:
        "//tensorflow/core/platform:errors",
        "//tensorflow/core/platform:macros",
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/macros.h"
  TF_CHECK_OK(FromProto(proto));
}

Status ResourceHandle::BuildResourceHandle(const ResourceHandleProto& proto,
                                           ResourceHandle* out) {
  if (out == nullptr)
    return errors::Internal(
        "BuildResourceHandle() was called with nullptr for the output");
  return out->FromProto(proto);
Status ResourceHandle::FromProto(const ResourceHandleProto& proto) {
    PartialTensorShape shape;
    Status s = PartialTensorShape::BuildPartialTensorShape(
        dtype_and_shape.shape(), &shape);
    if (!s.ok()) {
      return s;
    }
  return Status::OK();
  return proto.ParseFromString(s) && FromProto(proto).ok();
    if (!ps[i].FromProto(proto).ok()) {
      return false;
    }
  // Use this factory method if the `proto` comes from user controlled input, to
  // prevent a denial of service.
  static Status BuildResourceHandle(const ResourceHandleProto& proto,
                                    ResourceHandle* out);

  Status FromProto(const ResourceHandleProto& proto);
// Separate implementation for `ResourceHandle` to handle the case when the
// proto for the resource is invalid. See `resource_handle.h` constructor and
// static factory builder.
template <>
TensorBuffer* FromProtoField<ResourceHandle>(Allocator* a,
                                             const TensorProto& in, int64_t n) {
  CHECK_GT(n, 0);
  Buffer<ResourceHandle>* buf = new Buffer<ResourceHandle>(a, n);
  ResourceHandle* data = buf->template base<ResourceHandle>();
  if (data == nullptr) {
    buf->Unref();
    return nullptr;
  }
  const int64_t in_n = ProtoHelper<ResourceHandle>::NumElements(in);
  if (in_n <= 0) {
    std::fill_n(data, n, ResourceHandle());
  } else {
    // If tensor shape says we have n < in_n elements in the output tensor
    // then make sure to only decode the first n out of the in_n elements in the
    // in tensors. In all other cases, we decode all in_n elements of in and set
    // the remaining elements up to n to be the default ResourceHandle() value.
    const int64_t real_n = n < in_n ? n : in_n;
    for (int64_t i = 0; i < real_n; ++i) {
      Status s = ResourceHandle::BuildResourceHandle(in.resource_handle_val(i),
                                                     &data[i]);
      if (!s.ok()) {
        LOG(ERROR) << "Could not decode resource handle from proto \""
                   << in.resource_handle_val(i).ShortDebugString()
                   << "\", returned status: " << s.ToString();
        buf->Unref();
        return nullptr;
      }
    }
    for (int64_t i = in_n; i < n; ++i) {
      data[i] = ResourceHandle();
    }
  }
  return buf;
}