CVE-2020-15191: Unchecked Return Value (CWE-252)
value: In Tensorflow before versions 2.2.1 and 2.3.1, if a user passes an invalid argument to `dlpack.to_dlpack` the expected validations will cause variables to bind to `nullptr` while setting a `status` variable to the error condition. However, this `status` argument is not properly checked. Hence, code following these methods will bind references to null pointers. This is undefined behavior and reported as an error if compiling with `-fsanitize=null`. The issue is patched in commit 22e07fb204386768e5bcbea563641ea11f96ceb8 and is released in TensorFlow versions 2.2.1, or 2.3.1.
CWE description: The product does not check the return value from a method or function, which can prevent it from detecting unexpected states and conditions.
CVSS v3 base score: 5.3
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/c/eager/dlpack.cc, tensorflow/python/dlpack/BUILD, tensorflow/python/dlpack/dlpack_test.py, tensorflow/python/tfe_wrapper.cc
Commit date: 2020-09-18 19:19:43 -0700
Commit message: Fix multiple vulnerabilities in `tf.experimental.dlpack.to_dlpack`.

We have a use after free caused by memory coruption, a segmentation fault caused by memory corruption, several memory leaks and an undefined behavior when taking the reference of a nullptr.

PiperOrigin-RevId: 332568894
Change-Id: Ife0fc05e103b35325094ae5d822ee5fdea764572

Diff: diff --git a/tensorflow/c/eager/dlpack.cc b/tensorflow/c/eager/dlpack.cc
index 30d2009dc6a..df8e9ace997 100644
--- a/tensorflow/c/eager/dlpack.cc
+++ b/tensorflow/c/eager/dlpack.cc
@@ -247,25 +247,40 @@ void TFE_CallDLManagedTensorDeleter(void* dlm_ptr) {
     dlMTensor->deleter(dlMTensor);
   }
 }
 
 void* TFE_HandleToDLPack(TFE_TensorHandle* h, TF_Status* status) {
+  auto tf_dlm_context = GetDlContext(h, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  auto* tf_dlm_data = TFE_TensorHandleDevicePointer(h, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
   const Tensor* tensor = GetTensorFromHandle(h, status);
   TF_DataType data_type = static_cast<TF_DataType>(tensor->dtype());
-  TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
 
+  auto tf_dlm_type = GetDlDataType(data_type, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
   auto* tf_dlm_tensor_ctx = new TfDlManagedTensorCtx(tensor_ref);
   tf_dlm_tensor_ctx->reference = tensor_ref;
 
   DLManagedTensor* dlm_tensor = &tf_dlm_tensor_ctx->tensor;
   dlm_tensor->manager_ctx = tf_dlm_tensor_ctx;
   dlm_tensor->deleter = &DLManagedTensorDeleter;
-  dlm_tensor->dl_tensor.ctx = GetDlContext(h, status);
+  dlm_tensor->dl_tensor.ctx = tf_dlm_context;
   int ndim = tensor->dims();
   dlm_tensor->dl_tensor.ndim = ndim;
-  dlm_tensor->dl_tensor.data = TFE_TensorHandleDevicePointer(h, status);
-  dlm_tensor->dl_tensor.dtype = GetDlDataType(data_type, status);
+  dlm_tensor->dl_tensor.data = tf_dlm_data;
+  dlm_tensor->dl_tensor.dtype = tf_dlm_type;
 
   std::vector<int64_t>* shape_arr = &tf_dlm_tensor_ctx->shape;
   std::vector<int64_t>* stride_arr = &tf_dlm_tensor_ctx->strides;
   shape_arr->resize(ndim);
   stride_arr->resize(ndim, 1);
@@ -274,17 +289,18 @@ void* TFE_HandleToDLPack(TFE_TensorHandle* h, TF_Status* status) {
   }
   for (int i = ndim - 2; i >= 0; --i) {
     (*stride_arr)[i] = (*shape_arr)[i + 1] * (*stride_arr)[i + 1];
   }
 
-  dlm_tensor->dl_tensor.shape = &(*shape_arr)[0];
+  dlm_tensor->dl_tensor.shape = shape_arr->data();
   // There are two ways to represent compact row-major data
   // 1) nullptr indicates tensor is compact and row-majored.
   // 2) fill in the strides array as the real case for compact row-major data.
   // Here we choose option 2, since some frameworks didn't handle the strides
   // argument properly.
-  dlm_tensor->dl_tensor.strides = &(*stride_arr)[0];
+  dlm_tensor->dl_tensor.strides = stride_arr->data();
+
   dlm_tensor->dl_tensor.byte_offset =
       0;  // TF doesn't handle the strides and byte_offsets here
   return static_cast<void*>(dlm_tensor);
 }
 
diff --git a/tensorflow/python/dlpack/BUILD b/tensorflow/python/dlpack/BUILD
index 31ae7f4b435..7d865029bfb 100644
--- a/tensorflow/python/dlpack/BUILD
+++ b/tensorflow/python/dlpack/BUILD
@@ -17,11 +17,10 @@ py_library(
 
 cuda_py_test(
     name = "dlpack_test",
     srcs = ["dlpack_test.py"],
     srcs_version = "PY2AND3",
-    tags = ["noasan"],  # TODO(b/159774807)
     deps = [
         ":dlpack",
         "//tensorflow/python/eager:test",
         "@absl_py//absl/testing:absltest",
         "@absl_py//absl/testing:parameterized",
diff --git a/tensorflow/python/dlpack/dlpack_test.py b/tensorflow/python/dlpack/dlpack_test.py
index e32fbfc86ed..31cc6adbff4 100644
--- a/tensorflow/python/dlpack/dlpack_test.py
+++ b/tensorflow/python/dlpack/dlpack_test.py
@@ -18,13 +18,15 @@ from __future__ import division
 from __future__ import print_function
 
 from absl.testing import parameterized
 import numpy as np
 
+
 from tensorflow.python.dlpack import dlpack
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.platform import test
 from tensorflow.python.ops import array_ops
 
 int_dtypes = [
@@ -103,9 +105,15 @@ class DLPackTest(parameterized.TestCase, test.TestCase):
     self.assertRaisesRegex(Exception, ".* is not supported by dlpack",
                            UnsupportedQint16)
     self.assertRaisesRegex(Exception, ".* is not supported by dlpack",
                            UnsupportedComplex64)
 
+  def testMustPassTensorArgumentToDLPack(self):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "The argument to `to_dlpack` must be a TF tensor, not Python object"):
+      dlpack.to_dlpack([1])
+
 
 if __name__ == "__main__":
   ops.enable_eager_execution()
   test.main()
diff --git a/tensorflow/python/tfe_wrapper.cc b/tensorflow/python/tfe_wrapper.cc
index e76d9e872e4..87064f1f857 100644
--- a/tensorflow/python/tfe_wrapper.cc
+++ b/tensorflow/python/tfe_wrapper.cc
@@ -1356,13 +1356,20 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
         py::return_value_policy::reference);
 
   // DLPack functions
   m.def("TFE_ToDlpackCapsule", [](py::handle& o) {
     PyObject* eager_tensor_pyobject_ptr = o.ptr();
-    TFE_TensorHandle* thandle = EagerTensor_Handle(eager_tensor_pyobject_ptr);
     tensorflow::Safe_TF_StatusPtr status =
         tensorflow::make_safe(TF_NewStatus());
+
+    if (!EagerTensor_CheckExact(eager_tensor_pyobject_ptr)) {
+      status->status = tensorflow::errors::InvalidArgument(
+          "The argument to `to_dlpack` must be a TF tensor, not Python object");
+      tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
+    }
+
+    TFE_TensorHandle* thandle = EagerTensor_Handle(eager_tensor_pyobject_ptr);
     void* dlm_ptr = tensorflow::TFE_HandleToDLPack(thandle, status.get());
     tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
 
     py::capsule capsule(
         dlm_ptr, tensorflow::kDlTensorCapsuleName, [](PyObject* capsule) {

Vulnerable code:   TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
  dlm_tensor->dl_tensor.ctx = GetDlContext(h, status);
  dlm_tensor->dl_tensor.data = TFE_TensorHandleDevicePointer(h, status);
  dlm_tensor->dl_tensor.dtype = GetDlDataType(data_type, status);
  dlm_tensor->dl_tensor.shape = &(*shape_arr)[0];
  dlm_tensor->dl_tensor.strides = &(*stride_arr)[0];
    tags = ["noasan"],  # TODO(b/159774807)
    TFE_TensorHandle* thandle = EagerTensor_Handle(eager_tensor_pyobject_ptr);
Fixed code:   auto tf_dlm_context = GetDlContext(h, status);
  if (!status->status.ok()) {
    return nullptr;
  }

  auto* tf_dlm_data = TFE_TensorHandleDevicePointer(h, status);
  if (!status->status.ok()) {
    return nullptr;
  }

  auto tf_dlm_type = GetDlDataType(data_type, status);
  if (!status->status.ok()) {
    return nullptr;
  }

  TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
  dlm_tensor->dl_tensor.ctx = tf_dlm_context;
  dlm_tensor->dl_tensor.data = tf_dlm_data;
  dlm_tensor->dl_tensor.dtype = tf_dlm_type;
  dlm_tensor->dl_tensor.shape = shape_arr->data();
  dlm_tensor->dl_tensor.strides = stride_arr->data();


from tensorflow.python.framework import errors
  def testMustPassTensorArgumentToDLPack(self):
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "The argument to `to_dlpack` must be a TF tensor, not Python object"):
      dlpack.to_dlpack([1])


    if (!EagerTensor_CheckExact(eager_tensor_pyobject_ptr)) {
      status->status = tensorflow::errors::InvalidArgument(
          "The argument to `to_dlpack` must be a TF tensor, not Python object");
      tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
    }

    TFE_TensorHandle* thandle = EagerTensor_Handle(eager_tensor_pyobject_ptr);