CVE-2022-23564: Reachable Assertion (CWE-617)
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.
CWE description: 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.
CVSS v3 base score: 6.5
Published: 2022-02-04T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/framework/BUILD, tensorflow/core/framework/resource_handle.cc, tensorflow/core/framework/resource_handle.h, tensorflow/core/framework/tensor.cc
Commit date: 2021-11-06 10:27:15 -0700
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

Diff: diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 838f3f52a4bfb..3dd93fd902866 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -732,11 +732,13 @@ cc_library(
         ":type_index",
         ":types_proto_cc",
         "//tensorflow/core/lib/core:errors",
         "//tensorflow/core/lib/strings:strcat",
         "//tensorflow/core/platform:casts",
+        "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:intrusive_ptr",
+        "//tensorflow/core/platform:macros",
         "//tensorflow/core/platform:statusor",
         "//tensorflow/core/platform:tensor_coding",
         "//tensorflow/core/platform:types",
         "//tensorflow/core/util:managed_stack_trace",
         "@com_google_absl//absl/strings:str_format",
diff --git a/tensorflow/core/framework/resource_handle.cc b/tensorflow/core/framework/resource_handle.cc
index eb6c9706f0460..74b196dd62028 100644
--- a/tensorflow/core/framework/resource_handle.cc
+++ b/tensorflow/core/framework/resource_handle.cc
@@ -15,22 +15,33 @@ limitations under the License.
 
 #include "tensorflow/core/framework/resource_handle.h"
 
 #include "absl/strings/str_format.h"
 #include "tensorflow/core/framework/resource_handle.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/macros.h"
 
 namespace tensorflow {
 
 // Must be declared here for pre-C++17 compatibility.
 /* static */ constexpr const char* ResourceHandle::ANONYMOUS_NAME;
 
 ResourceHandle::ResourceHandle() {}
 
 ResourceHandle::ResourceHandle(const ResourceHandleProto& proto) {
-  FromProto(proto);
+  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);
 }
 
 ResourceHandle::~ResourceHandle() {}
 
 void ResourceHandle::AsProto(ResourceHandleProto* proto) const {
@@ -44,36 +55,40 @@ void ResourceHandle::AsProto(ResourceHandleProto* proto) const {
     dtype_and_shape->set_dtype(dtype_and_shape_pair.dtype);
     dtype_and_shape_pair.shape.AsProto(dtype_and_shape->mutable_shape());
   }
 }
 
-void ResourceHandle::FromProto(const ResourceHandleProto& proto) {
+Status ResourceHandle::FromProto(const ResourceHandleProto& proto) {
   set_device(proto.device());
   set_container(proto.container());
   set_name(proto.name());
   set_hash_code(proto.hash_code());
   set_maybe_type_name(proto.maybe_type_name());
   std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes;
   for (const auto& dtype_and_shape : proto.dtypes_and_shapes()) {
     DataType dtype = dtype_and_shape.dtype();
-    PartialTensorShape shape(dtype_and_shape.shape());
+    PartialTensorShape shape;
+    Status s = PartialTensorShape::BuildPartialTensorShape(
+        dtype_and_shape.shape(), &shape);
+    if (!s.ok()) {
+      return s;
+    }
     dtypes_and_shapes.push_back(DtypeAndPartialTensorShape{dtype, shape});
   }
   dtypes_and_shapes_ = std::move(dtypes_and_shapes);
+  return Status::OK();
 }
 
 string ResourceHandle::SerializeAsString() const {
   ResourceHandleProto proto;
   AsProto(&proto);
   return proto.SerializeAsString();
 }
 
 bool ResourceHandle::ParseFromString(const string& s) {
   ResourceHandleProto proto;
-  const bool status = proto.ParseFromString(s);
-  if (status) FromProto(proto);
-  return status;
+  return proto.ParseFromString(s) && FromProto(proto).ok();
 }
 
 string ResourceHandle::DebugString() const {
   return strings::StrCat("device: ", device(), " container: ", container(),
                          " name: ", name(), " hash_code: ", hash_code(),
@@ -138,11 +153,13 @@ bool DecodeResourceHandleList(std::unique_ptr<port::StringListDecoder> d,
   ResourceHandleProto proto;
   for (int i = 0; i < n; ++i) {
     if (!proto.ParseFromArray(d->Data(sizes[i]), sizes[i])) {
       return false;
     }
-    ps[i].FromProto(proto);
+    if (!ps[i].FromProto(proto).ok()) {
+      return false;
+    }
   }
   return true;
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/framework/resource_handle.h b/tensorflow/core/framework/resource_handle.h
index 0db8e2fcf4d21..46926720a0057 100644
--- a/tensorflow/core/framework/resource_handle.h
+++ b/tensorflow/core/framework/resource_handle.h
@@ -44,10 +44,15 @@ class ResourceHandle {
  public:
   ResourceHandle();
   ResourceHandle(const ResourceHandleProto& proto);
   ~ResourceHandle();
 
+  // 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);
+
   // Unique name for the device containing the resource.
   const std::string& device() const { return device_; }
 
   void set_device(const std::string& device) { device_ = device; }
 
@@ -89,11 +94,11 @@ class ResourceHandle {
     return definition_stack_trace_;
   }
 
   // Conversion to and from ResourceHandleProto
   void AsProto(ResourceHandleProto* proto) const;
-  void FromProto(const ResourceHandleProto& proto);
+  Status FromProto(const ResourceHandleProto& proto);
 
   // Serialization via ResourceHandleProto
   std::string SerializeAsString() const;
   bool ParseFromString(const std::string& s);
 
diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc
index 5e26bf05c03ca..8ae9fd0051652 100644
--- a/tensorflow/core/framework/tensor.cc
+++ b/tensorflow/core/framework/tensor.cc
@@ -535,10 +535,50 @@ TensorBuffer* FromProtoField(Allocator* a, const TensorProto& in, int64_t n) {
   }
 
   return buf;
 }
 
+// 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;
+}
+
 template <>
 TensorBuffer* FromProtoField<Variant>(Allocator* a, const TensorProto& in,
                                       int64_t n) {
   CHECK_GT(n, 0);
   Buffer<Variant>* buf = new Buffer<Variant>(a, n);

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