Skip to content

Commit 8faa6ea

Browse files
poulsbotensorflower-gardener
authored andcommitted
Fix tf.raw_ops.ImageProjectiveTransformV2 vulnerability with large output_shape.
Note: This fix will have to be cherry picked in r2.10, r2.9, and r2.8. PiperOrigin-RevId: 479125772
1 parent f264ff5 commit 8faa6ea

File tree

2 files changed

+28
-4
lines changed

2 files changed

+28
-4
lines changed

Diff for: tensorflow/core/kernels/image/image_ops.cc

+5-4
Original file line numberDiff line numberDiff line change
@@ -96,11 +96,12 @@ void DoImageProjectiveTransformOp(OpKernelContext* ctx,
9696
}
9797

9898
Tensor* output_t;
99+
TensorShape output_shape;
99100
OP_REQUIRES_OK(
100-
ctx, ctx->allocate_output(0,
101-
TensorShape({images_t.dim_size(0), out_height,
102-
out_width, images_t.dim_size(3)}),
103-
&output_t));
101+
ctx, TensorShape::BuildTensorShape({images_t.dim_size(0), out_height,
102+
out_width, images_t.dim_size(3)},
103+
&output_shape));
104+
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output_t));
104105
auto output = output_t->tensor<T, 4>();
105106
auto images = images_t.tensor<T, 4>();
106107
auto transform = transform_t.matrix<float>();

Diff for: tensorflow/python/ops/image_ops_test.py

+23
Original file line numberDiff line numberDiff line change
@@ -2335,6 +2335,29 @@ def testInvalidInput(self):
23352335
self.evaluate(v)
23362336

23372337

2338+
class ImageProjectiveTransformV2(test_util.TensorFlowTestCase):
2339+
2340+
def testShapeTooLarge(self):
2341+
interpolation = "BILINEAR"
2342+
fill_mode = "REFLECT"
2343+
images = constant_op.constant(
2344+
0.184634328, shape=[2, 5, 8, 3], dtype=dtypes.float32)
2345+
transforms = constant_op.constant(
2346+
0.378575385, shape=[2, 8], dtype=dtypes.float32)
2347+
output_shape = constant_op.constant([1879048192, 1879048192],
2348+
shape=[2],
2349+
dtype=dtypes.int32)
2350+
with self.assertRaisesRegex(errors.InvalidArgumentError,
2351+
r"Encountered overflow when multiplying"):
2352+
self.evaluate(
2353+
gen_image_ops.ImageProjectiveTransformV2(
2354+
images=images,
2355+
transforms=transforms,
2356+
output_shape=output_shape,
2357+
interpolation=interpolation,
2358+
fill_mode=fill_mode))
2359+
2360+
23382361
class InternalPadToBoundingBoxTest(test_util.TensorFlowTestCase,
23392362
parameterized.TestCase):
23402363

0 commit comments

Comments
 (0)