[pull] master from tensorflow:master

tensorflow/c/c_api_experimental.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/c/c_api_experimental.h"
 
 #include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/substitute.h"
 #include "absl/synchronization/notification.h"
 #include "tensorflow/c/c_api.h"
@@ -55,7 +56,6 @@ using tensorflow::FunctionDef;
 using tensorflow::Node;
 using tensorflow::NodeBuilder;
 using tensorflow::Status;
-using tensorflow::errors::InvalidArgument;
 
 namespace {
 typedef std::unique_ptr<TF_Function, decltype(&TF_DeleteFunction)>
@@ -500,6 +500,14 @@ TFE_TensorHandle* TFE_NewTensorHandleFromScalar(TF_DataType data_type,
     return nullptr;
   }
 
+  size_t expected_len = static_cast<size_t>(tensorflow::DataTypeSize(dtype));
+  if (len > expected_len) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("len (", len, ") is larger than expected scalar size for ",
+                     dtype, " (", expected_len, " bytes)"));
+    return nullptr;
+  }
+
   tensorflow::Tensor tensor(dtype, tensorflow::TensorShape({}));
   std::memcpy(tensorflow::TensorCApi::Buffer(tensor)->data(), data, len);
 
@@ -728,7 +736,7 @@ void TF_ImportGraphDefOptionsSetValidateColocationConstraints(
 TF_Library* TF_LoadPluggableDeviceLibrary(const char* library_filename,
                                           TF_Status* status) {
 #if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD)
-  status->status = tensorflow::errors::Unimplemented(
+  status->status = absl::UnimplementedError(
       "PluggableDevice plugin functionality is not supported on mobile");
   return nullptr;
 #else

tensorflow/c/eager/BUILD

@@ -926,6 +926,7 @@ tf_cuda_library(
         "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_client",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@xla//xla/tsl/c:tsl_status_internal",

tensorflow/c/eager/c_api_unified_experimental_graph.cc

@@ -17,9 +17,11 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_operation.h"
 #include "tensorflow/c/eager/c_api_internal.h"
 #include "tensorflow/c/eager/c_api_unified_experimental.h"
 #include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
@@ -33,7 +35,6 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
-#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/strcat.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -178,6 +179,11 @@ class GraphOperation : public TracingOperation {
     TF_RETURN_IF_ERROR(StatusFromTF_Status(s));
     TF_DeleteStatus(s);
     *num_retvals = TF_OperationNumOutputs(operation);
+    if (*num_retvals > retvals.size()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "retvals span capacity (", retvals.size(),
+          ") is smaller than operation output count (", *num_retvals, ")"));
+    }
     for (int i = 0; i < *num_retvals; ++i) {
       retvals[i] = new GraphTensor({operation, i}, g_);
     }

tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc

@@ -159,6 +159,11 @@ static int64_t LoadBufferFromGCS(const std::string& path, size_t offset,
   TF_SetStatus(status, TF_OK, "");
   VLOG(1) << absl::StrFormat("Successful read of %s @ %u of size: %u", path,
                              offset, read);
+  if (read < 0) {
+    TF_SetStatus(status, TF_INTERNAL, "Invalid negative content-length");
+    return -1;
+  }
+  read = std::min<int64_t>(read, buffer_size);
   stream.read(buffer, read);
   read = stream.gcount();
   if (read < buffer_size) {

tensorflow/c/tf_tensor.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/status.h"
 #include "tsl/platform/ctstring.h"
 #include "tsl/platform/ctstring_internal.h"

tensorflow/core/data/BUILD

@@ -101,6 +101,7 @@ cc_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
     ],
 )
 

tensorflow/core/data/compression_utils.cc

@@ -22,14 +22,18 @@ limitations under the License.
 
 #include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/common_runtime/dma_helper.h"
 #include "tensorflow/core/framework/dataset.pb.h"
+#include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/framework/variant_op_registry.h"
-#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/snappy.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
@@ -158,8 +162,16 @@ absl::Status UncompressElement(const CompressedElement& compressed,
     if (metadata.dtype() == DT_STRING) {
       ++num_string_tensors;
       num_string_tensor_strings += metadata.uncompressed_bytes_size();
-    } else if (!DataTypeCanUseMemcpy(metadata.dtype())) {
-      total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
+    } else {
+      TensorShape shape(metadata.tensor_shape());
+      int64_t num_elements = shape.num_elements();
+      if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
+        return absl::InvalidArgumentError(
+            "Missing uncompressed_bytes metadata for non-empty tensor");
+      }
+      if (!DataTypeCanUseMemcpy(metadata.dtype()) && num_elements > 0) {
+        total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
+      }
     }
   }
 
@@ -175,22 +187,39 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   char* nonmemcpyable_pos = nonmemcpyable.mdata();
   for (const auto& metadata : compressed.component_metadata()) {
     if (DataTypeCanUseMemcpy(metadata.dtype())) {
+      TensorShape shape(metadata.tensor_shape());
+      int64_t num_elements = shape.num_elements();
       out->emplace_back(metadata.dtype(), metadata.tensor_shape());
       TensorBuffer* buffer = DMAHelper::buffer(&out->back());
-      if (buffer) {
+      if (buffer && num_elements > 0) {
+        if (metadata.uncompressed_bytes(0) > buffer->size()) {
+          return absl::InvalidArgumentError(absl::StrCat(
+              "uncompressed_bytes (", metadata.uncompressed_bytes(0),
+              ") exceeds allocated buffer size (", buffer->size(), ")"));
+        }
         iov.Add(buffer->data(), metadata.uncompressed_bytes(0));
       }
     } else if (metadata.dtype() == DT_STRING) {
       out->emplace_back(metadata.dtype(), metadata.tensor_shape());
       const auto& flats = out->back().unaligned_flat<tstring>();
+      if (metadata.uncompressed_bytes_size() > flats.size()) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "uncompressed_bytes count (", metadata.uncompressed_bytes_size(),
+            ") exceeds allocated string tensor elements count (", flats.size(),
+            ")"));
+      }
       for (int i = 0; i < metadata.uncompressed_bytes_size(); ++i) {
         flats.data()[i].resize(metadata.uncompressed_bytes(i));
         iov.Add(flats.data()[i].mdata(), metadata.uncompressed_bytes(i));
       }
     } else {
+      TensorShape shape(metadata.tensor_shape());
+      int64_t num_elements = shape.num_elements();
       out->emplace_back();
-      iov.Add(nonmemcpyable_pos, metadata.uncompressed_bytes(0));
-      nonmemcpyable_pos += metadata.uncompressed_bytes(0);
+      if (num_elements > 0) {
+        iov.Add(nonmemcpyable_pos, metadata.uncompressed_bytes(0));
+        nonmemcpyable_pos += metadata.uncompressed_bytes(0);
+      }
     }
   }
 

tensorflow/core/framework/BUILD

@@ -935,6 +935,7 @@ cc_library(
         "//tensorflow/core/platform:macros",
         "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
     ],
 )
 

tensorflow/core/framework/shape_inference.h

@@ -18,9 +18,9 @@ limitations under the License.
 #include <vector>
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/node_def_util.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/macros.h"
 
@@ -296,6 +296,7 @@ class InferenceContext {
   // This requires idx to be in the [0, num_inputs) range. If the merge is
   // successful, return true. Return false otherwise.
   bool MergeInput(int idx, ShapeHandle shape) {
+    if (idx < 0 || idx >= num_inputs()) return false;
     ShapeHandle new_shape;
     if (!Merge(inputs_[idx], shape, &new_shape).ok()) return false;
     inputs_[idx] = new_shape;

tensorflow/core/kernels/BUILD

@@ -943,6 +943,8 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
     ],
 )
@@ -1474,6 +1476,8 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/util:ragged_to_dense_util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
     ],
 )
 
@@ -2820,7 +2824,11 @@ tf_kernel_library(
         "gpu_device_array_gpu.h",
     ],
     prefix = "dynamic_stitch_op",
-    deps = DYNAMIC_DEPS + [":loose_headers"],
+    deps = DYNAMIC_DEPS + [
+        ":loose_headers",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
 )
 
 cc_library(
@@ -4572,7 +4580,7 @@ tf_kernel_library(
 tf_kernel_library(
     name = "nth_element_op",
     prefix = "nth_element_op",
-    deps = NN_DEPS,
+    deps = NN_DEPS + ["@com_google_absl//absl/status"],
 )
 
 tf_kernel_library(
@@ -4592,6 +4600,7 @@ tf_kernel_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
     ],
@@ -4605,6 +4614,8 @@ tf_kernel_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
     ],
 )
@@ -6022,6 +6033,7 @@ tf_kernel_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
         "@icu//:common",
     ],

tensorflow/core/kernels/bincount_op.cc

@@ -15,16 +15,17 @@ limitations under the License.
 
 // See docs in ../ops/math_ops.cc.
 
-#include <atomic>
-
-#include "tensorflow/core/platform/errors.h"
 #define EIGEN_USE_THREADS
 
+#include "tensorflow/core/kernels/bincount_op.h"
+
+#include <atomic>
+
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/kernels/bincount_op.h"
 #include "tensorflow/core/kernels/fill_functor.h"
 #include "tensorflow/core/kernels/sparse_utils.h"
 #include "tensorflow/core/lib/core/threadpool.h"
@@ -302,9 +303,9 @@ class DenseBincountOp : public OpKernel {
                     weights.shape().DebugString())));
 
     Tidx size = size_t.scalar<Tidx>()();
-    OP_REQUIRES(
-        ctx, size >= 0,
-        errors::InvalidArgument("size (", size, ") must be non-negative"));
+    OP_REQUIRES(ctx, size >= 0,
+                absl::InvalidArgumentError(
+                    absl::StrCat("size (", size, ") must be non-negative")));
 
     Tensor* out_t;
     functor::SetZeroFunctor<Device, T> fill;
@@ -402,13 +403,18 @@ class SparseBincountOp : public OpKernel {
                 absl::InvalidArgumentError(absl::StrCat(
                     "Shape must be rank 0 but is rank ", size_t.dims())));
     Tidx size = size_t.scalar<Tidx>()();
-    OP_REQUIRES(
-        ctx, size >= 0,
-        errors::InvalidArgument("size (", size, ") must be non-negative"));
+    OP_REQUIRES(ctx, size >= 0,
+                absl::InvalidArgumentError(
+                    absl::StrCat("size (", size, ") must be non-negative")));
     OP_REQUIRES_OK(ctx, sparse_utils::ValidateSparseTensor<int64_t>(
                             indices, values, dense_shape,
                             sparse_utils::IndexValidation::kUnordered));
 
+    OP_REQUIRES(ctx, dense_shape.NumElements() > 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "dense_shape must have at least 1 dimension, got ",
+                    dense_shape.NumElements())));
+
     bool is_1d = dense_shape.NumElements() == 1;
 
     Tensor* out_t;
@@ -493,9 +499,9 @@ class RaggedBincountOp : public OpKernel {
                 absl::InvalidArgumentError(absl::StrCat(
                     "Shape must be rank 0 but is rank ", size_t.dims())));
     Tidx size = size_t.scalar<Tidx>()();
-    OP_REQUIRES(
-        ctx, size >= 0,
-        errors::InvalidArgument("size (", size, ") must be non-negative"));
+    OP_REQUIRES(ctx, size >= 0,
+                absl::InvalidArgumentError(
+                    absl::StrCat("size (", size, ") must be non-negative")));
 
     int num_rows = splits.size() - 1;
     int num_values = values.size();

tensorflow/core/kernels/decode_compressed_op.cc

@@ -19,12 +19,12 @@ limitations under the License.
 #include <climits>
 
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/io/zlib_compression_options.h"
 #include "tensorflow/core/lib/io/zlib_inputstream.h"