Skip to content

Commit 6f4d3e8

Browse files
Prevent crash due to integer overflow followed by allocating negative sized array.
PiperOrigin-RevId: 414891322 Change-Id: I5df390e0dc1d9f115209293708950cdf9306931c
1 parent adbbabd commit 6f4d3e8

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

Diff for: tensorflow/core/kernels/count_ops.cc

+9
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16+
#include <limits>
17+
1618
#include "absl/container/flat_hash_map.h"
1719
#include "tensorflow/core/framework/op_kernel.h"
1820
#include "tensorflow/core/framework/op_requires.h"
@@ -23,6 +25,9 @@ limitations under the License.
2325

2426
namespace tensorflow {
2527

28+
// Don't allocate too large `BatchedMap<T>` objects
29+
static int kMaxBatches = std::numeric_limits<int>::max();
30+
2631
template <class T>
2732
using BatchedMap = std::vector<absl::flat_hash_map<int64_t, T>>;
2833

@@ -235,6 +240,10 @@ class SparseCount : public OpKernel {
235240

236241
bool is_1d = shape.NumElements() == 1;
237242
int num_batches = is_1d ? 1 : shape_vector(0);
243+
OP_REQUIRES(
244+
context, 0 < num_batches && num_batches < kMaxBatches,
245+
errors::InvalidArgument("Cannot allocate ", num_batches,
246+
" batches, is the dense shape too wide?"));
238247

239248
const auto values_values = values.flat<T>();
240249
const auto weight_values = weights.flat<W>();

0 commit comments

Comments
 (0)