Fix MPO index placement after Adobe JPEG headers

Tests/test_file_mpo.py

@@ -2,7 +2,7 @@
 
 import warnings
 from io import BytesIO
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
@@ -22,6 +22,9 @@
     skip_unless_feature,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 test_files = ["Tests/images/sugarshack.mpo", "Tests/images/frozenpond.mpo"]
 
 pytestmark = skip_unless_feature("jpg")
@@ -374,6 +377,51 @@ def test_save_all() -> None:
     assert "mp" not in jpg.info
 
 
+def test_save_all_cmyk() -> None:
+    first = Image.new("CMYK", (1, 1))
+    second = Image.new("CMYK", (1, 1), (255, 255, 255, 255))
+    im_reloaded = roundtrip(first, save_all=True, append_images=[second])
+    assert_image_equal(first, im_reloaded)
+
+    im_reloaded.seek(1)
+    assert_image_equal(second, im_reloaded)
+
+
+@pytest.mark.parametrize(
+    "mode,options",
+    [
+        ("RGB", {}),
+        ("CMYK", {}),
+        ("RGB", {"keep_rgb": True, "dpi": (72, 72)}),
+    ],
+)
+def test_save_all_write_only(
+    tmp_path: Path, mode: str, options: dict[str, Any]
+) -> None:
+    first = Image.new(mode, (16, 8))
+    second = Image.new(mode, (24, 16))
+    references = []
+    for frame in (first, second):
+        buffer = BytesIO()
+        frame.save(buffer, "JPEG", **options)
+        references.append(buffer.getvalue())
+    path = tmp_path / "write-only.mpo"
+    with path.open("wb") as fp:
+        first.save(fp, "MPO", save_all=True, append_images=[second], **options)
+
+    with Image.open(path) as reopened:
+        assert reopened.format == "MPO"
+        assert isinstance(reopened, MpoImagePlugin.MpoImageFile)
+        assert reopened.n_frames == 2
+        assert reopened.size == first.size
+        with Image.open(BytesIO(references[0])) as expected:
+            assert_image_equal(expected, reopened)
+        reopened.seek(1)
+        assert reopened.size == second.size
+        with Image.open(BytesIO(references[1])) as expected:
+            assert_image_equal(expected, reopened)
+
+
 def test_save_xmp() -> None:
     im = Image.new("RGB", (1, 1))
     im2 = Image.new("RGB", (1, 1), "#f00")

docs/releasenotes/13.1.0.rst

@@ -49,6 +49,13 @@ TODO
 Other changes
 =============
 
+Saving MPO images with Adobe JPEG headers
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Saving CMYK images as MPO, or RGB images with ``keep_rgb=True``, wrote the
+MP Index at an incorrect offset after the Adobe JPEG header. The resulting
+files were reopened as a single JPEG instead of an MPO. This has been fixed.
+
 Resizing I;16B and I;16N images
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

src/PIL/MpoImagePlugin.py

@@ -40,6 +40,24 @@
     from typing import IO, Any
 
 
+class _MarkerCapture:
+    """Forward JPEG data while retaining only the marker's search window."""
+
+    def __init__(self, fp: IO[bytes], start: int) -> None:
+        self.fp = fp
+        self.start = start
+        self.position = 0
+        self.window = bytearray()
+
+    def write(self, data: bytes) -> int:
+        begin = max(0, self.start - self.position)
+        end = min(len(data), self.start + 20 - self.position)
+        if end > begin:
+            self.window.extend(data[begin:end])
+        self.position += len(data)
+        return self.fp.write(data)
+
+
 def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     JpegImagePlugin._save(im, fp, filename)
 
@@ -50,8 +68,9 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
         _save(im, fp, filename)
         return
 
-    mpf_offset = 28
+    mpf_offset = 18
     prev_offset = 0
+    capture: _MarkerCapture | None = None
     offsets: list[int] = []
     im_sequences = [im, *append_images]
     total = sum(getattr(seq, "n_frames", 1) for seq in im_sequences)
@@ -72,7 +91,9 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
                         im_frame.encoderinfo["exif"] = exif
                     mpf_offset += 4 + len(exif)
 
-                JpegImagePlugin._save(im_frame, fp, filename)
+                # Do not expose fileno: JPEG data must pass through write().
+                capture = _MarkerCapture(fp, mpf_offset)
+                JpegImagePlugin._save(im_frame, cast("IO[bytes]", capture), filename)
             else:
                 encoderinfo = im_frame._attach_default_encoderinfo(im)
                 im_frame.save(fp, "JPEG")
@@ -81,6 +102,9 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
             offsets.append(offset - prev_offset)
             prev_offset = offset
 
+    assert capture is not None
+    mpf_offset += 8 + capture.window.rindex(b"\xff\xe2")
+
     ifd = TiffImagePlugin.ImageFileDirectory_v2()
     ifd[0xB000] = b"0100"
     ifd[0xB001] = len(offsets)