Merge pull request #55876 from tensorflow/r2.6-13d38a07ce9

r2.6 cherry-pick: 13d38a0 "Fix tf.raw_ops.UnsortedSegmentJoin vulnerability with invalid num_segments."
tensorflow · May 4, 2022 · d379e5c · d379e5c
2 parents f33b5ef + 0a7bd11
commit d379e5c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/unsorted_segment_join_op.cc b/tensorflow/core/kernels/unsorted_segment_join_op.cc
@@ -92,6 +92,9 @@ class UnsortedSegmentJoinOp : public OpKernel {
     const Tensor& num_segments_tensor = context->input(2);
     OP_REQUIRES(context, num_segments_tensor.NumElements() != 0,
                 errors::InvalidArgument("Number of segments cannot be empty."));
+    OP_REQUIRES(context,
+                TensorShapeUtils::IsScalar(num_segments_tensor.shape()),
+                errors::InvalidArgument("Number of segments must be a scalar"));
     auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();
 
     OP_REQUIRES(context, segment_dims != 0,