CVE-2021-41227: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. In affected versions the `ImmutableConst` operation in TensorFlow can be tricked into reading arbitrary memory contents. This is because the `tstring` TensorFlow string class has a special case for memory mapped strings but the operation itself does not offer any support for this datatype. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit 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 reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 5.5
Published: 2021-11-05T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/immutable_constant_op.cc, tensorflow/core/kernels/immutable_constant_op_test.cc
Commit date: 2021-10-05 13:46:10 -0700
Commit message: Add error checking to ImmutableConst OP that strings are not yet supported.

PiperOrigin-RevId: 401065359
Change-Id: I9dd2bd2a2c36f22f4a05153daf6ebdc4613469d2

Diff: diff --git a/tensorflow/core/kernels/immutable_constant_op.cc b/tensorflow/core/kernels/immutable_constant_op.cc
index 19aa865c1fbe4..df0d76ce633e9 100644
--- a/tensorflow/core/kernels/immutable_constant_op.cc
+++ b/tensorflow/core/kernels/immutable_constant_op.cc
@@ -98,10 +98,13 @@ void ImmutableConstantOp::Compute(OpKernelContext* ctx) {
   std::unique_ptr<MemmappedTensorAllocator> allocator(
       new MemmappedTensorAllocator());
 
   OP_REQUIRES_OK(ctx,
                  allocator->InitializeFromRegion(region_name_, ctx->env()));
+  OP_REQUIRES(ctx, dtype_ != DT_STRING,
+              errors::Unimplemented("Sorry, DT_STRING is not currently "
+                                    "supported for ImmutableConstOp."));
   ctx->set_output(0, Tensor(allocator.get(), dtype_, shape_));
   OP_REQUIRES_OK(ctx, allocator->allocation_status());
   // Allocator is owned by the tensor from this point.
   allocator.release()->set_delete_on_deallocate();
 }
diff --git a/tensorflow/core/kernels/immutable_constant_op_test.cc b/tensorflow/core/kernels/immutable_constant_op_test.cc
index d52a8b55a35d7..40ce8918a39ad 100644
--- a/tensorflow/core/kernels/immutable_constant_op_test.cc
+++ b/tensorflow/core/kernels/immutable_constant_op_test.cc
@@ -144,11 +144,12 @@ TEST(ImmutableConstantOpTest, ExecutionError) {
   EXPECT_EQ(
       session->Run({}, {result.node()->name() + ":0"}, {}, &outputs).code(),
       error::INTERNAL);
 }
 
-Status CreateTempFile(Env* env, float value, uint64 size, string* filename) {
+Status CreateTempFileFloat(Env* env, float value, uint64 size,
+                           string* filename) {
   const string dir = testing::TmpDir();
   *filename = io::JoinPath(dir, strings::StrCat("file_", value));
   std::unique_ptr<WritableFile> file;
   TF_RETURN_IF_ERROR(env->NewWritableFile(*filename, &file));
   for (uint64 i = 0; i < size; ++i) {
@@ -164,12 +165,12 @@ TEST(ImmutableConstantOpTest, FromFile) {
   const TensorShape kFileTensorShape({1000, 1});
   Env* env = Env::Default();
   auto root = Scope::NewRootScope().ExitOnError();
 
   string two_file, three_file;
-  TF_ASSERT_OK(CreateTempFile(env, 2.0f, 1000, &two_file));
-  TF_ASSERT_OK(CreateTempFile(env, 3.0f, 1000, &three_file));
+  TF_ASSERT_OK(CreateTempFileFloat(env, 2.0f, 1000, &two_file));
+  TF_ASSERT_OK(CreateTempFileFloat(env, 3.0f, 1000, &three_file));
   auto node1 = ops::ImmutableConst(root, DT_FLOAT, kFileTensorShape, two_file);
   auto node2 =
       ops::ImmutableConst(root, DT_FLOAT, kFileTensorShape, three_file);
   auto result = ops::MatMul(root, node1, node2, ops::MatMul::TransposeB(true));
 
@@ -188,7 +189,41 @@ TEST(ImmutableConstantOpTest, FromFile) {
   EXPECT_EQ(outputs.front().flat<float>()(0), 2.0f * 3.0f);
   EXPECT_EQ(outputs.front().flat<float>()(1), 2.0f * 3.0f);
   EXPECT_EQ(outputs.front().flat<float>()(2), 2.0f * 3.0f);
 }
 
+Status CreateTempFileBadString(Env* env, char value, uint64 size,
+                               const string suffix, string* filename) {
+  const string dir = testing::TmpDir();
+  *filename = io::JoinPath(dir, strings::StrCat("file_", suffix));
+  std::unique_ptr<WritableFile> file;
+  TF_RETURN_IF_ERROR(env->NewWritableFile(*filename, &file));
+  TF_RETURN_IF_ERROR(file->Append(std::string(size, value)));
+  TF_RETURN_IF_ERROR(file->Close());
+  return Status::OK();
+}
+
+TEST(ImmutableConstantOpTest, FromFileStringUnimplmented) {
+  const TensorShape kFileTensorShape({1});
+  Env* env = Env::Default();
+  auto root = Scope::NewRootScope().ExitOnError();
+
+  string bad_file;
+  TF_ASSERT_OK(CreateTempFileBadString(env, '\xe2', 128, "bad_e2", &bad_file));
+  auto result =
+      ops::ImmutableConst(root, DT_STRING, kFileTensorShape, bad_file);
+  GraphDef graph_def;
+  TF_ASSERT_OK(root.ToGraphDef(&graph_def));
+  SessionOptions session_options;
+  session_options.env = Env::Default();
+  std::unique_ptr<Session> session(NewSession(session_options));
+  ASSERT_TRUE(session != nullptr) << "Failed to create session";
+  TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph";
+  std::vector<Tensor> outputs;
+  // Check that the run returned error.
+  EXPECT_EQ(
+      session->Run({}, {result.node()->name() + ":0"}, {}, &outputs).code(),
+      error::UNIMPLEMENTED);
+}
+
 }  // namespace
 }  // namespace tensorflow

Vulnerable code: Status CreateTempFile(Env* env, float value, uint64 size, string* filename) {
  TF_ASSERT_OK(CreateTempFile(env, 2.0f, 1000, &two_file));
  TF_ASSERT_OK(CreateTempFile(env, 3.0f, 1000, &three_file));
Fixed code:   OP_REQUIRES(ctx, dtype_ != DT_STRING,
              errors::Unimplemented("Sorry, DT_STRING is not currently "
                                    "supported for ImmutableConstOp."));
Status CreateTempFileFloat(Env* env, float value, uint64 size,
                           string* filename) {
  TF_ASSERT_OK(CreateTempFileFloat(env, 2.0f, 1000, &two_file));
  TF_ASSERT_OK(CreateTempFileFloat(env, 3.0f, 1000, &three_file));
Status CreateTempFileBadString(Env* env, char value, uint64 size,
                               const string suffix, string* filename) {
  const string dir = testing::TmpDir();
  *filename = io::JoinPath(dir, strings::StrCat("file_", suffix));
  std::unique_ptr<WritableFile> file;
  TF_RETURN_IF_ERROR(env->NewWritableFile(*filename, &file));
  TF_RETURN_IF_ERROR(file->Append(std::string(size, value)));
  TF_RETURN_IF_ERROR(file->Close());
  return Status::OK();
}

TEST(ImmutableConstantOpTest, FromFileStringUnimplmented) {
  const TensorShape kFileTensorShape({1});
  Env* env = Env::Default();
  auto root = Scope::NewRootScope().ExitOnError();

  string bad_file;
  TF_ASSERT_OK(CreateTempFileBadString(env, '\xe2', 128, "bad_e2", &bad_file));
  auto result =
      ops::ImmutableConst(root, DT_STRING, kFileTensorShape, bad_file);
  GraphDef graph_def;
  TF_ASSERT_OK(root.ToGraphDef(&graph_def));
  SessionOptions session_options;
  session_options.env = Env::Default();
  std::unique_ptr<Session> session(NewSession(session_options));
  ASSERT_TRUE(session != nullptr) << "Failed to create session";
  TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph";
  std::vector<Tensor> outputs;
  // Check that the run returned error.
  EXPECT_EQ(
      session->Run({}, {result.node()->name() + ":0"}, {}, &outputs).code(),
      error::UNIMPLEMENTED);
}