Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix integer overflow for multiframe gifs.
  • Loading branch information
API92 committed Dec 27, 2022
1 parent 4f27d19 commit 8dc723f
Show file tree
Hide file tree
Showing 5 changed files with 13 additions and 11 deletions.
4 changes: 2 additions & 2 deletions tensorflow/core/kernels/image/decode_image_op.cc
Expand Up @@ -452,12 +452,12 @@ class DecodeImageV2Op : public OpKernel {
// allocation til after dtype conversion is done. `gif`::Decode` supports
// uint8 only.
Tensor* output = nullptr;
int buffer_size = 0;
ptrdiff_t buffer_size = 0;
string error_string;
uint8* buffer = gif::Decode(
input.data(), input.size(),
[&](int num_frames, int width, int height, int channels) -> uint8* {
buffer_size = num_frames * height * width * channels;
buffer_size = ptrdiff_t(num_frames) * height * width * channels;

Status status;
// By the existing API, we support decoding GIF with `decode_jpeg` or
Expand Down
14 changes: 7 additions & 7 deletions tensorflow/core/lib/gif/gif_io.cc
Expand Up @@ -105,7 +105,7 @@ uint8* Decode(const void* srcdata, int datasize,
uint8* const dstdata =
allocate_output(target_num_frames, width, height, channel);
if (!dstdata) return nullptr;
for (int k = 0; k < target_num_frames; k++) {
for (ptrdiff_t k = 0; k < target_num_frames; k++) {
uint8* this_dst = dstdata + k * width * channel * height;

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

if (k > 0) {
uint8* last_dst = dstdata + (k - 1) * width * channel * height;
for (int i = 0; i < height; ++i) {
for (ptrdiff_t i = 0; i < height; ++i) {
uint8* p_dst = this_dst + i * width * channel;
uint8* l_dst = last_dst + i * width * channel;
for (int j = 0; j < width; ++j) {
for (ptrdiff_t j = 0; j < width; ++j) {
p_dst[j * channel + 0] = l_dst[j * channel + 0];
p_dst[j * channel + 1] = l_dst[j * channel + 1];
p_dst[j * channel + 2] = l_dst[j * channel + 2];
Expand All @@ -141,9 +141,9 @@ uint8* Decode(const void* srcdata, int datasize,
// If the first frame does not fill the entire canvas then fill the
// unoccupied canvas with zeros (black).
if (k == 0) {
for (int i = 0; i < height; ++i) {
for (ptrdiff_t i = 0; i < height; ++i) {
uint8* p_dst = this_dst + i * width * channel;
for (int j = 0; j < width; ++j) {
for (ptrdiff_t j = 0; j < width; ++j) {
p_dst[j * channel + 0] = 0;
p_dst[j * channel + 1] = 0;
p_dst[j * channel + 2] = 0;
Expand All @@ -165,9 +165,9 @@ uint8* Decode(const void* srcdata, int datasize,
return nullptr;
}

for (int i = imgTop; i < imgBottom; ++i) {
for (ptrdiff_t i = imgTop; i < imgBottom; ++i) {
uint8* p_dst = this_dst + i * width * channel;
for (int j = imgLeft; j < imgRight; ++j) {
for (ptrdiff_t j = imgLeft; j < imgRight; ++j) {
GifByteType color_index =
this_image->RasterBits[(i - img_desc->Top) * (img_desc->Width) +
(j - img_desc->Left)];
Expand Down
5 changes: 3 additions & 2 deletions tensorflow/core/lib/gif/gif_io_test.cc
Expand Up @@ -52,7 +52,7 @@ void TestDecodeGif(Env* env, DecodeGifTestCase testcase) {
w = width;
h = height;
c = channels;
return new uint8[frame_cnt * height * width * channels];
return new uint8[ptrdiff_t(frame_cnt) * height * width * channels];
},
&error_string));
ASSERT_NE(imgdata, nullptr);
Expand All @@ -72,7 +72,8 @@ TEST(GifTest, Gif) {
{testdata_path + "optimized.gif", 12, 20, 40, 3},
{testdata_path + "red_black.gif", 1, 16, 16, 3},
{testdata_path + "scan.gif", 12, 20, 40, 3},
{testdata_path + "squares.gif", 2, 16, 16, 3}});
{testdata_path + "squares.gif", 2, 16, 16, 3},
{testdata_path + "3g_multiframe.gif", 519, 1920, 1080, 3}});

for (const auto& tc : testcases) {
TestDecodeGif(env, tc);
Expand Down
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
1 change: 1 addition & 0 deletions tensorflow/core/lib/gif/testdata/BUILD
Expand Up @@ -16,6 +16,7 @@ filegroup(
"scan.gif",
"red_black.gif",
"squares.gif",
"3g_multiframe.gif",
"pendulum_sm.gif",
# Add groundtruth frames for `pendulum_sm.gif`.
# PNG format because it's lossless.
Expand Down

0 comments on commit 8dc723f

Please sign in to comment.