diff --git a/test/test_image.py b/test/test_image.py index 743b8cceff9..878e49b2f29 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -949,6 +949,101 @@ def le16(v): assert (out[1] == 0).all() +def _craft_gif(*, canvas, gct, bg, frame_size, pixels, lzw_min): + # Minimal single-frame GIF89a builder. `gct` is a flat RGB palette whose + # length implies the colour count; `lzw_min` is the LZW minimum code size, + # which the GIF format lets an encoder choose independently of the palette + # size. Codes are packed at a fixed width, which is valid here because + # these streams are far too short to trigger a code-width increase. + def le16(v): + return bytes([v & 0xFF, (v >> 8) & 0xFF]) + + n_colors = len(gct) // 3 + gct_bpp = n_colors.bit_length() - 1 + clear, eoi, width = 1 << lzw_min, (1 << lzw_min) + 1, lzw_min + 1 + + packed_codes, cur, nbits = bytearray(), 0, 0 + for code in [clear, *pixels, eoi]: + cur |= code << nbits + nbits += width + while nbits >= 8: + packed_codes.append(cur & 0xFF) + cur >>= 8 + nbits -= 8 + if nbits: + packed_codes.append(cur & 0xFF) + + sub_blocks = b"" + for i in range(0, len(packed_codes), 255): + chunk = packed_codes[i : i + 255] + sub_blocks += bytes([len(chunk)]) + bytes(chunk) + + return ( + b"GIF89a" + + le16(canvas) + + le16(canvas) + + bytes([0x80 | (gct_bpp - 1), bg, 0]) # LSD: GCT present, background index + + gct + + bytes([0x2C]) # image descriptor + + le16(0) + + le16(0) + + le16(frame_size) + + le16(frame_size) + + bytes([0]) # no local colormap + + bytes([lzw_min]) + + sub_blocks + + bytes([0, 0x3B]) # block terminator + trailer + ) + + +# A 2-colour global colormap: the palette allocation is only 6 bytes, so any +# index above 1 reads past it. +_TINY_GCT = bytes([17, 34, 51, 255, 255, 255]) + + +@pytest.mark.parametrize("scripted", (True, False)) +@pytest.mark.parametrize("bg, expected_bg_color", [(0, (17, 34, 51)), (255, (0, 0, 0))]) +def test_decode_gif_out_of_range_background_color(scripted, bg, expected_bg_color): + # Non-regression test: SBackGroundColor is an unvalidated byte (0-255) from + # the logical screen descriptor, but the global colormap may hold as few as + # 2 entries. Before the fix, bg=255 against a 2-colour palette read 759 + # bytes past the allocation and those heap bytes became the background + # colour of the output tensor (heap info leak, CWE-125). + encoded = _craft_gif(canvas=4, gct=_TINY_GCT, bg=bg, frame_size=1, pixels=[1], lzw_min=2) + f = torch.jit.script(decode_gif) if scripted else decode_gif + out = f(torch.frombuffer(bytearray(encoded), dtype=torch.uint8)) + + assert out.shape == (3, 4, 4) + # The 1x1 frame covers only the top-left pixel; the rest is background. + assert tuple(out[:, 0, 0].tolist()) == (255, 255, 255) + assert (out[:, 3, 3] == torch.tensor(expected_bg_color, dtype=torch.uint8)).all() + + +@pytest.mark.parametrize("scripted", (True, False)) +@pytest.mark.parametrize("n_colors", (2, 256)) +def test_decode_gif_out_of_range_palette_index(scripted, n_colors): + # Non-regression test: the LZW minimum code size is read independently of + # the colormap size, so a GIF may declare a 2-entry palette while emitting + # raster values up to 255. Before the fix, cmap->Colors[c] read up to 762 + # bytes past the palette and wrote those heap bytes straight into the + # output tensor, once per pixel (heap info leak, CWE-125). + gct = _TINY_GCT if n_colors == 2 else b"".join(bytes([i, i, i]) for i in range(256)) + encoded = _craft_gif(canvas=4, gct=gct, bg=0, frame_size=4, pixels=[255] * 16, lzw_min=8) + f = torch.jit.script(decode_gif) if scripted else decode_gif + out = f(torch.frombuffer(bytearray(encoded), dtype=torch.uint8)) + + assert out.shape == (3, 4, 4) + if n_colors == 2: + # Index 255 is out of range: the pixel is skipped and the background + # (palette entry 0) shows through. + expected = (17, 34, 51) + else: + # Control: the same raster stream against a full 256-entry palette must + # still resolve to entry 255, i.e. the fix must not clip valid indices. + expected = (255, 255, 255) + assert (out == torch.tensor(expected, dtype=torch.uint8)[:, None, None]).all() + + @pytest.mark.parametrize( "decode_fun, match", [ diff --git a/torchvision/csrc/io/image/cpu/decode_gif.cpp b/torchvision/csrc/io/image/cpu/decode_gif.cpp index 4e21fb61a5d..ec416fc998c 100644 --- a/torchvision/csrc/io/image/cpu/decode_gif.cpp +++ b/torchvision/csrc/io/image/cpu/decode_gif.cpp @@ -89,8 +89,14 @@ torch::stable::Tensor decode_gif(const torch::stable::Tensor& encoded_data) { STD_TORCH_CHECK( num_images > 0, "GIF file should contain at least one image!"); + // SBackGroundColor is an unvalidated byte (0-255) read straight from the + // logical screen descriptor, while SColorMap->ColorCount may be as small as + // 2. GIFLIB deliberately leaves this bound to the caller (see the comment at + // the end of DGifGetScreenDesc()), so we have to check it here. An + // out-of-range background index means "no background": we leave bg black. GifColorType bg = {0, 0, 0}; - if (gifFile->SColorMap) { + if (gifFile->SColorMap && + gifFile->SBackGroundColor < gifFile->SColorMap->ColorCount) { bg = gifFile->SColorMap->Colors[gifFile->SBackGroundColor]; } @@ -176,7 +182,11 @@ torch::stable::Tensor decode_gif(const torch::stable::Tensor& encoded_data) { continue; } auto c = img.RasterBits[h * desc.Width + w]; - if (c == gcb.TransparentColor) { + // The LZW minimum code size is read independently of the colour table + // size, so c may be greater than cmap->ColorCount on a crafted (or + // merely sloppy) GIF. Treat an out-of-range index as transparent + // rather than erroring out, which is what browsers and Pillow do. + if (c == gcb.TransparentColor || c >= cmap->ColorCount) { continue; } GifColorType rgb = cmap->Colors[c];