Cherrypick fixes for ImmutableConst op vuln

tensorflow/core/kernels/immutable_constant_op.cc

@@ -100,6 +100,9 @@ void ImmutableConstantOp::Compute(OpKernelContext* ctx) {
 
   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.

tensorflow/core/kernels/immutable_constant_op_test.cc

@@ -146,7 +146,8 @@ TEST(ImmutableConstantOpTest, ExecutionError) {
       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;
@@ -166,8 +167,8 @@ TEST(ImmutableConstantOpTest, FromFile) {
   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);
@@ -190,5 +191,39 @@ TEST(ImmutableConstantOpTest, FromFile) {
   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

tensorflow/core/platform/ctstring_internal.h

@@ -63,9 +63,9 @@ static inline uint32_t TF_swap32(uint32_t host_int) {
 #endif
 
 #if TF_TSTRING_LITTLE_ENDIAN
-#define TF_le32toh(x) TF_swap32(x)
-#else  // TF_TSTRING_LITTLE_ENDIAN
 #define TF_le32toh(x) x
+#else  // TF_TSTRING_LITTLE_ENDIAN
+#define TF_le32toh(x) TF_swap32(x)
 #endif  // TF_TSTRING_LITTLE_ENDIAN
 
 static inline size_t TF_align16(size_t i) { return (i + 0xF) & ~0xF; }

tensorflow/core/platform/ctstring_test.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <memory>
 #include <string>
 
+#include "tensorflow/core/platform/ctstring_internal.h"
 #include "tensorflow/core/platform/test.h"
 
 static const char kLongString[] =
@@ -380,3 +381,29 @@ TEST(TF_CTStringTest, ResizeReserve) {
     TF_TString_Dealloc(&s70);
   }
 }
+
+TEST(TF_CTStringTest, OffsetType) {
+  {
+    TF_TString s71;
+
+    TF_TString_Init(&s71);
+    size_t header_length = 24;
+    size_t size = 8;
+    TF_TString_ResizeUninitialized(&s71, header_length + size);
+    uint32_t save_size = s71.u.offset.size;
+    uint32_t save_offset = s71.u.offset.offset;
+    uint32_t save_count = s71.u.offset.count;
+
+    s71.u.offset.size = TF_TString_ToInternalSizeT(size, TF_TSTR_OFFSET);
+    s71.u.offset.offset = header_length;
+    s71.u.offset.count = 0;
+    EXPECT_EQ(size, TF_TString_GetSize(&s71));
+    EXPECT_EQ(TF_TSTR_OFFSET, TF_TString_GetType(&s71));
+
+    // restore state so string can be deallocated
+    s71.u.offset.size = save_size;
+    s71.u.offset.offset = save_offset;
+    s71.u.offset.count = save_count;
+    TF_TString_Dealloc(&s71);
+  }
+}