Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tf.raw_ops.EditDistance vulnerability with negative indices.
Check that indices are non-negative. Fix several identical code sites.
Clean up grammar in error message.

PiperOrigin-RevId: 445442017
  • Loading branch information
poulsbo authored and tensorflower-gardener committed Apr 29, 2022
1 parent a969f2b commit 30721cf
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 10 deletions.
20 changes: 10 additions & 10 deletions tensorflow/core/kernels/edit_distance_op.cc
Expand Up @@ -203,9 +203,9 @@ class EditDistanceOp : public OpKernel {
auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
output_strides.begin(), int64_t{0});
OP_REQUIRES(
ctx, loc < output_elements,
ctx, 0 <= loc && loc < output_elements,
errors::Internal("Got an inner product ", loc,
" which would require in writing to outside of "
" which would require writing to outside of "
"the buffer for the output tensor (max elements ",
output_elements, ")"));
output_t(loc) =
Expand All @@ -218,9 +218,9 @@ class EditDistanceOp : public OpKernel {
auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(),
output_strides.begin(), int64_t{0});
OP_REQUIRES(
ctx, loc < output_elements,
ctx, 0 <= loc && loc < output_elements,
errors::Internal("Got an inner product ", loc,
" which would require in writing to outside of "
" which would require writing to outside of "
"the buffer for the output tensor (max elements ",
output_elements, ")"));
output_t(loc) = hypothesis_seq.size();
Expand All @@ -232,9 +232,9 @@ class EditDistanceOp : public OpKernel {
auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
output_strides.begin(), int64_t{0});
OP_REQUIRES(
ctx, loc < output_elements,
ctx, 0 <= loc && loc < output_elements,
errors::Internal("Got an inner product ", loc,
" which would require in writing to outside of "
" which would require writing to outside of "
"the buffer for the output tensor (max elements ",
output_elements, ")"));
output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();
Expand All @@ -248,9 +248,9 @@ class EditDistanceOp : public OpKernel {
auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(),
output_strides.begin(), int64_t{0});
OP_REQUIRES(
ctx, loc < output_elements,
ctx, 0 <= loc && loc < output_elements,
errors::Internal("Got an inner product ", loc,
" which would require in writing to outside of the "
" which would require writing to outside of the "
"buffer for the output tensor (max elements ",
output_elements, ")"));
output_t(loc) = hypothesis_seq.size();
Expand All @@ -266,9 +266,9 @@ class EditDistanceOp : public OpKernel {
auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
output_strides.begin(), int64_t{0});
OP_REQUIRES(
ctx, loc < output_elements,
ctx, 0 <= loc && loc < output_elements,
errors::Internal("Got an inner product ", loc,
" which would require in writing to outside of the "
" which would require writing to outside of the "
"buffer for the output tensor (max elements ",
output_elements, ")"));
output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();
Expand Down
18 changes: 18 additions & 0 deletions tensorflow/python/kernel_tests/array_ops/edit_distance_op_test.py
Expand Up @@ -207,6 +207,24 @@ def testEditDistanceZeroLengthHypothesisAndTruth(self):
normalize=True,
expected_output=expected_output)

def testEditDistanceBadIndices(self):
hypothesis_indices = np.full((3, 3), -1250999896764, dtype=np.int64)
hypothesis_values = np.empty(3, dtype=np.int64)
hypothesis_shape = np.empty(3, dtype=np.int64)
truth_indices = np.full((3, 3), -1250999896764, dtype=np.int64)
truth_values = np.full([3], 2, dtype=np.int64)
truth_shape = np.full([3], 2, dtype=np.int64)
expected_output = [] # dummy; ignored

self._testEditDistance(
hypothesis=(hypothesis_indices, hypothesis_values, hypothesis_shape),
truth=(truth_indices, truth_values, truth_shape),
normalize=False,
expected_output=expected_output,
expected_err_re=(r"inner product -\d+ which would require writing "
"to outside of the buffer for the output tensor")
)


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

0 comments on commit 30721cf

Please sign in to comment.