Make `MakeShape` and `VectorToShape` robust to negative indices.

This avoids a CHECK failure if invalid input data is passed to some
ops. It also tightens up the Python class to reject negative dimensions.
Change: 110401558

1 files changed, 25 insertions, 11 deletions

diff --git a/tensorflow/core/kernels/conv_grad_ops.cc b/tensorflow/core/kernels/conv_grad_ops.cc
index fffafa2db4..ac0ed0a404 100644
--- a/tensorflow/core/kernels/conv_grad_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_ops.cc
@@ -264,16 +264,18 @@ typedef Eigen::GpuDevice GPUDevice;
           << ", strides = " << strides[1]
 
 namespace {
-TensorShape VectorToShape(const TTypes<int32>::ConstVec& sizes) {
-  TensorShape shape;
-
+Status VectorToShape(const TTypes<int32>::ConstVec& sizes, TensorShape* out) {
   using Index = TTypes<int32>::ConstVec::Index;
   const Index dims = sizes.size();
   for (Index i = 0; i < dims; ++i) {
-    shape.AddDim(sizes(i));
+    if (sizes(i) >= 0) {
+      out->AddDim(sizes(i));
+    } else {
+      return errors::InvalidArgument("Dimension ", sizes(i), " must be >= 0");
+    }
   }
 
-  return shape;
+  return Status::OK();
 }
 }  // namespace
 
@@ -309,7 +311,9 @@ class Conv2DFastBackpropInputOp : public OpKernel {
         errors::InvalidArgument(
             "Conv2DBackpropInput: input_sizes input must be 1-dim, not ",
             input_sizes.dims()));
-    TensorShape input_shape = VectorToShape(input_sizes.vec<int32>());
+    TensorShape input_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(input_sizes.vec<int32>(), &input_shape));
     const TensorShape& filter_shape = filter.shape();
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DBackpropInput");
@@ -359,7 +363,9 @@ class Conv2DCustomBackpropInputOp : public OpKernel {
         errors::InvalidArgument(
             "Conv2DBackpropInput: input_sizes input must be 1-dim, not ",
             input_sizes.dims()));
-    TensorShape input_shape = VectorToShape(input_sizes.vec<int32>());
+    TensorShape input_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(input_sizes.vec<int32>(), &input_shape));
     const TensorShape& filter_shape = filter.shape();
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DBackpropInput");
@@ -566,7 +572,9 @@ class Conv2DFastBackpropFilterOp : public OpKernel {
             "Conv2DBackpropFilter: filter_sizes input must be 1-dim, not ",
             filter_sizes.dims()));
     const TensorShape& input_shape = input.shape();
-    TensorShape filter_shape = VectorToShape(filter_sizes.vec<int32>());
+    TensorShape filter_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(filter_sizes.vec<int32>(), &filter_shape));
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DBackpropFilter");
     Tensor* filter_backprop = nullptr;
@@ -618,7 +626,9 @@ class Conv2DCustomBackpropFilterOp : public OpKernel {
             "not ",
             filter_sizes.dims()));
     const TensorShape& input_shape = input.shape();
-    TensorShape filter_shape = VectorToShape(filter_sizes.vec<int32>());
+    TensorShape filter_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(filter_sizes.vec<int32>(), &filter_shape));
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DCustomBackpropFilter");
     Tensor* filter_backprop;
@@ -790,7 +800,9 @@ class Conv2DSlowBackpropInputOp : public OpKernel {
         errors::InvalidArgument(
             "Conv2DBackpropInput: input_sizes input must be 1-dim, not ",
             input_sizes.dims()));
-    TensorShape input_shape = VectorToShape(input_sizes.vec<int32>());
+    TensorShape input_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(input_sizes.vec<int32>(), &input_shape));
     const TensorShape& filter_shape = filter.shape();
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DBackpropInput");
@@ -1067,7 +1079,9 @@ class Conv2DSlowBackpropFilterOp : public OpKernel {
             "Conv2DBackpropFilter: filter_sizes input must be 1-dim, not ",
             filter_sizes.dims()));
     const TensorShape& input_shape = input.shape();
-    TensorShape filter_shape = VectorToShape(filter_sizes.vec<int32>());
+    TensorShape filter_shape;
+    OP_REQUIRES_OK(context,
+                   VectorToShape(filter_sizes.vec<int32>(), &filter_shape));
 
     EXTRACT_AND_VERIFY_DIMENSIONS("Conv2DBackpropFilter");
     Tensor* filter_backprop = nullptr;