Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tf.raw_ops.TensorListConcat vulnerability (segfault).
Check that the element_shape input is valid.
Add graph/eager unit tests.

Note: This fix will have to be cherry picked in r2.10, r2.9, and r2.8.
PiperOrigin-RevId: 477544091
  • Loading branch information
poulsbo authored and tensorflower-gardener committed Sep 28, 2022
1 parent 5f0eb84 commit fc33f3d
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 2 deletions.
7 changes: 5 additions & 2 deletions tensorflow/core/kernels/list_kernels.h
Expand Up @@ -395,8 +395,11 @@ class TensorListConcat : public OpKernel {
void Compute(OpKernelContext* c) override {
PartialTensorShape element_shape_except_first_dim;
if (!element_shape_.unknown_rank()) {
element_shape_except_first_dim = PartialTensorShape(
gtl::ArraySlice<int64_t>(element_shape_.dim_sizes()).subspan(1));
auto dim_sizes = element_shape_.dim_sizes();
OP_REQUIRES(c, !dim_sizes.empty(),
errors::InvalidArgument("element_shape must not be empty"));
element_shape_except_first_dim =
PartialTensorShape(gtl::ArraySlice<int64_t>(dim_sizes).subspan(1));
}
// Check that the input Variant tensor is indeed a TensorList and has the
// correct element type.
Expand Down
Expand Up @@ -1514,6 +1514,15 @@ def testConcatWithUninitializedTensorsFailsIfNoInputLengths(self):
t = list_ops.tensor_list_concat(l, element_dtype=dtypes.float32)
self.evaluate(t)

@test_util.run_in_graph_and_eager_modes
def testConcatWithInvalidElementShape(self):
l = list_ops.tensor_list_reserve(
element_dtype=dtypes.float32, element_shape=[], num_elements=0)
with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
r"element_shape must not be empty"):
self.evaluate(gen_list_ops.tensor_list_concat(
input_handle=l, element_dtype=dtypes.float32, element_shape=[]))

def testEmptyTensorListInvalidShape(self):
with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
r"Shape must be at most rank 1 but is rank 2"):
Expand Down

0 comments on commit fc33f3d

Please sign in to comment.