Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix security vulnerability with LRNGradOp
PiperOrigin-RevId: 460738938
  • Loading branch information
sagunb authored and tensorflower-gardener committed Jul 13, 2022
1 parent 642f837 commit bd90b3e
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 1 deletion.
3 changes: 2 additions & 1 deletion tensorflow/core/kernels/lrn_op.cc
Expand Up @@ -668,7 +668,8 @@ class LRNGradOp : public OpKernel {
in_image.dim_size(0) == batch && in_image.dim_size(1) == rows &&
in_image.dim_size(2) == cols && in_image.dim_size(3) == depth &&
out_image.dim_size(0) == batch && out_image.dim_size(1) == rows &&
out_image.dim_size(2) == cols && out_image.dim_size(3) == depth,
out_image.dim_size(2) == cols && out_image.dim_size(3) == depth &&
out_image.dims() == 4,
errors::InvalidArgument(
"input_grads, input_image, and out_image should have the same "
"shape"));
Expand Down
37 changes: 37 additions & 0 deletions tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py
Expand Up @@ -20,11 +20,13 @@

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
import tensorflow.python.ops.nn_grad # pylint: disable=unused-import
from tensorflow.python.platform import test

Expand Down Expand Up @@ -111,6 +113,41 @@ def testGradientsZeroInput(self):
self.assertAllClose(r, expected)
self.assertShapeEqual(expected, grad)

@test_util.run_in_graph_and_eager_modes
def testIncompatibleInputAndOutputImageShapes(self):
depth_radius = 1
bias = 1.59018219
alpha = 0.117728651
beta = 0.404427052
input_grads = random_ops.random_uniform(
shape=[4, 4, 4, 4],
minval=-10000,
maxval=10000,
dtype=dtypes.float32,
seed=-2033)
input_image = random_ops.random_uniform(
shape=[4, 4, 4, 4],
minval=-10000,
maxval=10000,
dtype=dtypes.float32,
seed=-2033)
invalid_output_image = random_ops.random_uniform(
shape=[4, 4, 4, 4, 4, 4],
minval=-10000,
maxval=10000,
dtype=dtypes.float32,
seed=-2033)
with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
self.evaluate(
nn.lrn_grad(
input_grads=input_grads,
input_image=input_image,
output_image=invalid_output_image,
depth_radius=depth_radius,
bias=bias,
alpha=alpha,
beta=beta))

def _RunAndVerifyGradients(self, dtype):
with self.cached_session():
# random shape
Expand Down

0 comments on commit bd90b3e

Please sign in to comment.