From 8f97b7caffab7757782be2f19b6cd3e5c154432a Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Sat, 19 Sep 2020 18:18:51 -0700 Subject: [PATCH] [tflite] Ensure `ResolveAxis` properly handles negative inputs. In Python, a list `l` of length `n` allows indexing with negative indices, `l[i]`. The only constraint is that `n + i` becomes positive. Code in `ResolveAxis` assumes the constraints and only checks it using a `DCHECK`. But the macro is a no-op in non-debug builds and that can result in reading from negative offsets (buffer underflows). PiperOrigin-RevId: 332530683 Change-Id: I464e073fee618054ae3719a3679739007bb3f3bc --- tensorflow/lite/kernels/internal/reference/reference_ops.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h index 5f2e8331f59..7dda94dc620 100644 --- a/tensorflow/lite/kernels/internal/reference/reference_ops.h +++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h @@ -2658,6 +2658,9 @@ inline bool ResolveAxis(const int num_dims, const int* axis, // Handle negative index. int current = axis[idx] < 0 ? (axis[idx] + num_dims) : axis[idx]; TFLITE_DCHECK(current >= 0 && current < num_dims); + if (current < 0 || current >= num_dims) { + return false; + } bool is_dup = false; for (int j = 0; j < *out_num_axis; ++j) { if (out_axis[j] == current) { -- GitLab