Skip to content

Commit 8dc723f

Browse files
committed
Fix integer overflow for multiframe gifs.
1 parent 4f27d19 commit 8dc723f

File tree

5 files changed

+13
-11
lines changed

5 files changed

+13
-11
lines changed

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -452,12 +452,12 @@ class DecodeImageV2Op : public OpKernel {
452452
// allocation til after dtype conversion is done. `gif`::Decode` supports
453453
// uint8 only.
454454
Tensor* output = nullptr;
455-
int buffer_size = 0;
455+
ptrdiff_t buffer_size = 0;
456456
string error_string;
457457
uint8* buffer = gif::Decode(
458458
input.data(), input.size(),
459459
[&](int num_frames, int width, int height, int channels) -> uint8* {
460-
buffer_size = num_frames * height * width * channels;
460+
buffer_size = ptrdiff_t(num_frames) * height * width * channels;
461461

462462
Status status;
463463
// By the existing API, we support decoding GIF with `decode_jpeg` or

Diff for: tensorflow/core/lib/gif/gif_io.cc

+7-7
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ uint8* Decode(const void* srcdata, int datasize,
105105
uint8* const dstdata =
106106
allocate_output(target_num_frames, width, height, channel);
107107
if (!dstdata) return nullptr;
108-
for (int k = 0; k < target_num_frames; k++) {
108+
for (ptrdiff_t k = 0; k < target_num_frames; k++) {
109109
uint8* this_dst = dstdata + k * width * channel * height;
110110

111111
SavedImage* this_image = &gif_file->SavedImages[k];
@@ -125,10 +125,10 @@ uint8* Decode(const void* srcdata, int datasize,
125125

126126
if (k > 0) {
127127
uint8* last_dst = dstdata + (k - 1) * width * channel * height;
128-
for (int i = 0; i < height; ++i) {
128+
for (ptrdiff_t i = 0; i < height; ++i) {
129129
uint8* p_dst = this_dst + i * width * channel;
130130
uint8* l_dst = last_dst + i * width * channel;
131-
for (int j = 0; j < width; ++j) {
131+
for (ptrdiff_t j = 0; j < width; ++j) {
132132
p_dst[j * channel + 0] = l_dst[j * channel + 0];
133133
p_dst[j * channel + 1] = l_dst[j * channel + 1];
134134
p_dst[j * channel + 2] = l_dst[j * channel + 2];
@@ -141,9 +141,9 @@ uint8* Decode(const void* srcdata, int datasize,
141141
// If the first frame does not fill the entire canvas then fill the
142142
// unoccupied canvas with zeros (black).
143143
if (k == 0) {
144-
for (int i = 0; i < height; ++i) {
144+
for (ptrdiff_t i = 0; i < height; ++i) {
145145
uint8* p_dst = this_dst + i * width * channel;
146-
for (int j = 0; j < width; ++j) {
146+
for (ptrdiff_t j = 0; j < width; ++j) {
147147
p_dst[j * channel + 0] = 0;
148148
p_dst[j * channel + 1] = 0;
149149
p_dst[j * channel + 2] = 0;
@@ -165,9 +165,9 @@ uint8* Decode(const void* srcdata, int datasize,
165165
return nullptr;
166166
}
167167

168-
for (int i = imgTop; i < imgBottom; ++i) {
168+
for (ptrdiff_t i = imgTop; i < imgBottom; ++i) {
169169
uint8* p_dst = this_dst + i * width * channel;
170-
for (int j = imgLeft; j < imgRight; ++j) {
170+
for (ptrdiff_t j = imgLeft; j < imgRight; ++j) {
171171
GifByteType color_index =
172172
this_image->RasterBits[(i - img_desc->Top) * (img_desc->Width) +
173173
(j - img_desc->Left)];

Diff for: tensorflow/core/lib/gif/gif_io_test.cc

+3-2
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ void TestDecodeGif(Env* env, DecodeGifTestCase testcase) {
5252
w = width;
5353
h = height;
5454
c = channels;
55-
return new uint8[frame_cnt * height * width * channels];
55+
return new uint8[ptrdiff_t(frame_cnt) * height * width * channels];
5656
},
5757
&error_string));
5858
ASSERT_NE(imgdata, nullptr);
@@ -72,7 +72,8 @@ TEST(GifTest, Gif) {
7272
{testdata_path + "optimized.gif", 12, 20, 40, 3},
7373
{testdata_path + "red_black.gif", 1, 16, 16, 3},
7474
{testdata_path + "scan.gif", 12, 20, 40, 3},
75-
{testdata_path + "squares.gif", 2, 16, 16, 3}});
75+
{testdata_path + "squares.gif", 2, 16, 16, 3},
76+
{testdata_path + "3g_multiframe.gif", 519, 1920, 1080, 3}});
7677

7778
for (const auto& tc : testcases) {
7879
TestDecodeGif(env, tc);

Diff for: tensorflow/core/lib/gif/testdata/3g_multiframe.gif

22.2 KB
Loading

Diff for: tensorflow/core/lib/gif/testdata/BUILD

+1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ filegroup(
1616
"scan.gif",
1717
"red_black.gif",
1818
"squares.gif",
19+
"3g_multiframe.gif",
1920
"pendulum_sm.gif",
2021
# Add groundtruth frames for `pendulum_sm.gif`.
2122
# PNG format because it's lossless.

0 commit comments

Comments
 (0)