Fix nullptr deref and heap OOB access in binary cwise ops.

PiperOrigin-RevId: 387936777 Change-Id: I608b8074cec36a982cca622b7144cb2c43e6e19f
tensorflow · Jul 31, 2021 · 93f428f · 93f428f
1 parent bc9c546
commit 93f428f
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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<Tin>();
     auto in1_flat = in1.flat<Tin>();
     const Device& eigen_device = ctx->eigen_device<Device>();