#323 JPEGSegmentImageInputStream now rewrites duplicate SOF/SOS ids.

And emits warnings when it applies rewrites.

imageio/imageio-jpeg/src/main/java/com/twelvemonkeys/imageio/plugins/jpeg/JPEGImageReader.java

@@ -331,7 +331,7 @@ public final class JPEGImageReader extends ImageReaderBase {
 
         // JPEGSegmentImageInputStream that filters out/skips bad/unnecessary segments
         delegate.setInput(imageInput != null
-                          ? new JPEGSegmentImageInputStream(imageInput)
+                          ? new JPEGSegmentImageInputStream(imageInput, new JPEGSegmentStreamWarningDelegate())
                           : null, seekForwardOnly, ignoreMetadata);
     }
 
@@ -701,6 +701,7 @@ public final class JPEGImageReader extends ImageReaderBase {
             this.segments = segments;
 
             if (DEBUG) {
+                System.out.println("segments: " + segments);
                 System.out.println("Read metadata in " + (System.currentTimeMillis() - start) + " ms");
             }
         }
@@ -1251,6 +1252,12 @@ public final class JPEGImageReader extends ImageReaderBase {
         }
     }
 
+    private class JPEGSegmentStreamWarningDelegate implements JPEGSegmentStreamWarningListener {
+        @Override
+        public void warningOccurred(String warning) {
+            processWarningOccurred(warning);
+        }
+    }
     protected static void showIt(final BufferedImage pImage, final String pTitle) {
         ImageReaderBase.showIt(pImage, pTitle);
     }

imageio/imageio-jpeg/src/main/java/com/twelvemonkeys/imageio/plugins/jpeg/JPEGSegmentImageInputStream.java

@@ -36,14 +36,13 @@ import javax.imageio.stream.ImageInputStreamImpl;
 import java.io.EOFException;
 import java.io.IOException;
 import java.nio.charset.Charset;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
+import java.util.*;
 
 import static com.twelvemonkeys.lang.Validate.notNull;
 
 /**
- * ImageInputStream implementation that filters out certain JPEG segments.
+ * ImageInputStream implementation that filters out or rewrites
+ * certain JPEG segments.
  *
  * @author <a href="mailto:harald.kuhr@gmail.com">Harald Kuhr</a>
  * @author last modified by $Author: haraldk$
@@ -51,18 +50,29 @@ import static com.twelvemonkeys.lang.Validate.notNull;
  */
 final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
     // TODO: Rewrite JPEGSegment (from metadata) to store stream pos/length, and be able to replay data, and use instead of Segment?
-    // TODO: Change order of segments, to make sure APP0/JFIF is always before APP14/Adobe? What about EXIF?
-    // TODO: Insert fake APP0/JFIF if needed by the reader?
-    // TODO: Sort out ICC_PROFILE issues (duplicate sequence numbers etc)?
+    // TODO: Support multiple JPEG streams (SOI...EOI, SOI...EOI, ...) in a single file
 
     final private ImageInputStream stream;
+    final private JPEGSegmentStreamWarningListener warningListener;
+
+    final private Set<Integer> componentIds = new LinkedHashSet<>(4);
 
     private final List<Segment> segments = new ArrayList<Segment>(64);
     private int currentSegment = -1;
     private Segment segment;
 
-    JPEGSegmentImageInputStream(final ImageInputStream stream) {
+
+    JPEGSegmentImageInputStream(final ImageInputStream stream, final JPEGSegmentStreamWarningListener warningListener) {
         this.stream = notNull(stream, "stream");
+        this.warningListener = notNull(warningListener, "warningListener");
+    }
+
+    JPEGSegmentImageInputStream(final ImageInputStream stream) {
+        this(stream, JPEGSegmentStreamWarningListener.NULL_LISTENER);
+    }
+
+    private void processWarningOccured(final String warning) {
+        warningListener.warningOccurred(warning);
     }
 
     private Segment fetchSegment() throws IOException {
@@ -104,12 +114,6 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
                     realPosition++;
                 }
 
-                if (trash != 0) {
-                    // NOTE: We previously allowed these bytes to pass through to the native reader, as it could cope
-                    // and issued the correct warning. However, the native metadata chokes on it, so we'll mask it out.
-                    // TODO: Issue warning from the JPEGImageReader, telling how many bytes we skipped
-                }
-
                 marker = 0xff00 | stream.readUnsignedByte();
 
                 // Skip over 0xff padding between markers
@@ -118,6 +122,12 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
                     marker = 0xff00 | stream.readUnsignedByte();
                 }
 
+                if (trash != 0) {
+                    // NOTE: We previously allowed these bytes to pass through to the native reader, as it could cope
+                    // and issued the correct warning. However, the native metadata chokes on it, so we'll mask it out.
+                    processWarningOccured(String.format("Corrupt JPEG data: %d extraneous bytes before marker 0x%02x", trash, marker & 0xff));
+                }
+
                 // We are now handling all important segments ourselves, except APP1/Exif and APP14/Adobe,
                 // as these segments affects image decoding.
                 boolean appSegmentMarker = isAppSegmentMarker(marker);
@@ -134,17 +144,8 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
                         segments.add(segment);
                     }
                     else {
-                        long length;
-
-                        if (marker == JPEG.SOS) {
-                            // Treat rest of stream as a single segment (scanning for EOI is too much work)
-                            // TODO: For progressive, there will be more than one SOS...
-                            length = Long.MAX_VALUE - realPosition;
-                        }
-                        else {
-                            // Length including length field itself
-                            length = 2 + stream.readUnsignedShort();
-                        }
+                        // Length including length field itself
+                        long length = 2 + stream.readUnsignedShort();
 
                         if (isApp14Adobe && length != 16) {
                             // Need to rewrite this segment, so that it gets length 16 and discard the remaining bytes...
@@ -156,13 +157,24 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
                             // multiple quality tables with varying precision)
                             int qtInfo = stream.read();
                             if ((qtInfo & 0x10) == 0x10) {
-                                // TODO: Warning!
+                                processWarningOccured("16 bit DQT encountered");
                                 segment = new DownsampledDQTReplacement(realPosition, segment.end(), length, qtInfo, stream);
                             }
                             else {
                                 segment = new Segment(marker, realPosition, segment.end(), length);
                             }
                         }
+                        else if (isSOFMarker(marker)) {
+                            // Replace duplicate SOFn component ids
+                            byte[] data = replaceDuplicateSOFnComponentIds(marker, length);
+                            segment = new ReplacementSegment(marker, realPosition, segment.end(), length, data);
+                        }
+                        else if (marker == JPEG.SOS) {
+                            // Replace duplicate SOS component selectors
+                            byte[] data = replaceDuplicateSOSComponentSelectors(length);
+
+                            segment = new ReplacementSegment(marker, realPosition, segment.end(), length, data);
+                        }
                         else {
                             segment = new Segment(marker, realPosition, segment.end(), length);
                         }
@@ -172,6 +184,12 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
 
                     currentSegment = segments.size() - 1;
 
+                    if (marker == JPEG.SOS) {
+                        // Treat rest of stream as a single segment (scanning for EOI is too much work)
+                        // TODO: For progressive, there will be more than one SOS...
+                        segments.add(new Segment(-1, segment.realEnd(), segment.end(), Long.MAX_VALUE - segment.realEnd()));
+                    }
+
                     if (streamPos >= segment.start && streamPos < segment.end()) {
                         segment.seek(stream, streamPos);
 
@@ -205,6 +223,76 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
         return segment;
     }
 
+    private byte[] replaceDuplicateSOSComponentSelectors(long length) throws IOException {
+        // See: http://www.hackerfactor.com/blog/index.php?/archives/588-JPEG-Patches.html
+        byte[] data = readSegment(JPEG.SOS, (int) length, stream);
+
+        // Detect duplicates
+        Set<Integer> componentSelectors = new LinkedHashSet<>(4);
+        boolean duplicatesFound = false;
+        int off = 5;
+        while (off < length - 3) {
+            int selector = data[off] & 0xff;
+            if (!componentSelectors.add(selector)) {
+                processWarningOccured(String.format("Duplicate component selector %d in SOS", selector));
+                duplicatesFound = true;
+            }
+
+            off += 2;
+        }
+
+        // Replace all with the component ids in order, as this is the most likely situation
+        if (duplicatesFound) {
+            off = 5;
+
+            Iterator<Integer> ids = componentIds.iterator();
+            while (off < length - 3) {
+                if (ids.hasNext()) {
+                    data[off] = (byte) (int) ids.next();
+                }
+                // Otherwise we'll have an undefined component selector...
+
+                off += 2;
+            }
+        }
+        return data;
+    }
+
+    private byte[] replaceDuplicateSOFnComponentIds(int marker, long length) throws IOException {
+        byte[] data = readSegment(marker, (int) length, stream);
+
+        int off = 10;
+        while (off < length) {
+            int id = data[off] & 0xff;
+            if (!componentIds.add(id)) {
+                processWarningOccured(String.format("Duplicate component ID %d in SOF", id));
+
+                id++;
+                while (!componentIds.add(id)) {
+                    id++;
+                }
+
+                data[off] = (byte) id;
+            }
+
+            off += 3;
+        }
+        return data;
+    }
+
+    private static byte[] readSegment(final int marker, final int length, final ImageInputStream stream) throws IOException {
+        byte[] data = new byte[length];
+
+        data[0] = (byte) ((marker >> 8) & 0xff);
+        data[1] = (byte) (marker & 0xff);
+        data[2] = (byte) (((length - 2) >> 8) & 0xff);
+        data[3] = (byte) ((length - 2) & 0xff);
+
+        stream.readFully(data, 4, length - 4);
+
+        return data;
+    }
+
     private static boolean isAppSegmentWithId(final String segmentId, final ImageInputStream stream) throws IOException {
         notNull(segmentId, "segmentId");
 
@@ -261,6 +349,30 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
         return marker >= JPEG.APP0 && marker <= JPEG.APP15;
     }
 
+    static boolean isSOFMarker(final int marker) {
+        switch (marker) {
+            case JPEG.SOF0:
+            case JPEG.SOF1:
+            case JPEG.SOF2:
+            case JPEG.SOF3:
+
+            case JPEG.SOF5:
+            case JPEG.SOF6:
+            case JPEG.SOF7:
+
+            case JPEG.SOF9:
+            case JPEG.SOF10:
+            case JPEG.SOF11:
+
+            case JPEG.SOF13:
+            case JPEG.SOF14:
+            case JPEG.SOF15:
+                return true;
+            default:
+                return false;
+        }
+    }
+
     private void repositionAsNecessary() throws IOException {
         if (segment == null || streamPos < segment.start || streamPos >= segment.end()) {
             try {
@@ -380,16 +492,17 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
         }
 
         private static byte[] createMarkerFixedLength(final ImageInputStream stream) throws IOException {
-            byte[] segmentData = new byte[16];
-
-            segmentData[0] = (byte) ((JPEG.APP14 >> 8) & 0xff);
-            segmentData[1] = (byte) (JPEG.APP14 & 0xff);
-            segmentData[2] = (byte) 0;
-            segmentData[3] = (byte) 14;
-
-            stream.readFully(segmentData, 4, segmentData.length - 4);
-
-            return segmentData;
+            return readSegment(JPEG.APP14, 16, stream);
+//            byte[] segmentData = new byte[16];
+//
+//            segmentData[0] = (byte) ((JPEG.APP14 >> 8) & 0xff);
+//            segmentData[1] = (byte) (JPEG.APP14 & 0xff);
+//            segmentData[2] = (byte) 0;
+//            segmentData[3] = (byte) 14;
+//
+//            stream.readFully(segmentData, 4, segmentData.length - 4);
+//
+//            return segmentData;
         }
     }
 
@@ -405,11 +518,10 @@ final class JPEGSegmentImageInputStream extends ImageInputStreamImpl {
         }
 
         private static byte[] createMarkerFixedLength(final int length, final int qtInfo, final ImageInputStream stream) throws IOException {
-            byte[] replacementData = new byte[length];
-
             int numQTs = length / 128;
             int newSegmentLength = 2 + (1 + 64) * numQTs; // Len + (qtInfo + qtSize) * numQTs
 
+            byte[] replacementData = new byte[length];
             replacementData[0] = (byte) ((JPEG.DQT >> 8) & 0xff);
             replacementData[1] = (byte) (JPEG.DQT & 0xff);
             replacementData[2] = (byte) ((newSegmentLength >> 8) & 0xff);

imageio/imageio-jpeg/src/main/java/com/twelvemonkeys/imageio/plugins/jpeg/JPEGSegmentStreamWarningListener.java

@@ -0,0 +1,13 @@
+package com.twelvemonkeys.imageio.plugins.jpeg;
+
+/**
+ * JPEGSegmentStreamWarningListener
+ */
+interface JPEGSegmentStreamWarningListener {
+    void warningOccurred(String warning);
+
+    JPEGSegmentStreamWarningListener NULL_LISTENER = new JPEGSegmentStreamWarningListener() {
+        @Override
+        public void warningOccurred(final String warning) {}
+    };
+}

imageio/imageio-jpeg/src/test/java/com/twelvemonkeys/imageio/plugins/jpeg/JPEGImageReaderTest.java

@@ -62,6 +62,7 @@ import static com.twelvemonkeys.imageio.util.IIOUtil.lookupProviderByName;
 import static org.junit.Assert.*;
 import static org.junit.Assume.assumeNoException;
 import static org.junit.Assume.assumeNotNull;
+import static org.mockito.AdditionalMatchers.and;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.*;
@@ -1608,10 +1609,16 @@ public class JPEGImageReaderTest extends ImageReaderAbstractTest<JPEGImageReader
         try {
             reader.setInput(ImageIO.createImageInputStream(getClassLoaderResource("/jpeg/exif-jfif-app13-app14ycck-3channel.jpg")));
 
+            IIOReadWarningListener listener = mock(IIOReadWarningListener.class);
+            reader.addIIOReadWarningListener(listener);
+
             assertEquals(310, reader.getWidth(0));
             assertEquals(384, reader.getHeight(0));
 
             BufferedImage image = reader.read(0, null);
+
+            verify(listener, times(1)).warningOccurred(eq(reader), matches("(?i).*Adobe App14.*(?-i)CMYK.*SOF.*"));
+
             assertNotNull(image);
             assertEquals(310, image.getWidth());
             assertEquals(384, image.getHeight());
@@ -1621,4 +1628,32 @@ public class JPEGImageReaderTest extends ImageReaderAbstractTest<JPEGImageReader
             reader.dispose();
         }
     }
+
+    @Test
+    public void testReadDuplicateComponentIds() throws IOException {
+        JPEGImageReader reader = createReader();
+
+        try {
+            reader.setInput(ImageIO.createImageInputStream(getClassLoaderResource("/jpeg/duplicate-component-ids.jpg")));
+
+            IIOReadWarningListener listener = mock(IIOReadWarningListener.class);
+            reader.addIIOReadWarningListener(listener);
+
+            assertEquals(367, reader.getWidth(0));
+            assertEquals(242, reader.getHeight(0));
+
+            BufferedImage image = reader.read(0, null);
+
+            verify(listener, times(1)).warningOccurred(eq(reader), and(matches("(?i).*duplicate component id.*(?-i)SOF.*"), contains("1")));
+            verify(listener, times(1)).warningOccurred(eq(reader), and(matches("(?i).*duplicate component selector.*(?-i)SOS.*"), contains("1")));
+
+            assertNotNull(image);
+            assertEquals(367, image.getWidth());
+            assertEquals(242, image.getHeight());
+            assertEquals(ColorSpace.TYPE_RGB, image.getColorModel().getColorSpace().getType());
+        }
+        finally {
+            reader.dispose();
+        }
+    }
 }

imageio/imageio-jpeg/src/test/java/com/twelvemonkeys/imageio/plugins/jpeg/JPEGSegmentImageInputStreamTest.java

@@ -59,7 +59,7 @@ public class JPEGSegmentImageInputStreamTest {
         ImageIO.setUseCache(false);
     }
 
-    protected URL getClassLoaderResource(final String pName) {
+    private URL getClassLoaderResource(final String pName) {
         return getClass().getResource(pName);
     }
 
@@ -119,7 +119,7 @@ public class JPEGSegmentImageInputStreamTest {
             length++;
         }
 
-        assertThat(length, new LessOrEqual<Long>(10203l)); // In no case should length increase
+        assertThat(length, new LessOrEqual<>(10203L)); // In no case should length increase
 
         assertEquals(9607L, length); // May change, if more chunks are passed to reader...
     }
@@ -149,7 +149,7 @@ public class JPEGSegmentImageInputStreamTest {
             length++;
         }
 
-        assertEquals(9281L, length); // Sanity check: same as file size
+        assertEquals(9281L, length); // Sanity check: same as file size, except..?
     }
 
     @Test
@@ -162,7 +162,7 @@ public class JPEGSegmentImageInputStreamTest {
         assertEquals(JPEG.APP1, appSegments.get(0).marker());
         assertEquals("Exif", appSegments.get(0).identifier());
 
-        stream.seek(0l);
+        stream.seek(0L);
 
         long length = 0;
         while (stream.read() != -1) {