Skip to content

Commit 8ee24e7

Browse files
[tflite] Ensure MatchingDim does not allow buffer overflow.
We check in `MatchingDim` that both arguments have the same dimensionality, however that is a `DCHECK` only enabled if building in debug mode. Hence, it could be possible to cause buffer overflows by passing in a tensor with larger dimensions as the second argument. To fix, we now make `MatchingDim` return the minimum of the two sizes. A much better fix would be to return a status object but that requires refactoring a large part of the codebase for minor benefits. PiperOrigin-RevId: 332526127 Change-Id: If627d0d2c80a685217b6e0d1e64b0872dbf1c5e4
1 parent 0b5662b commit 8ee24e7

File tree

1 file changed

+1
-1
lines changed
  • tensorflow/lite/kernels/internal

1 file changed

+1
-1
lines changed

Diff for: tensorflow/lite/kernels/internal/types.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ int MatchingArraySize(const ArrayType1& array1, int index1,
438438
inline int MatchingDim(const RuntimeShape& shape1, int index1,
439439
const RuntimeShape& shape2, int index2) {
440440
TFLITE_DCHECK_EQ(shape1.Dims(index1), shape2.Dims(index2));
441-
return shape1.Dims(index1);
441+
return std::min(shape1.Dims(index1), shape2.Dims(index2));
442442
}
443443

444444
template <typename... Args>

0 commit comments

Comments
 (0)