diff --git a/Tests/test_file_qoi.py b/Tests/test_file_qoi.py index b9becb24f77..3057dabc05e 100644 --- a/Tests/test_file_qoi.py +++ b/Tests/test_file_qoi.py @@ -1,5 +1,6 @@ from __future__ import annotations +from io import BytesIO from pathlib import Path import pytest @@ -32,6 +33,15 @@ def test_invalid_file() -> None: QoiImagePlugin.QoiImageFile(invalid_file) +@pytest.mark.parametrize("length", (14, 44, 48, 55)) +def test_truncated(length: int) -> None: + with open("Tests/images/pil123rgba.qoi", "rb") as fp: + data = fp.read()[:length] + with Image.open(BytesIO(data)) as im: + with pytest.raises(ValueError, match="not enough image data"): + im.load() + + def test_op_index() -> None: # QOI_OP_INDEX as the first chunk with Image.open("Tests/images/op_index.qoi") as im: diff --git a/src/PIL/QoiImagePlugin.py b/src/PIL/QoiImagePlugin.py index e7a6cee16db..3226c46c509 100644 --- a/src/PIL/QoiImagePlugin.py +++ b/src/PIL/QoiImagePlugin.py @@ -61,12 +61,20 @@ def decode(self, buffer: Image.DecoderInput) -> tuple[int, int]: bands = Image.getmodebands(self.mode) dest_length = self.state.xsize * self.state.ysize * bands while len(data) < dest_length: - byte = self.fd.read(1)[0] + byte_data = self.fd.read(1) + if not byte_data: + break + byte = byte_data[0] value: bytes | bytearray if byte == 0b11111110 and self._previous_pixel: # QOI_OP_RGB - value = bytearray(self.fd.read(3)) + self._previous_pixel[3:] + rgb_data = self.fd.read(3) + if len(rgb_data) < 3: + break + value = bytearray(rgb_data) + self._previous_pixel[3:] elif byte == 0b11111111: # QOI_OP_RGBA value = self.fd.read(4) + if len(value) < 4: + break else: op = byte >> 6 if op == 0: # QOI_OP_INDEX @@ -86,7 +94,10 @@ def decode(self, buffer: Image.DecoderInput) -> tuple[int, int]: ) ) elif op == 2 and self._previous_pixel: # QOI_OP_LUMA - second_byte = self.fd.read(1)[0] + second_byte_data = self.fd.read(1) + if not second_byte_data: + break + second_byte = second_byte_data[0] diff_green = (byte & 0b00111111) - 32 diff_red = ((second_byte & 0b11110000) >> 4) - 8 diff_blue = (second_byte & 0b00001111) - 8