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/tensor_slice.cc, tensorflow/core/framework/tensor_slice.h, tensorflow/core/framework/tensor_slice_test.cc, tensorflow/core/util/tensor_slice_reader.cc, tensorflow/core/util/tensor_slice_reader_test.cc
Commit date: 2021-08-24 17:45:52 -0700
Commit message: Add BuildTensorSlice for building from unvalidated TensorSliceProtos.

This avoids several sources of crashes and undefined behavior when loading
invalid checkpoints.

PiperOrigin-RevId: 392785704
Change-Id: Icd9713c768b882f3b58b427eddac376060696833

Diff: diff --git a/tensorflow/core/framework/tensor_slice.cc b/tensorflow/core/framework/tensor_slice.cc
index bde1c09e7fce5..0b3914359062e 100644
--- a/tensorflow/core/framework/tensor_slice.cc
+++ b/tensorflow/core/framework/tensor_slice.cc
@@ -12,11 +12,14 @@ 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 "tensorflow/core/framework/tensor_slice.h"
+
+#include <limits>
 #include <vector>
+
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
@@ -42,10 +45,38 @@ TensorSlice::TensorSlice(
     starts_.push_back(e.first);
     lengths_.push_back(e.second);
   }
 }
 
+Status TensorSlice::BuildTensorSlice(const TensorSliceProto& proto,
+                                     TensorSlice* output) {
+  output->Clear();
+  output->starts_.reserve(proto.extent_size());
+  output->lengths_.reserve(proto.extent_size());
+  for (const auto& e : proto.extent()) {
+    int64_t l = GetExtentLength(e);
+    if (e.start() != 0 || l != kFullExtent) {
+      if (e.start() < 0 || l <= 0) {
+        return errors::InvalidArgument(
+            "Expected non-negative start and positive length but got start = ",
+            e.start(), ", length = ", l, ": extent = ", e.ShortDebugString());
+      }
+      // Calculating the extent end must not cause signed integer overflow.
+      if (static_cast<uint64_t>(e.start()) + static_cast<uint64_t>(e.length()) >
+          std::numeric_limits<int64_t>::max()) {
+        return errors::InvalidArgument(
+            "Extent end exceeds the maximum possible size: extent = ",
+            e.ShortDebugString());
+      }
+    }
+    output->starts_.push_back(e.start());
+    output->lengths_.push_back(l);
+  }
+
+  return Status::OK();
+}
+
 Status TensorSlice::Parse(const string& str, TensorSlice* slice) {
   std::vector<string> items = str_util::Split(str, ':', str_util::SkipEmpty());
   slice->starts_.reserve(items.size());
   slice->lengths_.reserve(items.size());
   for (const string& x : items) {
diff --git a/tensorflow/core/framework/tensor_slice.h b/tensorflow/core/framework/tensor_slice.h
index 18bf066e80d66..f67cc0996bdb3 100644
--- a/tensorflow/core/framework/tensor_slice.h
+++ b/tensorflow/core/framework/tensor_slice.h
@@ -46,10 +46,16 @@ class TensorSlice {
   explicit TensorSlice(int dim);
   explicit TensorSlice(const TensorSliceProto& proto);
   explicit TensorSlice(
       std::initializer_list<std::pair<int64_t, int64_t>> extents);
 
+  // This factory methods should be used instead of the constructor that takes a
+  // `TensorSliceProto` if calling code cannot validate that the sizes specify a
+  // valid `TensorSlice`.
+  static Status BuildTensorSlice(const TensorSliceProto& proto,
+                                 TensorSlice* output);
+
   static Status Parse(const string& str, TensorSlice* output);
   static TensorSlice ParseOrDie(const string& str) {
     TensorSlice ret;
     Status s = Parse(str, &ret);
     if (!s.ok()) {
diff --git a/tensorflow/core/framework/tensor_slice_test.cc b/tensorflow/core/framework/tensor_slice_test.cc
index ae8fa902d4863..8226c580e3ebc 100644
--- a/tensorflow/core/framework/tensor_slice_test.cc
+++ b/tensorflow/core/framework/tensor_slice_test.cc
@@ -13,10 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/framework/tensor_slice.h"
 
+#include <limits>
+
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
@@ -123,10 +125,52 @@ TEST(TensorSliceTest, Serialization) {
         "'19223372036854775808,19223372036854775808': string = "
         "19223372036854775808,19223372036854775808"));
   }
 }
 
+// Testing `BuildTensorSlice` with valid and invalid input protos.
+TEST(TensorSliceTest, BuildTensorSlice) {
+  TensorSliceProto proto;
+  TensorSlice({{0, -1}, {0, 10}, {14, 1}}).AsProto(&proto);
+  TensorSlice s;
+
+  // Successful building.
+  {
+    TF_ASSERT_OK(TensorSlice::BuildTensorSlice(proto, &s));
+    EXPECT_EQ("-:0,10:14,1", s.DebugString());
+  }
+
+  // Failed building due to negative extent start.
+  {
+    TensorSliceProto invalid_proto = proto;
+    invalid_proto.mutable_extent(0)->set_start(-1);
+    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
+  }
+
+  // Failed building due to negative extent length.
+  {
+    TensorSliceProto invalid_proto = proto;
+    invalid_proto.mutable_extent(2)->set_length(-1);
+    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
+  }
+
+  // Failed building due to missing extent length.
+  {
+    TensorSliceProto invalid_proto = proto;
+    invalid_proto.mutable_extent(2)->clear_length();
+    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
+  }
+
+  // Failed building due to extent end overflowing.
+  {
+    TensorSliceProto invalid_proto = proto;
+    invalid_proto.mutable_extent(2)->set_length(
+        std::numeric_limits<int64_t>::max());
+    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
+  }
+}
+
 // Testing the slice intersection
 TEST(TensorSliceTest, Intersection) {
   // "EVERYTHING" intersects with everything
   {
     TensorSlice a = TensorSlice::ParseOrDie("-:-");
diff --git a/tensorflow/core/util/tensor_slice_reader.cc b/tensorflow/core/util/tensor_slice_reader.cc
index da5ac17ee0eae..fd98d1b4041a3 100644
--- a/tensorflow/core/util/tensor_slice_reader.cc
+++ b/tensorflow/core/util/tensor_slice_reader.cc
@@ -170,11 +170,13 @@ void TensorSliceReader::LoadShard(int shard) const {
   for (const SavedSliceMeta& ssm : sts.meta().tensor()) {
     TensorShape ssm_shape;
     status_ = TensorShape::BuildTensorShapeBase(ssm.shape(), &ssm_shape);
     if (!status_.ok()) return;
     for (const TensorSliceProto& tsp : ssm.slice()) {
-      TensorSlice ss_slice(tsp);
+      TensorSlice ss_slice;
+      status_ = TensorSlice::BuildTensorSlice(tsp, &ss_slice);
+      if (!status_.ok()) return;
       status_ = RegisterTensorSlice(ssm.name(), ssm_shape, ssm.type(), fname,
                                     ss_slice, &tensors_);
       if (!status_.ok()) return;
     }
   }
diff --git a/tensorflow/core/util/tensor_slice_reader_test.cc b/tensorflow/core/util/tensor_slice_reader_test.cc
index 382e29ab32198..53993862385e3 100644
--- a/tensorflow/core/util/tensor_slice_reader_test.cc
+++ b/tensorflow/core/util/tensor_slice_reader_test.cc
@@ -434,10 +434,33 @@ TEST(TensorSliceReaderTest, NegativeTensorShapeDimension) {
   TensorSliceReader reader(fname, OpenTableTensorSliceReader);
   // The negative dimension should cause loading to fail.
   EXPECT_FALSE(reader.status().ok());
 }
 
+TEST(TensorSliceReaderTest, InvalidTensorSlice) {
+  const string fname =
+      io::JoinPath(testing::TmpDir(), "invalid_slice_checkpoint");
+  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
+  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  TF_CHECK_OK(writer.Add("test", TensorShape({4, 5}),
+                         TensorSlice::ParseOrDie("0,2:-"), data));
+  TF_CHECK_OK(writer.Finish());
+
+  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
+    if (sts.has_meta()) {
+      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
+        tensor.mutable_slice(0)->mutable_extent(0)->set_length(-10);
+      }
+    }
+    return sts.SerializeAsString();
+  });
+
+  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
+  // The negative exent length should cause loading to fail.
+  EXPECT_FALSE(reader.status().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:       TensorSlice ss_slice(tsp);
Fixed code: 
#include <limits>

Status TensorSlice::BuildTensorSlice(const TensorSliceProto& proto,
                                     TensorSlice* output) {
  output->Clear();
  output->starts_.reserve(proto.extent_size());
  output->lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    int64_t l = GetExtentLength(e);
    if (e.start() != 0 || l != kFullExtent) {
      if (e.start() < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and positive length but got start = ",
            e.start(), ", length = ", l, ": extent = ", e.ShortDebugString());
      }
      // Calculating the extent end must not cause signed integer overflow.
      if (static_cast<uint64_t>(e.start()) + static_cast<uint64_t>(e.length()) >
          std::numeric_limits<int64_t>::max()) {
        return errors::InvalidArgument(
            "Extent end exceeds the maximum possible size: extent = ",
            e.ShortDebugString());
      }
    }
    output->starts_.push_back(e.start());
    output->lengths_.push_back(l);
  }

  return Status::OK();
}

  // This factory methods should be used instead of the constructor that takes a
  // `TensorSliceProto` if calling code cannot validate that the sizes specify a
  // valid `TensorSlice`.
  static Status BuildTensorSlice(const TensorSliceProto& proto,
                                 TensorSlice* output);

#include <limits>

// Testing `BuildTensorSlice` with valid and invalid input protos.
TEST(TensorSliceTest, BuildTensorSlice) {
  TensorSliceProto proto;
  TensorSlice({{0, -1}, {0, 10}, {14, 1}}).AsProto(&proto);
  TensorSlice s;

  // Successful building.
  {
    TF_ASSERT_OK(TensorSlice::BuildTensorSlice(proto, &s));
    EXPECT_EQ("-:0,10:14,1", s.DebugString());
  }

  // Failed building due to negative extent start.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(0)->set_start(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to negative extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to missing extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->clear_length();
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to extent end overflowing.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(
        std::numeric_limits<int64_t>::max());
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }
}

      TensorSlice ss_slice;
      status_ = TensorSlice::BuildTensorSlice(tsp, &ss_slice);
      if (!status_.ok()) return;
TEST(TensorSliceReaderTest, InvalidTensorSlice) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "invalid_slice_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_CHECK_OK(writer.Add("test", TensorShape({4, 5}),
                         TensorSlice::ParseOrDie("0,2:-"), data));
  TF_CHECK_OK(writer.Finish());

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

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  // The negative exent length should cause loading to fail.
  EXPECT_FALSE(reader.status().ok());
}