Create BoxReader helper class to parse JPEG2000 header

src/PIL/Jpeg2KImagePlugin.py

@@ -19,6 +19,79 @@ import struct
 from . import Image, ImageFile
 
 
+class BoxReader:
+    """
+    A small helper class to read fields stored in JPEG2000 header boxes
+    and to easily step into and read sub-boxes.
+    """
+
+    def __init__(self, fp, length=-1):
+        self.fp = fp
+        self.has_length = length >= 0
+        self.length = length
+        self.remaining_in_box = -1
+
+    def _can_read(self, num_bytes):
+        if self.remaining_in_box >= 0:
+            # Inside box contents: ensure read does not go past box boundaries
+            return num_bytes <= self.remaining_in_box
+        elif self.has_length:
+            # Outside box: ensure we don't read past the known file length
+            return self.fp.tell() + num_bytes <= self.length
+        else:
+            return True  # No length known, just read
+
+    def _read_bytes(self, num_bytes):
+        if not self._can_read(num_bytes):
+            raise SyntaxError("Not enough data in header")
+
+        data = self.fp.read(num_bytes)
+        if len(data) < num_bytes:
+            raise OSError(
+                f"Expected to read {num_bytes} bytes but only got {len(data)}."
+            )
+
+        if self.remaining_in_box > 0:
+            self.remaining_in_box -= num_bytes
+        return data
+
+    def read_fields(self, field_format):
+        size = struct.calcsize(field_format)
+        data = self._read_bytes(size)
+        return struct.unpack(field_format, data)
+
+    def read_boxes(self):
+        size = self.remaining_in_box
+        data = self._read_bytes(size)
+        return BoxReader(io.BytesIO(data), size)
+
+    def has_next_box(self):
+        if self.has_length:
+            return self.fp.tell() + self.remaining_in_box < self.length
+        else:
+            return True
+
+    def next_box_type(self):
+        # Skip the rest of the box if it has not been read
+        if self.remaining_in_box > 0:
+            self.fp.seek(self.remaining_in_box, os.SEEK_CUR)
+        self.remaining_in_box = -1
+
+        # Read the length and type of the next box
+        lbox, tbox = self.read_fields(">I4s")
+        if lbox == 1:
+            lbox = self.read_fields(">Q")[0]
+            hlen = 16
+        else:
+            hlen = 8
+
+        if lbox < hlen or not self._can_read(lbox - hlen):
+            raise SyntaxError("Invalid header length")
+
+        self.remaining_in_box = lbox - hlen
+        return tbox
+
+
 def _parse_codestream(fp):
     """Parse the JPEG 2000 codestream to extract the size and component
     count from the SIZ marker segment, returning a PIL (size, mode) tuple."""
@@ -58,50 +131,32 @@ def _parse_jp2_header(fp):
     color space information, returning a (size, mode, mimetype) tuple."""
 
     # Find the JP2 header box
+    reader = BoxReader(fp)
     header = None
     mimetype = None
-    while True:
-        lbox, tbox = struct.unpack(">I4s", fp.read(8))
-        if lbox == 1:
-            lbox = struct.unpack(">Q", fp.read(8))[0]
-            hlen = 16
-        else:
-            hlen = 8
-
-        if lbox < hlen:
-            raise SyntaxError("Invalid JP2 header length")
+    while reader.has_next_box():
+        tbox = reader.next_box_type()
 
         if tbox == b"jp2h":
-            header = fp.read(lbox - hlen)
+            header = reader.read_boxes()
             break
         elif tbox == b"ftyp":
-            if fp.read(4) == b"jpx ":
+            if reader.read_fields(">4s")[0] == b"jpx ":
                 mimetype = "image/jpx"
-            fp.seek(lbox - hlen - 4, os.SEEK_CUR)
-        else:
-            fp.seek(lbox - hlen, os.SEEK_CUR)
 
     if header is None:
-        raise SyntaxError("could not find JP2 header")
+        raise SyntaxError("Could not find JP2 header")
 
     size = None
     mode = None
     bpc = None
     nc = None
 
-    hio = io.BytesIO(header)
-    while True:
-        lbox, tbox = struct.unpack(">I4s", hio.read(8))
-        if lbox == 1:
-            lbox = struct.unpack(">Q", hio.read(8))[0]
-            hlen = 16
-        else:
-            hlen = 8
-
-        content = hio.read(lbox - hlen)
+    while header.has_next_box():
+        tbox = header.next_box_type()
 
         if tbox == b"ihdr":
-            height, width, nc, bpc, c, unkc, ipr = struct.unpack(">IIHBBBB", content)
+            height, width, nc, bpc, c, unkc, ipr = header.read_fields(">IIHBBBB")
             size = (width, height)
             if unkc:
                 if nc == 1 and (bpc & 0x7F) > 8:
@@ -116,9 +171,9 @@ def _parse_jp2_header(fp):
                     mode = "RGBA"
                 break
         elif tbox == b"colr":
-            meth, prec, approx = struct.unpack_from(">BBB", content)
+            meth, prec, approx = header.read_fields(">BBB")
             if meth == 1:
-                cs = struct.unpack_from(">I", content, 3)[0]
+                cs = header.read_fields(">I")[0]
                 if cs == 16:  # sRGB
                     if nc == 1 and (bpc & 0x7F) > 8:
                         mode = "I;16"
@@ -145,7 +200,7 @@ def _parse_jp2_header(fp):
                     break
 
     if size is None or mode is None:
-        raise SyntaxError("Malformed jp2 header")
+        raise SyntaxError("Malformed JP2 header")
 
     return (size, mode, mimetype)