CVE-2021-41203: Integer Overflow or Wraparound (CWE-190)
value: TensorFlow is an open source platform for machine learning. In affected versions an attacker can trigger undefined behavior, integer overflows, segfaults and `CHECK`-fail crashes if they can change saved checkpoints from outside of TensorFlow. This is because the checkpoints loading infrastructure is missing validation for invalid file formats. The fixes will be included in TensorFlow 2.7.0. We will also cherrypick these commits on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
CVSS v3 base score: 7.8
Published: 2021-11-05T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C
Files changed: tensorflow/core/util/saved_tensor_slice_util.h, tensorflow/core/util/tensor_slice_reader.h, tensorflow/core/util/tensor_slice_reader_test.cc
Commit date: 2021-08-25 13:49:17 -0700
Commit message: Avoid buffer overflow when loading tensors with insufficient data from checkpoints.

`CopyDataFromTensorSliceToTensorSlice` does not (and cannot conveniently)
provide any bounds checking on its own, so the size is instead checked prior
to passing unvalidated data to that function.

PiperOrigin-RevId: 392971286
Change-Id: If2073b36d4d5eedd386329f56729395fd7effee1

Diff: diff --git a/tensorflow/core/util/saved_tensor_slice_util.h b/tensorflow/core/util/saved_tensor_slice_util.h
index d0b08b3c294e0..b992061a15df9 100644
--- a/tensorflow/core/util/saved_tensor_slice_util.h
+++ b/tensorflow/core/util/saved_tensor_slice_util.h
@@ -57,10 +57,13 @@ Status ParseShapeAndSlice(const string& shape_and_slice, TensorShape* shape,
                           TensorSlice* slice, TensorShape* shape_slice);
 
 template <typename T>
 struct SaveTypeTraits;
 
+template <typename T>
+int TensorProtoDataSize(const TensorProto& t);
+
 template <typename T>
 const typename SaveTypeTraits<T>::SavedType* TensorProtoData(
     const TensorProto& t);
 
 template <typename T>
@@ -93,19 +96,27 @@ void Fill(T* data, size_t n, TensorProto* t);
   }
 
 #define TENSOR_PROTO_EXTRACT_TYPE(TYPE, FIELD, FTYPE)             \
   TENSOR_PROTO_EXTRACT_TYPE_HELPER(TYPE, FIELD, FTYPE, FTYPE)     \
   template <>                                                     \
+  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {    \
+    return t.FIELD##_val_size();                                  \
+  }                                                               \
+  template <>                                                     \
   inline void Fill(const TYPE* data, size_t n, TensorProto* t) {  \
     typename protobuf::RepeatedField<FTYPE> copy(data, data + n); \
     t->mutable_##FIELD##_val()->Swap(&copy);                      \
   }
 
 // Complex needs special treatment since proto doesn't have native complex
 #define TENSOR_PROTO_EXTRACT_TYPE_COMPLEX(TYPE, FIELD, FTYPE)       \
   TENSOR_PROTO_EXTRACT_TYPE_HELPER(TYPE, FIELD, FTYPE, TYPE)        \
   template <>                                                       \
+  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {      \
+    return t.FIELD##_val_size() / 2;                                \
+  }                                                                 \
+  template <>                                                       \
   inline void Fill(const TYPE* data, size_t n, TensorProto* t) {    \
     const FTYPE* sub = reinterpret_cast<const FTYPE*>(data);        \
     typename protobuf::RepeatedField<FTYPE> copy(sub, sub + 2 * n); \
     t->mutable_##FIELD##_val()->Swap(&copy);                        \
   }
@@ -134,10 +145,15 @@ TENSOR_PROTO_EXTRACT_TYPE(quint16, int, int32);
 // Custom implementation for qint32, based on the one for int32.
 
 template <>
 struct SaveTypeTraits<qint32> : SaveTypeTraits<int32> {};
 
+template <>
+inline int TensorProtoDataSize<qint32>(const TensorProto& t) {
+  return t.int_val_size();
+}
+
 template <>
 inline const int32* TensorProtoData<qint32>(const TensorProto& t) {
   static_assert(SaveTypeTraits<qint32>::supported,
                 "Specified type qint32 not supported for Restore");
   return reinterpret_cast<const int32*>(t.int_val().data());
@@ -156,10 +172,15 @@ struct SaveTypeTraits<Eigen::half> {
   static constexpr bool supported = true;
   typedef int SavedType;
   typedef protobuf::RepeatedField<int32> RepeatedField;
 };
 
+template <>
+inline int TensorProtoDataSize<Eigen::half>(const TensorProto& t) {
+  return t.half_val_size();
+}
+
 template <>
 inline const int* TensorProtoData<Eigen::half>(const TensorProto& t) {
   return t.half_val().data();
 }
 
@@ -185,10 +206,15 @@ struct SaveTypeTraits<tstring> {
   static constexpr bool supported = true;
   typedef const string* SavedType;
   typedef protobuf::RepeatedPtrField<string> RepeatedField;
 };
 
+template <>
+inline int TensorProtoDataSize<tstring>(const TensorProto& t) {
+  return t.string_val_size();
+}
+
 template <>
 inline const string* const* TensorProtoData<tstring>(const TensorProto& t) {
   static_assert(SaveTypeTraits<tstring>::supported,
                 "Specified type tstring not supported for Restore");
   return t.string_val().data();
diff --git a/tensorflow/core/util/tensor_slice_reader.h b/tensorflow/core/util/tensor_slice_reader.h
index 0fb2e11bf8dd0..bc0a91523fe36 100644
--- a/tensorflow/core/util/tensor_slice_reader.h
+++ b/tensorflow/core/util/tensor_slice_reader.h
@@ -179,10 +179,26 @@ bool TensorSliceReader::CopySliceData(const string& name,
     if (!ParseProtoUnlimited(&sts, value)) {
       VLOG(1) << "Failed to parse the record for tensor " << name << ", slice "
               << slice_s.DebugString() << ": computed key = " << key;
       return false;
     }
+    // Ensure the TensorSlice contains the expected amount of data.
+    TensorShape shp_s;
+    Status s = slice_s.SliceTensorShape(tss->shape(), &shp_s);
+    if (!s.ok()) {
+      VLOG(1) << "Failed to slice tensor " << name << ", slice "
+              << slice_s.DebugString() << ": " << s;
+      return false;
+    }
+    if (checkpoint::TensorProtoDataSize<T>(sts.data().data()) !=
+        shp_s.num_elements()) {
+      VLOG(1) << "Tensor " << name << ", slice " << slice_s.DebugString()
+              << " had an unexpected amount of data: expected = "
+              << shp_s.num_elements() << ", got = "
+              << checkpoint::TensorProtoDataSize<T>(sts.data().data());
+      return false;
+    }
     CopyDataFromTensorSliceToTensorSlice(
         tss->shape(), slice_s, slice,
         checkpoint::TensorProtoData<T>(sts.data().data()), data);
   }
   return true;
diff --git a/tensorflow/core/util/tensor_slice_reader_test.cc b/tensorflow/core/util/tensor_slice_reader_test.cc
index 53993862385e3..efb2d9d9748fd 100644
--- a/tensorflow/core/util/tensor_slice_reader_test.cc
+++ b/tensorflow/core/util/tensor_slice_reader_test.cc
@@ -457,10 +457,37 @@ TEST(TensorSliceReaderTest, InvalidTensorSlice) {
   TensorSliceReader reader(fname, OpenTableTensorSliceReader);
   // The negative exent length should cause loading to fail.
   EXPECT_FALSE(reader.status().ok());
 }
 
+TEST(TensorSliceReaderTest, MissingTensorData) {
+  const string fname =
+      io::JoinPath(testing::TmpDir(), "missing_data_checkpoint");
+  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
+  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  TF_ASSERT_OK(writer.Add("test", TensorShape({4, 5}),
+                          TensorSlice::ParseOrDie("0,2:-"), data));
+  TF_ASSERT_OK(writer.Finish());
+
+  MutateSavedTensorSlices(fname, [&](SavedTensorSlices sts) {
+    if (sts.has_data()) {
+      // Replace the data with only 4 elements.
+      Fill(data, 4, sts.mutable_data()->mutable_data());
+    }
+    return sts.SerializeAsString();
+  });
+
+  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
+  TF_ASSERT_OK(reader.status());
+
+  // The tensor should be present, but loading it should fail due to the missing
+  // data.
+  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
+  std::unique_ptr<Tensor> tensor;
+  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
+}
+
 void CachedTensorSliceReaderTesterHelper(
     const TensorSliceWriter::CreateBuilderFunction& create_function,
     const TensorSliceReader::OpenTableFunction& open_function) {
   const string fname_base = io::JoinPath(testing::TmpDir(), "float_checkpoint");
 

Fixed code: template <typename T>
int TensorProtoDataSize(const TensorProto& t);

  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {    \
    return t.FIELD##_val_size();                                  \
  }                                                               \
  template <>                                                     \
  inline int TensorProtoDataSize<TYPE>(const TensorProto& t) {      \
    return t.FIELD##_val_size() / 2;                                \
  }                                                                 \
  template <>                                                       \
template <>
inline int TensorProtoDataSize<qint32>(const TensorProto& t) {
  return t.int_val_size();
}

template <>
inline int TensorProtoDataSize<Eigen::half>(const TensorProto& t) {
  return t.half_val_size();
}

template <>
inline int TensorProtoDataSize<tstring>(const TensorProto& t) {
  return t.string_val_size();
}

    // Ensure the TensorSlice contains the expected amount of data.
    TensorShape shp_s;
    Status s = slice_s.SliceTensorShape(tss->shape(), &shp_s);
    if (!s.ok()) {
      VLOG(1) << "Failed to slice tensor " << name << ", slice "
              << slice_s.DebugString() << ": " << s;
      return false;
    }
    if (checkpoint::TensorProtoDataSize<T>(sts.data().data()) !=
        shp_s.num_elements()) {
      VLOG(1) << "Tensor " << name << ", slice " << slice_s.DebugString()
              << " had an unexpected amount of data: expected = "
              << shp_s.num_elements() << ", got = "
              << checkpoint::TensorProtoDataSize<T>(sts.data().data());
      return false;
    }
TEST(TensorSliceReaderTest, MissingTensorData) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "missing_data_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_ASSERT_OK(writer.Add("test", TensorShape({4, 5}),
                          TensorSlice::ParseOrDie("0,2:-"), data));
  TF_ASSERT_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [&](SavedTensorSlices sts) {
    if (sts.has_data()) {
      // Replace the data with only 4 elements.
      Fill(data, 4, sts.mutable_data()->mutable_data());
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  TF_ASSERT_OK(reader.status());

  // The tensor should be present, but loading it should fail due to the missing
  // data.
  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
  std::unique_ptr<Tensor> tensor;
  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
}