From 93f428fd1768df147171ed674fee1fc5ab8309ec Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Fri, 30 Jul 2021 21:42:36 -0700 Subject: [PATCH] Fix nullptr deref and heap OOB access in binary cwise ops. PiperOrigin-RevId: 387936777 Change-Id: I608b8074cec36a982cca622b7144cb2c43e6e19f --- tensorflow/core/kernels/cwise_ops_common.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorflow/core/kernels/cwise_ops_common.h b/tensorflow/core/kernels/cwise_ops_common.h index 9adc628421d04..4f2c83322ba00 100644 --- a/tensorflow/core/kernels/cwise_ops_common.h +++ b/tensorflow/core/kernels/cwise_ops_common.h @@ -265,6 +265,11 @@ class SimpleBinaryOp : public OpKernel { void Compute(OpKernelContext* ctx) override { const Tensor& in0 = ctx->input(0); const Tensor& in1 = ctx->input(1); + OP_REQUIRES( + ctx, in0.NumElements() == in1.NumElements(), + errors::InvalidArgument("The two arguments to a cwise op must have " + "same number of elements, got ", + in0.NumElements(), " and ", in1.NumElements())); auto in0_flat = in0.flat(); auto in1_flat = in1.flat(); const Device& eigen_device = ctx->eigen_device();