From bc9c546ce7015c57c2f15c168b3d9201de679a1d Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Fri, 30 Jul 2021 21:37:59 -0700
Subject: [PATCH] Prevent heap oob access in `resource_variable_ops.cc`

PiperOrigin-RevId: 387936433
Change-Id: I9e71ddaa8dbd51ec6afbf163a6b3b591f193b4f6
---
 tensorflow/core/kernels/resource_variable_ops.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index b9c883c7..1d7e7c3c 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -643,6 +643,11 @@ class ResourceGatherOp : public OpKernel {
     OP_REQUIRES(
         c, TensorShapeUtils::IsVectorOrHigher(params.shape()),
         errors::InvalidArgument("params must be at least 1 dimensional"));
+    OP_REQUIRES(
+        c, params.shape().dims() >= batch_dims_,
+	 errors::InvalidArgument("params must have at least ", batch_dims_,
+                                " (batch_dims) dimensions but it has shape ",
+				 params.shape().DebugString()));
 
     // Check that we have enough index space
     const int64 N = indices.NumElements();
-- 
2.27.0