Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix an invalid address vulnerability in tf.raw_ops.RaggedBincount.
PiperOrigin-RevId: 368293153
Change-Id: I4b4e493d3fd05e7dc55a55de3a041a80a4f275c3
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 13, 2021
1 parent 8b677d7 commit eebb96c
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions tensorflow/core/kernels/bincount_op.cc
Expand Up @@ -420,6 +420,15 @@ class RaggedBincountOp : public OpKernel {
int num_values = values.size();
int batch_idx = 0;

OP_REQUIRES(ctx, splits(0) == 0,
errors::InvalidArgument("Splits must start with 0, not with ",
splits(0)));

OP_REQUIRES(ctx, splits(num_rows) == num_values,
errors::InvalidArgument(
"Splits must end with the number of values, got ",
splits(num_rows), " instead of ", num_values));

Tensor* out_t;
OP_REQUIRES_OK(
ctx, ctx->allocate_output(0, TensorShape({num_rows, size}), &out_t));
Expand Down

1 comment on commit eebb96c

@Rayyan335
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tensorflow/core/kernels/bincount_op.cc

Please sign in to comment.