CVE-2022-35983: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. If `Save` or `SaveSlices` is run over tensors of an unsupported `dtype`, it results in a `CHECK` fail that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 5dd7b86b84a864b834c6fa3d7f9f51c87efa99d4. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
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: 7.5
Published: 2022-09-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/util/tensor_slice_writer.cc, tensorflow/core/util/tensor_slice_writer.h, tensorflow/core/util/tensor_slice_writer_test.cc
Commit date: 2022-07-18 10:56:09 -0700
Commit message: Fix tf.raw_ops.SaveSlices vulnerability with unsupported dtypes.

Check that given dtype is supported and emit a descriptive error if not.

PiperOrigin-RevId: 461660795

Diff: diff --git a/tensorflow/core/util/tensor_slice_writer.cc b/tensorflow/core/util/tensor_slice_writer.cc
index ca907916fd3f1..75197a52c6d55 100644
--- a/tensorflow/core/util/tensor_slice_writer.cc
+++ b/tensorflow/core/util/tensor_slice_writer.cc
@@ -129,10 +129,20 @@ Status TensorSliceWriter::Finish() {
   return s;
 }
 
 /* static */
 size_t TensorSliceWriter::MaxBytesPerElement(DataType dt) {
+  size_t max_bytes_per_element =
+      TensorSliceWriter::MaxBytesPerElementOrZero(dt);
+  if (max_bytes_per_element == 0) {
+    LOG(FATAL) << "MaxBytesPerElement not implemented for dtype: " << dt;
+  }
+  return max_bytes_per_element;
+}
+
+/* static */
+size_t TensorSliceWriter::MaxBytesPerElementOrZero(DataType dt) {
   switch (dt) {
     case DT_FLOAT:
       return 4;
     case DT_DOUBLE:
       return 8;
@@ -168,13 +178,12 @@ size_t TensorSliceWriter::MaxBytesPerElement(DataType dt) {
       return 3;
     case DT_INVALID:
     case DT_STRING:
     case DT_BFLOAT16:
     default:
-      LOG(FATAL) << "MaxBytesPerElement not implemented for dtype: " << dt;
+      return 0;
   }
-  return 0;
 }
 
 template <>
 Status TensorSliceWriter::SaveData(const tstring* data, int64_t num_elements,
                                    SavedSlice* ss) {
diff --git a/tensorflow/core/util/tensor_slice_writer.h b/tensorflow/core/util/tensor_slice_writer.h
index 3436966866aab..34aa7c0b91673 100644
--- a/tensorflow/core/util/tensor_slice_writer.h
+++ b/tensorflow/core/util/tensor_slice_writer.h
@@ -66,10 +66,12 @@ class TensorSliceWriter {
   static Status SaveData(const T* data, int64_t num_elements, SavedSlice* ss);
 
   static size_t MaxBytesPerElement(DataType dt);
 
  private:
+  static size_t MaxBytesPerElementOrZero(DataType dt);
+
   static constexpr size_t kMaxMessageBytes = 1LL << 31;
   // Filling in the TensorProto in a SavedSlice will add the following
   // header bytes, in addition to the data:
   // - 1 byte: TensorProto tag and wire format
   // - <= 5 bytes: TensorProto length
@@ -160,13 +162,19 @@ Status TensorSliceWriter::Add(const string& name, const TensorShape& shape,
 }
 
 template <typename T>
 Status TensorSliceWriter::SaveData(const T* data, int64_t num_elements,
                                    SavedSlice* ss) {
-  size_t size_bound =
-      ss->ByteSize() + kTensorProtoHeaderBytes +
-      (MaxBytesPerElement(DataTypeToEnum<T>::value) * num_elements);
+  size_t max_bytes_per_element =
+      MaxBytesPerElementOrZero(DataTypeToEnum<T>::value);
+  if (max_bytes_per_element == 0) {
+    return errors::InvalidArgument(
+        "Tensor slice serialization not implemented for dtype ",
+        DataTypeToEnum<T>::value);
+  }
+  size_t size_bound = ss->ByteSize() + kTensorProtoHeaderBytes +
+                      (max_bytes_per_element * num_elements);
   if (size_bound > kMaxMessageBytes) {
     return errors::InvalidArgument(
         "Tensor slice is too large to serialize (conservative estimate: ",
         size_bound, " bytes)");
   }
diff --git a/tensorflow/core/util/tensor_slice_writer_test.cc b/tensorflow/core/util/tensor_slice_writer_test.cc
index c9740730a66a2..5120be5a8b5e3 100644
--- a/tensorflow/core/util/tensor_slice_writer_test.cc
+++ b/tensorflow/core/util/tensor_slice_writer_test.cc
@@ -13,21 +13,23 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/util/tensor_slice_writer.h"
 
+#include <algorithm>
 #include <array>
+#include <memory>
+#include <vector>
 
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/core/stringpiece.h"
-#include "tensorflow/core/lib/io/path.h"
-#include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/util/saved_tensor_slice_util.h"
 #include "tensorflow/core/util/tensor_slice_reader.h"
 
 namespace tensorflow {
@@ -360,8 +362,19 @@ TEST(TensorSliceWriteTest, SizeErrors) {
     EXPECT_TRUE(absl::StrContains(s.error_message(),
                                   "Tensor slice is too large to serialize"));
   }
 }
 
+TEST(TensorSliceWriterTest, InvalidInput) {
+  SavedSlice ss;
+  std::array<uint32_t, 1> data;
+  std::fill(data.begin(), data.end(), 1234);
+  Status s = TensorSliceWriter::SaveData(data.data(), data.size(), &ss);
+  EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
+  EXPECT_TRUE(absl::StrContains(
+      s.error_message(),
+      "Tensor slice serialization not implemented for dtype"));
+}
+
 }  // namespace checkpoint
 
 }  // namespace tensorflow

Vulnerable code:       LOG(FATAL) << "MaxBytesPerElement not implemented for dtype: " << dt;
  return 0;
  size_t size_bound =
      ss->ByteSize() + kTensorProtoHeaderBytes +
      (MaxBytesPerElement(DataTypeToEnum<T>::value) * num_elements);
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
Fixed code:   size_t max_bytes_per_element =
      TensorSliceWriter::MaxBytesPerElementOrZero(dt);
  if (max_bytes_per_element == 0) {
    LOG(FATAL) << "MaxBytesPerElement not implemented for dtype: " << dt;
  }
  return max_bytes_per_element;
}

/* static */
size_t TensorSliceWriter::MaxBytesPerElementOrZero(DataType dt) {
      return 0;
  static size_t MaxBytesPerElementOrZero(DataType dt);

  size_t max_bytes_per_element =
      MaxBytesPerElementOrZero(DataTypeToEnum<T>::value);
  if (max_bytes_per_element == 0) {
    return errors::InvalidArgument(
        "Tensor slice serialization not implemented for dtype ",
        DataTypeToEnum<T>::value);
  }
  size_t size_bound = ss->ByteSize() + kTensorProtoHeaderBytes +
                      (max_bytes_per_element * num_elements);
#include <algorithm>
#include <memory>
#include <vector>
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
TEST(TensorSliceWriterTest, InvalidInput) {
  SavedSlice ss;
  std::array<uint32_t, 1> data;
  std::fill(data.begin(), data.end(), 1234);
  Status s = TensorSliceWriter::SaveData(data.data(), data.size(), &ss);
  EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
  EXPECT_TRUE(absl::StrContains(
      s.error_message(),
      "Tensor slice serialization not implemented for dtype"));
}