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/framework/BUILD, tensorflow/core/framework/tensor.cc, tensorflow/core/framework/tensor.h, tensorflow/core/util/tensor_slice_reader.cc, tensorflow/core/util/tensor_slice_reader_test.cc
Commit date: 2021-08-24 10:51:32 -0700
Commit message: Prevent crashes when loading tensor slices with unsupported types.

Also fix the `Tensor(const TensorShape&)` constructor swapping the LOG(FATAL)
messages for the unset and unsupported types.

PiperOrigin-RevId: 392695027
Change-Id: I4beda7db950db951d273e3259a7c8534ece49354

Diff: diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 88550b3205ff5..11fe921173ae4 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -833,10 +833,11 @@ tf_cuda_library(
         "//tensorflow/core/lib/gtl:inlined_vector",
         "//tensorflow/core/lib/hash",
         "//tensorflow/core/lib/strings:str_util",
         "//tensorflow/core/lib/strings:strcat",
         "//tensorflow/core/platform:abi",
+        "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:logging",
         "//tensorflow/core/platform:macros",
         "//tensorflow/core/platform:platform_port",
         "//tensorflow/core/platform:protobuf",
         "//tensorflow/core/platform:strcat",
diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc
index 2abd8fce1c427..5e26bf05c03ca 100644
--- a/tensorflow/core/framework/tensor.cc
+++ b/tensorflow/core/framework/tensor.cc
@@ -50,10 +50,11 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/tensor_coding.h"
 #include "tensorflow/core/platform/types.h"
@@ -721,15 +722,15 @@ bool Tensor::RefCountIsOne() const {
 }
 
 // The macro CASES() expands to a switch statement conditioned on
 // TYPE_ENUM. Each case expands the STMTS after a typedef for T.
 #define SINGLE_ARG(...) __VA_ARGS__
-#define CASE(TYPE, STMTS)             \
-  case DataTypeToEnum<TYPE>::value: { \
-    typedef TYPE T;                   \
-    STMTS;                            \
-    break;                            \
+#define CASE(TYPE, STMTS)               \
+  case DataTypeToEnum<TYPE>::value: {   \
+    typedef TF_ATTRIBUTE_UNUSED TYPE T; \
+    STMTS;                              \
+    break;                              \
   }
 #define CASES_WITH_DEFAULT(TYPE_ENUM, STMTS, INVALID, DEFAULT) \
   switch (TYPE_ENUM) {                                         \
     CASE(float, SINGLE_ARG(STMTS))                             \
     CASE(double, SINGLE_ARG(STMTS))                            \
@@ -761,13 +762,12 @@ bool Tensor::RefCountIsOne() const {
       DEFAULT;                                                 \
       break;                                                   \
   }
 
 #define CASES(TYPE_ENUM, STMTS)                                      \
-  CASES_WITH_DEFAULT(TYPE_ENUM, STMTS,                               \
-                     LOG(FATAL) << "Unexpected type: " << TYPE_ENUM; \
-                     , LOG(FATAL) << "Type not set";)
+  CASES_WITH_DEFAULT(TYPE_ENUM, STMTS, LOG(FATAL) << "Type not set"; \
+                     , LOG(FATAL) << "Unexpected type: " << TYPE_ENUM;)
 
 Tensor::Tensor(Allocator* a, DataType type, const TensorShape& shape)
     : shape_(shape), buf_(nullptr) {
   set_dtype(type);
   CHECK_NOTNULL(a);
@@ -793,10 +793,20 @@ Tensor::Tensor(Allocator* a, DataType type, const TensorShape& shape,
     LogMemory::RecordTensorAllocation("Unknown (with attributes)",
                                       LogMemory::UNKNOWN_STEP_ID, *this);
   }
 }
 
+Status Tensor::BuildTensor(DataType type, const TensorShape& shape,
+                           Tensor* out_tensor) {
+  // Avoid crashes due to invalid or unsupported types.
+  CASES_WITH_DEFAULT(
+      type, {}, return errors::InvalidArgument("Type not set"),
+      return errors::InvalidArgument("Unexpected type: ", DataType_Name(type)));
+  *out_tensor = Tensor(type, shape);
+  return Status::OK();
+}
+
 // NOTE(mrry): The default allocator for a Tensor (when none is specified) is
 // the default CPU allocator for NUMA zone 0. Accessing that currently involves
 // acquiring a lock, which guards initialization of the per-NUMA zone
 // allocators, and becomes highly contended.
 //
diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h
index dd66f995ecbcd..d4ffb1106e40c 100644
--- a/tensorflow/core/framework/tensor.h
+++ b/tensorflow/core/framework/tensor.h
@@ -168,10 +168,19 @@ class Tensor {
   /// Like Tensor(), returns a 1-dimensional, 0-element Tensor with
   /// IsInitialized() returning True. See the Tensor() documentation
   /// for details.
   explicit Tensor(DataType type);
 
+  /// \brief Initializes a tensor with the input `type` and `shape`, or returns
+  /// an error and leaves `out_tensor` unmodified. This factory method should be
+  /// used instead of the corresponding constructor if calling code cannot
+  /// validate that the `DataType` is valid and supported.
+  ///
+  /// The underlying buffer is allocated using a `CPUAllocator`.
+  static Status BuildTensor(DataType type, const TensorShape& shape,
+                            Tensor* out_tensor);
+
  private:
   // A tag type for selecting the `Tensor` constructor overload that creates a
   // scalar tensor in host memory.
   struct host_scalar_tag {};
 
diff --git a/tensorflow/core/util/tensor_slice_reader.cc b/tensorflow/core/util/tensor_slice_reader.cc
index 58c4c22ce7a9e..00911b13f3491 100644
--- a/tensorflow/core/util/tensor_slice_reader.cc
+++ b/tensorflow/core/util/tensor_slice_reader.cc
@@ -246,11 +246,13 @@ Status TensorSliceReader::GetTensor(
     type = tss->type();
     shape = tss->shape();
     slice = tss->Slices().begin()->second.slice;
   }
 
-  std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor(type, shape));
+  std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor);
+  Status s = tensorflow::Tensor::BuildTensor(type, shape, t.get());
+  if (!s.ok()) return s;
   bool success = false;
 
 #define READER_COPY(dt)                                                  \
   case dt:                                                               \
     success = CopySliceData(name, slice,                                 \
diff --git a/tensorflow/core/util/tensor_slice_reader_test.cc b/tensorflow/core/util/tensor_slice_reader_test.cc
index e6e65fc2e6c96..9bd3063d4ebec 100644
--- a/tensorflow/core/util/tensor_slice_reader_test.cc
+++ b/tensorflow/core/util/tensor_slice_reader_test.cc
@@ -11,27 +11,32 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <utility>
-
 #include "tensorflow/core/util/tensor_slice_reader.h"
 
+#include <utility>
+#include <vector>
+
 #include "tensorflow/core/framework/types.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/iterator.h"
 #include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/io/table.h"
+#include "tensorflow/core/lib/io/table_builder.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/version.h"
+#include "tensorflow/core/util/saved_tensor_slice.pb.h"
 #include "tensorflow/core/util/saved_tensor_slice_util.h"
 #include "tensorflow/core/util/tensor_slice_reader_cache.h"
 #include "tensorflow/core/util/tensor_slice_writer.h"
 
 namespace tensorflow {
@@ -307,10 +312,106 @@ TEST_SIMPLE_INT(int32, int32)
 TEST_SIMPLE_INT(int64_t, int64_t)
 TEST_SIMPLE_INT(int16, int32)
 TEST_SIMPLE_INT(int8, int32)
 TEST_SIMPLE_INT(uint8, int32)
 
+// Modifies the SavedTensorSlices messages in a checkpoint to allow creating
+// malformed or unsupported checkpoints.
+void MutateSavedTensorSlices(
+    const std::string& fname,
+    const std::function<std::string(SavedTensorSlices)>& mutator) {
+  table::Options options;
+  options.compression = table::kNoCompression;
+
+  // Read all entres from the table.
+  std::vector<std::pair<std::string, std::string>> entries;
+  {
+    std::unique_ptr<RandomAccessFile> file;
+    TF_CHECK_OK(Env::Default()->NewRandomAccessFile(fname, &file));
+    uint64 file_size;
+    TF_CHECK_OK(Env::Default()->GetFileSize(fname, &file_size));
+    table::Table* t;
+    TF_CHECK_OK(table::Table::Open(options, file.get(), file_size, &t));
+    std::unique_ptr<table::Table> table(t);
+    std::unique_ptr<table::Iterator> it(table->NewIterator());
+    for (it->Seek(""); it->Valid(); it->Next()) {
+      entries.emplace_back(it->key(), it->value());
+    }
+    TF_CHECK_OK(it->status());
+  }
+
+  // Rewrite the table, mutating each value.
+  {
+    std::unique_ptr<WritableFile> file;
+    TF_CHECK_OK(Env::Default()->NewWritableFile(fname, &file));
+    table::TableBuilder builder(options, file.get());
+    for (const auto& entry : entries) {
+      SavedTensorSlices sts;
+      CHECK(sts.ParseFromString(entry.second));
+      builder.Add(entry.first, mutator(std::move(sts)));
+    }
+    TF_CHECK_OK(builder.Finish());
+    TF_CHECK_OK(file->Close());
+  }
+}
+
+TEST(TensorSliceReaderTest, MissingTensorType) {
+  const string fname = io::JoinPath(testing::TmpDir(), "invalid_checkpoint");
+  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
+  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  TensorShape shape({4, 5});
+  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
+  TF_CHECK_OK(writer.Add("test", shape, slice, data));
+  TF_CHECK_OK(writer.Finish());
+
+  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
+    if (sts.has_meta()) {
+      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
+        tensor.clear_type();
+      }
+    }
+    return sts.SerializeAsString();
+  });
+
+  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
+  TF_CHECK_OK(reader.status());
+
+  // The tensor should be present, but loading it should fail due to the
+  // unset (invalid) type.
+  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
+  std::unique_ptr<Tensor> tensor;
+  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
+}
+
+TEST(TensorSliceReaderTest, UnsupportedTensorType) {
+  const string fname = io::JoinPath(testing::TmpDir(), "int32_ref_checkpoint");
+  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
+  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  TensorShape shape({4, 5});
+  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
+  TF_CHECK_OK(writer.Add("test", shape, slice, data));
+  TF_CHECK_OK(writer.Finish());
+
+  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
+    if (sts.has_meta()) {
+      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
+        tensor.set_type(DT_INT32_REF);
+      }
+    }
+    return sts.SerializeAsString();
+  });
+
+  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
+  TF_CHECK_OK(reader.status());
+
+  // The tensor should be present, but loading it should fail due to the
+  // unsupported type.
+  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");
 

Vulnerable code: #define CASE(TYPE, STMTS)             \
  case DataTypeToEnum<TYPE>::value: { \
    typedef TYPE T;                   \
    STMTS;                            \
    break;                            \
  CASES_WITH_DEFAULT(TYPE_ENUM, STMTS,                               \
                     LOG(FATAL) << "Unexpected type: " << TYPE_ENUM; \
                     , LOG(FATAL) << "Type not set";)
  std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor(type, shape));
#include <utility>

Fixed code:         "//tensorflow/core/platform:errors",
#include "tensorflow/core/platform/errors.h"
#define CASE(TYPE, STMTS)               \
  case DataTypeToEnum<TYPE>::value: {   \
    typedef TF_ATTRIBUTE_UNUSED TYPE T; \
    STMTS;                              \
    break;                              \
  CASES_WITH_DEFAULT(TYPE_ENUM, STMTS, LOG(FATAL) << "Type not set"; \
                     , LOG(FATAL) << "Unexpected type: " << TYPE_ENUM;)
Status Tensor::BuildTensor(DataType type, const TensorShape& shape,
                           Tensor* out_tensor) {
  // Avoid crashes due to invalid or unsupported types.
  CASES_WITH_DEFAULT(
      type, {}, return errors::InvalidArgument("Type not set"),
      return errors::InvalidArgument("Unexpected type: ", DataType_Name(type)));
  *out_tensor = Tensor(type, shape);
  return Status::OK();
}

  /// \brief Initializes a tensor with the input `type` and `shape`, or returns
  /// an error and leaves `out_tensor` unmodified. This factory method should be
  /// used instead of the corresponding constructor if calling code cannot
  /// validate that the `DataType` is valid and supported.
  ///
  /// The underlying buffer is allocated using a `CPUAllocator`.
  static Status BuildTensor(DataType type, const TensorShape& shape,
                            Tensor* out_tensor);

  std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor);
  Status s = tensorflow::Tensor::BuildTensor(type, shape, t.get());
  if (!s.ok()) return s;
#include <utility>
#include <vector>

#include "tensorflow/core/lib/io/iterator.h"
#include "tensorflow/core/lib/io/table.h"
#include "tensorflow/core/lib/io/table_builder.h"
#include "tensorflow/core/util/saved_tensor_slice.pb.h"
// Modifies the SavedTensorSlices messages in a checkpoint to allow creating
// malformed or unsupported checkpoints.
void MutateSavedTensorSlices(
    const std::string& fname,
    const std::function<std::string(SavedTensorSlices)>& mutator) {
  table::Options options;
  options.compression = table::kNoCompression;

  // Read all entres from the table.
  std::vector<std::pair<std::string, std::string>> entries;
  {
    std::unique_ptr<RandomAccessFile> file;
    TF_CHECK_OK(Env::Default()->NewRandomAccessFile(fname, &file));
    uint64 file_size;
    TF_CHECK_OK(Env::Default()->GetFileSize(fname, &file_size));
    table::Table* t;
    TF_CHECK_OK(table::Table::Open(options, file.get(), file_size, &t));
    std::unique_ptr<table::Table> table(t);
    std::unique_ptr<table::Iterator> it(table->NewIterator());
    for (it->Seek(""); it->Valid(); it->Next()) {
      entries.emplace_back(it->key(), it->value());
    }
    TF_CHECK_OK(it->status());
  }

  // Rewrite the table, mutating each value.
  {
    std::unique_ptr<WritableFile> file;
    TF_CHECK_OK(Env::Default()->NewWritableFile(fname, &file));
    table::TableBuilder builder(options, file.get());
    for (const auto& entry : entries) {
      SavedTensorSlices sts;
      CHECK(sts.ParseFromString(entry.second));
      builder.Add(entry.first, mutator(std::move(sts)));
    }
    TF_CHECK_OK(builder.Finish());
    TF_CHECK_OK(file->Close());
  }
}

TEST(TensorSliceReaderTest, MissingTensorType) {
  const string fname = io::JoinPath(testing::TmpDir(), "invalid_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TensorShape shape({4, 5});
  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(writer.Add("test", shape, slice, data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.clear_type();
      }
    }
    return sts.SerializeAsString();
  });

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

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

TEST(TensorSliceReaderTest, UnsupportedTensorType) {
  const string fname = io::JoinPath(testing::TmpDir(), "int32_ref_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TensorShape shape({4, 5});
  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(writer.Add("test", shape, slice, data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.set_type(DT_INT32_REF);
      }
    }
    return sts.SerializeAsString();
  });

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

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