Vulnerability: Null pointer dereference in `CompressElement` ### Impact
It is possible to trigger a null pointer dereference in TensorFlow by passing an invalid input to `tf.raw_ops.CompressElement`:

```python
import tensorflow as tf

tf.raw_ops.CompressElement(components=[[]])
```
  
The [implementation](https://github.com/tensorflow/tensorflow/blob/47a06f40411a69c99f381495f490536972152ac0/tensorflow/core/data/compression_utils.cc#L34) was accessing the size of a buffer obtained from the return of a separate function call before validating that said buffer is valid.

### Patches
We have patched the issue in GitHub commit [5dc7f6981fdaf74c8c5be41f393df705841fb7c5](https://github.com/tensorflow/tensorflow/commit/5dc7f6981fdaf74c8c5be41f393df705841fb7c5).

The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for  more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by members of the Aivul Team from Qihoo 360. Concurrently, it was resolved in `master` branch as it was also discovered internally and fixed before the report was handled.
Vulnerability ID: GHSA-c9qf-r67m-p7cg
CWE: CWE-476
Severity score: 7.7
Project: https://github.com/tensorflow/tensorflow
Published: 2021-08-25
File: tensorflow/core/data/compression_utils.cc
Commit message: Fix accessing possible nullptr in tensorflow::data::CompressElement and UncompressElement which are used in tf.data.service.

PiperOrigin-RevId: 373920841
Change-Id: Ia88d78aee09fa19bb53a0f163fd19620d0c68743
Commit author: A. Unique TensorFlower

Vulnerable version (before fix):
 Status CompressElement(const std::vector<Tensor>& element,   int64 total_size = 0;   for (auto& component : element) {     if (DataTypeCanUseMemcpy(component.dtype())) {       // Some datatypes can be memcopied, allowing us to save two copies       // (AsProtoTensorContent and SerializeToArray).       total_size += DMAHelper::buffer(&component)->size();     } else {       non_memcpy_components.emplace_back();       component.AsProtoTensorContent(&non_memcpy_components.back()); Status CompressElement(const std::vector<Tensor>& element,     component.shape().AsProto(metadata->mutable_tensor_shape());     if (DataTypeCanUseMemcpy(component.dtype())) {       const TensorBuffer* buffer = DMAHelper::buffer(&component);       memcpy(position, buffer->data(), buffer->size());       metadata->set_tensor_size_bytes(buffer->size());     } else {       TensorProto& proto = non_memcpy_components[non_memcpy_component_index++];       proto.SerializeToArray(position, proto.ByteSizeLong()); Status UncompressElement(const CompressedElement& compressed,     if (DataTypeCanUseMemcpy(metadata.dtype())) {       out->emplace_back(metadata.dtype(), metadata.tensor_shape());       TensorBuffer* buffer = DMAHelper::buffer(&out->back());       iov[i].iov_base = buffer->data();       iov[i].iov_len = buffer->size();     } else {       // Allocate an empty Tensor. We will fill it out later after       // uncompressing into the tensor_proto_str.

Fixed version (after fix):
 Status CompressElement(const std::vector<Tensor>& element,   int64 total_size = 0;   for (auto& component : element) {     if (DataTypeCanUseMemcpy(component.dtype())) {       const TensorBuffer* buffer = DMAHelper::buffer(&component);       if (buffer) {         total_size += buffer->size();       }     } else {       non_memcpy_components.emplace_back();       component.AsProtoTensorContent(&non_memcpy_components.back()); Status CompressElement(const std::vector<Tensor>& element,     component.shape().AsProto(metadata->mutable_tensor_shape());     if (DataTypeCanUseMemcpy(component.dtype())) {       const TensorBuffer* buffer = DMAHelper::buffer(&component);       if (buffer) {         memcpy(position, buffer->data(), buffer->size());         metadata->set_tensor_size_bytes(buffer->size());       }     } else {       TensorProto& proto = non_memcpy_components[non_memcpy_component_index++];       proto.SerializeToArray(position, proto.ByteSizeLong()); Status UncompressElement(const CompressedElement& compressed,     if (DataTypeCanUseMemcpy(metadata.dtype())) {       out->emplace_back(metadata.dtype(), metadata.tensor_shape());       TensorBuffer* buffer = DMAHelper::buffer(&out->back());       if (buffer) {         iov[i].iov_base = buffer->data();         iov[i].iov_len = buffer->size();       } else {         iov[i].iov_base = nullptr;         iov[i].iov_len = 0;       }     } else {       // Allocate an empty Tensor. We will fill it out later after       // uncompressing into the tensor_proto_str.

Patch:
patch: @@ -29,9 +29,10 @@ Status CompressElement(const std::vector<Tensor>& element,
   int64 total_size = 0;
   for (auto& component : element) {
     if (DataTypeCanUseMemcpy(component.dtype())) {
-      // Some datatypes can be memcopied, allowing us to save two copies
-      // (AsProtoTensorContent and SerializeToArray).
-      total_size += DMAHelper::buffer(&component)->size();
+      const TensorBuffer* buffer = DMAHelper::buffer(&component);
+      if (buffer) {
+        total_size += buffer->size();
+      }
     } else {
       non_memcpy_components.emplace_back();
       component.AsProtoTensorContent(&non_memcpy_components.back());
@@ -53,8 +54,10 @@ Status CompressElement(const std::vector<Tensor>& element,
     component.shape().AsProto(metadata->mutable_tensor_shape());
     if (DataTypeCanUseMemcpy(component.dtype())) {
       const TensorBuffer* buffer = DMAHelper::buffer(&component);
-      memcpy(position, buffer->data(), buffer->size());
-      metadata->set_tensor_size_bytes(buffer->size());
+      if (buffer) {
+        memcpy(position, buffer->data(), buffer->size());
+        metadata->set_tensor_size_bytes(buffer->size());
+      }
     } else {
       TensorProto& proto = non_memcpy_components[non_memcpy_component_index++];
       proto.SerializeToArray(position, proto.ByteSizeLong());
@@ -94,8 +97,13 @@ Status UncompressElement(const CompressedElement& compressed,
     if (DataTypeCanUseMemcpy(metadata.dtype())) {
       out->emplace_back(metadata.dtype(), metadata.tensor_shape());
       TensorBuffer* buffer = DMAHelper::buffer(&out->back());
-      iov[i].iov_base = buffer->data();
-      iov[i].iov_len = buffer->size();
+      if (buffer) {
+        iov[i].iov_base = buffer->data();
+        iov[i].iov_len = buffer->size();
+      } else {
+        iov[i].iov_base = nullptr;
+        iov[i].iov_len = 0;
+      }
     } else {
       // Allocate an empty Tensor. We will fill it out later after
       // uncompressing into the tensor_proto_str.
path: tensorflow/core/data/compression_utils.cc