CVE-2021-29519: Access of Resource Using Incompatible Type ('Type Confusion') (CWE-843)
value: TensorFlow is an end-to-end open source platform for machine learning. The API of `tf.raw_ops.SparseCross` allows combinations which would result in a `CHECK`-failure and denial of service. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/3d782b7d47b1bf2ed32bd4a246d6d6cadc4c903d/tensorflow/core/kernels/sparse_cross_op.cc#L114-L116) is tricked to consider a tensor of type `tstring` which in fact contains integral elements. Fixing the type confusion by preventing mixing `DT_STRING` and `DT_INT64` types solves this issue. 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.
CWE description: The product allocates or initializes a resource such as a pointer, object, or variable using one type, but it later accesses that resource using a type that is incompatible with the original type.
CVSS v3 base score: 5.5
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/sparse_cross_op.cc
Commit date: 2021-04-15 13:08:45 -0700
Commit message: Fix `tf.raw_ops.SparseCross` failing CHECK.

PiperOrigin-RevId: 368701671
Change-Id: Id805729dd9ba0bda36e4bb309408129b55fb649d

Diff: diff --git a/tensorflow/core/kernels/sparse_cross_op.cc b/tensorflow/core/kernels/sparse_cross_op.cc
index 583235b4a309..43b3bedc7450 100644
--- a/tensorflow/core/kernels/sparse_cross_op.cc
+++ b/tensorflow/core/kernels/sparse_cross_op.cc
@@ -25,10 +25,11 @@ limitations under the License.
 #include "tensorflow/core/framework/op_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/fingerprint.h"
 #include "tensorflow/core/platform/strong_hash.h"
 #include "tensorflow/core/util/work_sharder.h"
@@ -458,14 +459,23 @@ int64 CalculateBatchSize(const OpInputList& shapes_list_in,
 
 // Validates input tensors.
 Status ValidateInput(const OpInputList& indices_list_in,
                      const OpInputList& values_list_in,
                      const OpInputList& shapes_list_in,
-                     const OpInputList& dense_list_in) {
+                     const OpInputList& dense_list_in,
+                     const DataType& internal_type) {
   const auto size = indices_list_in.size();
+  // Only perform internal_type check for SparseCrossOp.
+  // Check if the internal_type is not invalid before doing so.
+  bool check_type = internal_type != DT_INVALID;
   // Validates indices_list_in OpInputList.
   for (int i = 0; i < size; i++) {
+    if (check_type && indices_list_in[i].dtype() != DT_INT64) {
+      return errors::InvalidArgument("Input indices should be of type ",
+                                     DT_INT64, " but received ",
+                                     indices_list_in[i].dtype());
+    }
     if (!TensorShapeUtils::IsMatrix(indices_list_in[i].shape())) {
       return errors::InvalidArgument(
           "Input indices should be a matrix but received shape ",
           indices_list_in[i].shape().DebugString(), " at position ", i);
     }
@@ -480,10 +490,18 @@ Status ValidateInput(const OpInputList& indices_list_in,
   if (values_list_in.size() != size) {
     return errors::InvalidArgument("Expected ", size, " input values, got ",
                                    values_list_in.size());
   }
   for (int i = 0; i < size; i++) {
+    // Make sure to avoid the expected type to be string, but input values to be
+    // int64.
+    if (check_type && internal_type == DT_STRING &&
+        values_list_in[i].dtype() == DT_INT64) {
+      return errors::InvalidArgument("Input values should be of internal type ",
+                                     internal_type, " but received ",
+                                     values_list_in[i].dtype());
+    }
     if (!TensorShapeUtils::IsVector(values_list_in[i].shape())) {
       return errors::InvalidArgument(
           "Input values should be a vector but received shape ",
           values_list_in[i].shape().DebugString(), " at position ", i);
     }
@@ -500,10 +518,15 @@ Status ValidateInput(const OpInputList& indices_list_in,
   if (shapes_list_in.size() != size) {
     return errors::InvalidArgument("Expected ", size, " input shapes, got ",
                                    shapes_list_in.size());
   }
   for (int i = 0; i < size; i++) {
+    if (check_type && shapes_list_in[i].dtype() != DT_INT64) {
+      return errors::InvalidArgument("Input shape should be of type ", DT_INT64,
+                                     " but received ",
+                                     shapes_list_in[i].dtype());
+    }
     if (!TensorShapeUtils::IsVector(shapes_list_in[i].shape())) {
       return errors::InvalidArgument(
           "Input shapes should be a vector but received shape ",
           shapes_list_in[i].shape().DebugString(), " at position ", i);
     }
@@ -515,10 +538,18 @@ Status ValidateInput(const OpInputList& indices_list_in,
     }
   }
 
   // Validates dense_list_in OpInputList
   for (int i = 0; i < dense_list_in.size(); ++i) {
+    // Make sure to avoid the expected type to be string, but input values to be
+    // int64.
+    if (check_type && internal_type == DT_STRING &&
+        dense_list_in[i].dtype() == DT_INT64) {
+      return errors::InvalidArgument("Dense inputs should be of internal type ",
+                                     internal_type, " but received ",
+                                     dense_list_in[i].dtype());
+    }
     if (!TensorShapeUtils::IsMatrix(dense_list_in[i].shape())) {
       return errors::InvalidArgument(
           "Dense inputs should be a matrix but received shape ",
           dense_list_in[i].shape().DebugString(), " at position ", i);
     }
@@ -696,10 +727,11 @@ class SparseCrossOp : public OpKernel {
     // Read signed_hash_key_ as int64 since uint64 attributes are not
     // supported by REGISTER_OP.
     int64 signed_hash_key_;
     OP_REQUIRES_OK(context, context->GetAttr("hash_key", &signed_hash_key_));
     hash_key_ = static_cast<uint64>(signed_hash_key_);
+    OP_REQUIRES_OK(context, context->GetAttr("internal_type", &internal_type_));
   }
 
   void Compute(OpKernelContext* context) override {
     OpInputList indices_list_in;
     OP_REQUIRES_OK(context, context->input_list("indices", &indices_list_in));
@@ -709,12 +741,14 @@ class SparseCrossOp : public OpKernel {
     OP_REQUIRES_OK(context, context->input_list("shapes", &shapes_list_in));
     OpInputList dense_list_in;
     OP_REQUIRES_OK(context,
                    context->input_list("dense_inputs", &dense_list_in));
 
-    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
-                                          shapes_list_in, dense_list_in));
+    DataType internal_type = internal_type_;
+    OP_REQUIRES_OK(
+        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
+                               dense_list_in, internal_type));
 
     std::vector<std::unique_ptr<ColumnInterface<InternalType>>> columns =
         GenerateColumnsFromInput<InternalType>(indices_list_in, values_list_in,
                                                shapes_list_in, dense_list_in);
 
@@ -754,10 +788,11 @@ class SparseCrossOp : public OpKernel {
   }
 
  private:
   int64 num_buckets_;
   uint64 hash_key_;
+  DataType internal_type_;
 };
 
 class SparseCrossV2Op : public OpKernel {
  public:
   explicit SparseCrossV2Op(OpKernelConstruction* context) : OpKernel(context) {}
@@ -771,12 +806,15 @@ class SparseCrossV2Op : public OpKernel {
     OP_REQUIRES_OK(context, context->input_list("shapes", &shapes_list_in));
     OpInputList dense_list_in;
     OP_REQUIRES_OK(context,
                    context->input_list("dense_inputs", &dense_list_in));
 
-    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
-                                          shapes_list_in, dense_list_in));
+    // Set internal_type to invalid_type so that the check will be ignored.
+    DataType internal_type = DT_INVALID;
+    OP_REQUIRES_OK(
+        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
+                               dense_list_in, internal_type));
 
     const Tensor* sep_t;
     OP_REQUIRES_OK(context, context->input("sep", &sep_t));
     const tstring separator = sep_t->scalar<tstring>()();
 
@@ -830,12 +868,15 @@ class SparseCrossHashedOp : public OpKernel {
     OP_REQUIRES_OK(context, context->input_list("shapes", &shapes_list_in));
     OpInputList dense_list_in;
     OP_REQUIRES_OK(context,
                    context->input_list("dense_inputs", &dense_list_in));
 
-    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
-                                          shapes_list_in, dense_list_in));
+    // Set internal_type to invalid_type so that the check will be ignored.
+    DataType internal_type = DT_INVALID;
+    OP_REQUIRES_OK(
+        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
+                               dense_list_in, internal_type));
 
     const Tensor* num_buckets_t;
     OP_REQUIRES_OK(context, context->input("num_buckets", &num_buckets_t));
     const int64 num_buckets = num_buckets_t->scalar<int64>()();
 

Vulnerable code:                      const OpInputList& dense_list_in) {
    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
                                          shapes_list_in, dense_list_in));
    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
                                          shapes_list_in, dense_list_in));
    OP_REQUIRES_OK(context, ValidateInput(indices_list_in, values_list_in,
                                          shapes_list_in, dense_list_in));
Fixed code: #include "tensorflow/core/framework/types.pb.h"
                     const OpInputList& dense_list_in,
                     const DataType& internal_type) {
  // Only perform internal_type check for SparseCrossOp.
  // Check if the internal_type is not invalid before doing so.
  bool check_type = internal_type != DT_INVALID;
    if (check_type && indices_list_in[i].dtype() != DT_INT64) {
      return errors::InvalidArgument("Input indices should be of type ",
                                     DT_INT64, " but received ",
                                     indices_list_in[i].dtype());
    }
    // Make sure to avoid the expected type to be string, but input values to be
    // int64.
    if (check_type && internal_type == DT_STRING &&
        values_list_in[i].dtype() == DT_INT64) {
      return errors::InvalidArgument("Input values should be of internal type ",
                                     internal_type, " but received ",
                                     values_list_in[i].dtype());
    }
    if (check_type && shapes_list_in[i].dtype() != DT_INT64) {
      return errors::InvalidArgument("Input shape should be of type ", DT_INT64,
                                     " but received ",
                                     shapes_list_in[i].dtype());
    }
    // Make sure to avoid the expected type to be string, but input values to be
    // int64.
    if (check_type && internal_type == DT_STRING &&
        dense_list_in[i].dtype() == DT_INT64) {
      return errors::InvalidArgument("Dense inputs should be of internal type ",
                                     internal_type, " but received ",
                                     dense_list_in[i].dtype());
    }
    OP_REQUIRES_OK(context, context->GetAttr("internal_type", &internal_type_));
    DataType internal_type = internal_type_;
    OP_REQUIRES_OK(
        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
                               dense_list_in, internal_type));
  DataType internal_type_;
    // Set internal_type to invalid_type so that the check will be ignored.
    DataType internal_type = DT_INVALID;
    OP_REQUIRES_OK(
        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
                               dense_list_in, internal_type));
    // Set internal_type to invalid_type so that the check will be ignored.
    DataType internal_type = DT_INVALID;
    OP_REQUIRES_OK(
        context, ValidateInput(indices_list_in, values_list_in, shapes_list_in,
                               dense_list_in, internal_type));