Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix security vulnerability in EditDistance op shape function.
PiperOrigin-RevId: 504367470
  • Loading branch information
ishark authored and tensorflower-gardener committed Jan 24, 2023
1 parent 7533da4 commit 08b8e18
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 2 deletions.
14 changes: 13 additions & 1 deletion tensorflow/core/ops/array_ops.cc
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/mirror_pad_mode.h"
#include "tensorflow/core/util/padding.h"
Expand Down Expand Up @@ -1072,13 +1073,24 @@ REGISTER_OP("EditDistance")
// or else the output shape is unknown.
return shape_inference::UnknownShape(c);
}

if (hypothesis_shape_t->NumElements() != truth_shape_t->NumElements()) {
return errors::InvalidArgument(
"Num elements of hypothesis_shape does not match truth_shape: ",
hypothesis_shape_t->NumElements(), " vs. ",
truth_shape_t->NumElements());
}
if (hypothesis_shape_t->NumElements() < 2) {
return errors::InvalidArgument(
"Input Hypothesis SparseTensors must have rank at least 2, but "
"hypothesis_shape rank is: ",
hypothesis_shape_t->NumElements());
}
if (truth_shape_t->NumElements() < 2) {
return errors::InvalidArgument(
"Input Truth SparseTensors must have rank at least 2, but "
"truth_shape rank is: ",
truth_shape_t->NumElements());
}

auto h_values = hypothesis_shape_t->flat<int64_t>();
auto t_values = truth_shape_t->flat<int64_t>();
Expand Down
Expand Up @@ -15,8 +15,9 @@
"""Tests for tensorflow.kernels.edit_distance_op."""

import numpy as np

from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import array_ops
Expand Down Expand Up @@ -225,6 +226,66 @@ def testEditDistanceBadIndices(self):
"to outside of the buffer for the output tensor|"
r"Dimension -\d+ must be >= 0"))

def testEmptyShapeWithEditDistanceRaisesError(self):
para = {
"hypothesis_indices": [[]],
"hypothesis_values": ["tmp/"],
"hypothesis_shape": [],
"truth_indices": [[]],
"truth_values": [""],
"truth_shape": [],
"normalize": False,
}

# Check edit distance raw op with empty shape in eager mode.
with self.assertRaisesRegex(
(errors.InvalidArgumentError, ValueError),
(
r"Input Hypothesis SparseTensors must have rank at least 2, but"
" hypothesis_shape rank is: 0|Input SparseTensors must have rank "
"at least 2, but truth_shape rank is: 0"
),
):
array_ops.gen_array_ops.EditDistance(**para)

# Check raw op with tf.function
@def_function.function
def TestFunction():
"""Wrapper function for edit distance call."""
array_ops.gen_array_ops.EditDistance(**para)

with self.assertRaisesRegex(
ValueError,
(
"Input Hypothesis SparseTensors must have rank at least 2, but"
" hypothesis_shape rank is: 0"
),
):
TestFunction()

# Check with python wrapper API
hypothesis_indices = [[]]
hypothesis_values = [0]
hypothesis_shape = []
truth_indices = [[]]
truth_values = [1]
truth_shape = []
expected_output = [] # dummy ignored

with self.assertRaisesRegex(
ValueError,
(
"Input Hypothesis SparseTensors must have rank at least 2, but"
" hypothesis_shape rank is: 0"
),
):
self._testEditDistance(
hypothesis=(hypothesis_indices, hypothesis_values, hypothesis_shape),
truth=(truth_indices, truth_values, truth_shape),
normalize=False,
expected_output=expected_output,
)


if __name__ == "__main__":
test.main()

0 comments on commit 08b8e18

Please sign in to comment.