CVE-2022-29210: Out-of-bounds Write (CWE-787)
value: TensorFlow is an open source platform for machine learning. In version 2.8.0, the `TensorKey` hash function used total estimated `AllocatedBytes()`, which (a) is an estimate per tensor, and (b) is a very poor hash function for constants (e.g. `int32_t`). It also tried to access individual tensor bytes through `tensor.data()` of size `AllocatedBytes()`. This led to ASAN failures because the `AllocatedBytes()` is an estimate of total bytes allocated by a tensor, including any pointed-to constructs (e.g. strings), and does not refer to contiguous bytes in the `.data()` buffer. The discoverers could not use this byte vector anyway because types such as `tstring` include pointers, whereas they needed to hash the string values themselves. This issue is patched in Tensorflow versions 2.9.0 and 2.8.1.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 5.5
Published: 2022-05-21T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/core/framework/tensor_key.h, tensorflow/python/kernel_tests/data_structures/BUILD
Commit date: 2022-05-03 12:51:58 -0700
Commit message: Fix TensorKey hash function.

The original hash function only used total estimated `AllocatedBytes()`, which (a) is an estimate per tensor, and (b) is a very poor hash function for constants (e.g. `int32_t`).
It also tried to access individual tensor bytes through `tensor.data()` of size `AllocatedBytes()`.  This led to ASAN failures because the `AllocatedBytes()` is an estimate of total bytes allocated by a tensor, including any pointed-to constructs (e.g. strings), and does not refer to contiguous bytes in the `.data()` buffer.  We couldn't use this byte vector anyways, since types like `tstring` include pointers, whereas we need to hash the string values themselves.

Modified the hash function to more closely mirror the `==` operator.  This correctly handles `tstring` and any numeric types that do have contiguous storage.  Other types are currently left as unimplemented.

PiperOrigin-RevId: 446265413

Diff: diff --git a/tensorflow/core/framework/tensor_key.h b/tensorflow/core/framework/tensor_key.h
index 243205dc50f83..3bde6fce6214c 100644
--- a/tensorflow/core/framework/tensor_key.h
+++ b/tensorflow/core/framework/tensor_key.h
@@ -14,10 +14,11 @@ limitations under the License.
 ==============================================================================*/
 #ifndef TENSORFLOW_CORE_FRAMEWORK_TENSOR_KEY_H_
 #define TENSORFLOW_CORE_FRAMEWORK_TENSOR_KEY_H_
 
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
 
 namespace tensorflow {
 
 class TensorKey : public Tensor {
  public:
@@ -30,20 +31,22 @@ class TensorKey : public Tensor {
     if (t1.dtype() != t2.dtype() || t1.shape() != t2.shape()) {
       return false;
     }
     if (DataTypeCanUseMemcpy(t1.dtype())) {
       return t1.tensor_data() == t2.tensor_data();
-    }
-    if (t1.dtype() == DT_STRING) {
+    } else if (t1.dtype() == DT_STRING) {
       const auto s1 = t1.unaligned_flat<tstring>();
       const auto s2 = t2.unaligned_flat<tstring>();
       for (int64_t i = 0, n = t1.NumElements(); i < n; ++i) {
         if (TF_PREDICT_FALSE(s1(i) != s2(i))) {
           return false;
         }
       }
       return true;
+    } else {
+      DCHECK(false) << "Unimplemented dtype " << DataTypeString(t1.dtype())
+                    << std::endl;
     }
     return false;
   }
 
   friend bool operator!=(const TensorKey& t1, const TensorKey& t2) {
@@ -51,18 +54,23 @@ class TensorKey : public Tensor {
   }
 
   // Needed for absl hash function.
   template <typename H>
   friend H AbslHashValue(H h, const TensorKey& k) {
-    const uint8* d = static_cast<uint8*>(k.data());
-    size_t s = k.AllocatedBytes();
-    std::vector<uint8> vec;
-    vec.reserve(s);
-    for (int i = 0; i < s; i++) {
-      vec.push_back(d[i]);
+    if (DataTypeCanUseMemcpy(k.dtype())) {
+      return H::combine(std::move(h), k.tensor_data());
+    } else if (k.dtype() == DT_STRING) {
+      const auto strs = k.unaligned_flat<tstring>();
+      for (int64_t i = 0, n = k.NumElements(); i < n; ++i) {
+        h = H::combine(std::move(h), strs(i));
+      }
+      return h;
+    } else {
+      DCHECK(false) << "Unimplemented dtype " << DataTypeString(k.dtype())
+                    << std::endl;
     }
-    return H::combine(std::move(h), s);
+    return h;
   }
 };
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/python/kernel_tests/data_structures/BUILD b/tensorflow/python/kernel_tests/data_structures/BUILD
index 79c46dd4dc18f..4f9def3b6ab00 100644
--- a/tensorflow/python/kernel_tests/data_structures/BUILD
+++ b/tensorflow/python/kernel_tests/data_structures/BUILD
@@ -163,12 +163,10 @@ tf_py_test(
     size = "small",
     srcs = ["map_ops_test.py"],
     grpc_enabled = True,
     tags = [
         "no_windows",  # TODO(b/192259628)
-        "noasan",  # TODO(b/164696004)
-        "notsan",  # TODO(b/164696004)
     ],
     deps = [
         "//tensorflow/python:array_ops",
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:framework_for_generated_wrappers",

Vulnerable code:     }
    if (t1.dtype() == DT_STRING) {
    const uint8* d = static_cast<uint8*>(k.data());
    size_t s = k.AllocatedBytes();
    std::vector<uint8> vec;
    vec.reserve(s);
    for (int i = 0; i < s; i++) {
      vec.push_back(d[i]);
    return H::combine(std::move(h), s);
        "noasan",  # TODO(b/164696004)
        "notsan",  # TODO(b/164696004)
Fixed code: #include "tensorflow/core/framework/types.h"
    } else if (t1.dtype() == DT_STRING) {
    } else {
      DCHECK(false) << "Unimplemented dtype " << DataTypeString(t1.dtype())
                    << std::endl;
    if (DataTypeCanUseMemcpy(k.dtype())) {
      return H::combine(std::move(h), k.tensor_data());
    } else if (k.dtype() == DT_STRING) {
      const auto strs = k.unaligned_flat<tstring>();
      for (int64_t i = 0, n = k.NumElements(); i < n; ++i) {
        h = H::combine(std::move(h), strs(i));
      }
      return h;
    } else {
      DCHECK(false) << "Unimplemented dtype " << DataTypeString(k.dtype())
                    << std::endl;
    return h;