Merge pull request #3258 from SixLabors/fix/jpeg-marker-recovery-v4

[V4] Recover JPEG restart markers through the shared reader

Author: JimBobSquarePants

src/ImageSharp/Formats/Jpeg/Components/Decoder/ArithmeticDecodingTable.cs

@@ -5,10 +5,19 @@ namespace SixLabors.ImageSharp.Formats.Jpeg.Components.Decoder;
 
 internal class ArithmeticDecodingTable
 {
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ArithmeticDecodingTable"/> class with the JPEG conditioning defaults.
+    /// </summary>
+    /// <param name="tableClass">The table class: zero for DC or one for AC.</param>
+    /// <param name="identifier">The table identifier from the scan or DAC segment.</param>
     public ArithmeticDecodingTable(byte tableClass, byte identifier)
     {
         this.TableClass = tableClass;
         this.Identifier = identifier;
+
+        // ITU T.81 F.1.4.4.1.4 specifies DC bounds L = 0 and U = 1 when DAC is absent.
+        // U occupies the upper nibble, so the encoded DC value is 0x10. F.1.4.4.2.1 specifies AC Kx = 5.
+        this.Configure(tableClass == 0 ? (byte)0x10 : (byte)5);
     }
 
     public byte TableClass { get; }

src/ImageSharp/Formats/Jpeg/Components/Decoder/ArithmeticScanDecoder.cs

@@ -202,6 +202,9 @@ public ArithmeticScanDecoder(BufferedReadStream stream, SpectralConverter conver
         this.ct = -16; // Force reading 2 initial bytes to fill C.
     }
 
+    /// <inheritdoc/>
+    public SegmentIntegrityHandling SegmentIntegrityHandling { get; set; }
+
     /// <inheritdoc/>
     public int ResetInterval
     {
@@ -276,6 +279,13 @@ public void InjectFrameData(JpegFrame frame, IRawJpegData jpegData)
         this.spectralConverter.InjectFrameData(frame, jpegData);
     }
 
+    /// <summary>
+    /// Gets the declared conditioning table or stores the JPEG defaults for an undeclared table.
+    /// </summary>
+    /// <param name="arithmeticDecodingTables">The conditioning tables for this image.</param>
+    /// <param name="isDcTable">True for a DC table. False for an AC table.</param>
+    /// <param name="identifier">The table identifier selected by the component.</param>
+    /// <returns>The declared table or a table with the standard default values.</returns>
     private static ArithmeticDecodingTable GetArithmeticTable(List<ArithmeticDecodingTable> arithmeticDecodingTables, bool isDcTable, int identifier)
     {
         int tableClass = isDcTable ? 0 : 1;
@@ -288,7 +298,11 @@ private static ArithmeticDecodingTable GetArithmeticTable(List<ArithmeticDecodin
             }
         }
 
-        return null;
+        // Omitted DAC entries are valid. Cache the defaults so later scans reuse the table without another allocation.
+        // A later DAC segment replaces this entry through the existing table-update path.
+        ArithmeticDecodingTable table = new((byte)tableClass, (byte)identifier);
+        arithmeticDecodingTables.Add(table);
+        return table;
     }
 
     private ArithmeticStatistics CreateOrGetStatisticsBin(bool dc, int identifier, bool reset = false)
@@ -1100,9 +1114,10 @@ private bool HandleRestart()
     {
         if (this.restartInterval > 0 && (--this.todo) == 0)
         {
-            if (this.scanBuffer.Marker == JpegConstants.Markers.XFF)
+            // Use the same marker recovery and integrity policy as Huffman-coded scans.
+            if (this.scanBuffer.Marker == JpegConstants.Markers.XFF || this.scanBuffer.HasReservedMarker())
             {
-                if (!this.scanBuffer.FindNextMarker())
+                if (!this.scanBuffer.FindNextMarker(this.SegmentIntegrityHandling))
                 {
                     return false;
                 }
@@ -1119,8 +1134,7 @@ private bool HandleRestart()
                 component.AcStatistics?.Reset();
             }
 
-            this.Reset();
-
+            // Reset clears the detected marker. Classify it first so a real segment marker can be read again.
             if (this.scanBuffer.HasRestartMarker())
             {
                 this.Reset();

src/ImageSharp/Formats/Jpeg/Components/Decoder/HuffmanScanDecoder.cs

@@ -85,6 +85,9 @@ public HuffmanScanDecoder(
         this.acHuffmanTables = new HuffmanTable[maxTables];
     }
 
+    /// <inheritdoc/>
+    public SegmentIntegrityHandling SegmentIntegrityHandling { get; set; }
+
     /// <summary>
     /// Sets reset interval determined by RST markers.
     /// </summary>
@@ -745,9 +748,10 @@ private bool HandleRestart()
     {
         if (this.restartInterval > 0 && (--this.todo) == 0)
         {
-            if (this.scanBuffer.Marker == JpegConstants.Markers.XFF)
+            // A reserved marker cannot complete a restart interval. The shared reader applies the recovery policy.
+            if (this.scanBuffer.Marker == JpegConstants.Markers.XFF || this.scanBuffer.HasReservedMarker())
             {
-                if (!this.scanBuffer.FindNextMarker())
+                if (!this.scanBuffer.FindNextMarker(this.SegmentIntegrityHandling))
                 {
                     return false;
                 }

src/ImageSharp/Formats/Jpeg/Components/Decoder/IJpegScanDecoder.cs

@@ -10,6 +10,11 @@ namespace SixLabors.ImageSharp.Formats.Jpeg.Components.Decoder;
 /// </summary>
 internal interface IJpegScanDecoder
 {
+    /// <summary>
+    /// Sets the recovery policy for errors in compressed scan data.
+    /// </summary>
+    public SegmentIntegrityHandling SegmentIntegrityHandling { set; }
+
     /// <summary>
     /// Sets the reset interval.
     /// </summary>

src/ImageSharp/Formats/Jpeg/Components/Decoder/JpegBitReader.cs

@@ -84,6 +84,13 @@ public void Reset()
     [MethodImpl(InliningOptions.ShortMethod)]
     public readonly bool HasBadMarker() => this.Marker != JpegConstants.Markers.XFF && !this.HasRestartMarker();
 
+    /// <summary>
+    /// Determines whether the detected marker is TEM or a reserved code before SOF0.
+    /// </summary>
+    /// <returns>True for TEM and reserved marker codes.</returns>
+    [MethodImpl(InliningOptions.ShortMethod)]
+    public readonly bool HasReservedMarker() => this.Marker < JpegConstants.Markers.SOF0;
+
     [MethodImpl(InliningOptions.AlwaysInline)]
     public void FillBuffer()
     {
@@ -177,39 +184,56 @@ private ulong GetBytes()
         return temp;
     }
 
+    /// <summary>
+    /// Finds the next marker at a restart boundary and applies the selected image-data recovery policy.
+    /// </summary>
+    /// <param name="integrityHandling">The recovery policy for compressed image data.</param>
+    /// <returns>True if a marker was found. False if the stream ended first.</returns>
     [MethodImpl(InliningOptions.ShortMethod)]
-    public bool FindNextMarker()
+    public bool FindNextMarker(SegmentIntegrityHandling integrityHandling)
     {
-        while (true)
+        do
         {
-            int b = this.stream.ReadByte();
-            if (b == -1)
+            // The bit reader can already hold a marker from its entropy-data lookahead.
+            // Libjpeg treats codes before SOF0 as damaged image data and scans forward to another marker.
+            // Only IgnoreImageData permits that recovery. Real segment markers stay with the segment parser.
+            if (this.HasReservedMarker() && integrityHandling != SegmentIntegrityHandling.IgnoreImageData)
             {
-                return false;
+                JpegThrowHelper.ThrowInvalidImageContentException("Unexpected marker before a JPEG restart marker.");
             }
 
-            // Found a marker.
-            if (b == JpegConstants.Markers.XFF)
+            while (true)
             {
-                while (b == JpegConstants.Markers.XFF)
+                int b = this.stream.ReadByte();
+                if (b == -1)
                 {
-                    // Loop here to discard any padding FF bytes on terminating marker.
-                    b = this.stream.ReadByte();
-                    if (b == -1)
-                    {
-                        return false;
-                    }
+                    return false;
                 }
 
-                // Found a valid marker. Exit loop
-                if (b != 0)
+                if (b == JpegConstants.Markers.XFF)
                 {
-                    this.Marker = (byte)b;
-                    this.MarkerPosition = this.stream.Position - 2;
-                    return true;
+                    // Multiple FF bytes are marker padding. FF 00 is a stuffed data byte, not a marker.
+                    while (b == JpegConstants.Markers.XFF)
+                    {
+                        b = this.stream.ReadByte();
+                        if (b == -1)
+                        {
+                            return false;
+                        }
+                    }
+
+                    if (b != 0)
+                    {
+                        this.Marker = (byte)b;
+                        this.MarkerPosition = this.stream.Position - 2;
+                        break;
+                    }
                 }
             }
         }
+        while (this.HasReservedMarker());
+
+        return true;
     }
 
     [MethodImpl(InliningOptions.AlwaysInline)]

src/ImageSharp/Formats/Jpeg/JpegConstants.cs

@@ -271,6 +271,14 @@ internal static class Markers
         /// </remarks>
         /// </summary>
         public const byte RST7 = 0xD7;
+
+        /// <summary>
+        /// Temporary private marker for arithmetic coding.
+        /// </summary>
+        /// <remarks>
+        /// This marker stands alone. It has no length field or payload.
+        /// </remarks>
+        public const byte TEM = 0x01;
     }
 
     /// <summary>

src/ImageSharp/Formats/Jpeg/JpegDecoderCore.cs

@@ -188,8 +188,9 @@ public static JpegFileMarker FindNextFileMarker(BufferedReadStream stream)
                     }
                 }
 
-                // Found a valid marker. Exit loop
-                if (b is not 0 and (< JpegConstants.Markers.RST0 or > JpegConstants.Markers.RST7))
+                // FF 00 is a stuffed data byte. RSTn and TEM have no length field (ITU T.81, Table B.1).
+                // Skip these standalone markers so the next marker's bytes are not read as a length.
+                if (b is not 0 and not JpegConstants.Markers.TEM and (< JpegConstants.Markers.RST0 or > JpegConstants.Markers.RST7))
                 {
                     return new JpegFileMarker((byte)(uint)b, stream.Position - 2);
                 }
@@ -324,6 +325,9 @@ internal void ParseStream(BufferedReadStream stream, SpectralConverter spectralC
 
         this.scanDecoder ??= new HuffmanScanDecoder(stream, spectralConverter, cancellationToken);
 
+        // Apply the image-data policy to the active decoder, including one supplied by TIFF.
+        this.scanDecoder.SegmentIntegrityHandling = this.Options.SegmentIntegrityHandling;
+
         this.Metadata ??= new ImageMetadata();
 
         Span<byte> markerBuffer = stackalloc byte[2];
@@ -386,7 +390,11 @@ internal void ParseStream(BufferedReadStream stream, SpectralConverter spectralC
                     case JpegConstants.Markers.SOF10:
                     case JpegConstants.Markers.SOF13:
                     case JpegConstants.Markers.SOF14:
-                        this.scanDecoder = new ArithmeticScanDecoder(stream, spectralConverter, cancellationToken);
+                        this.scanDecoder = new ArithmeticScanDecoder(stream, spectralConverter, cancellationToken)
+                        {
+                            SegmentIntegrityHandling = this.Options.SegmentIntegrityHandling
+                        };
+
                         if (this.resetInterval.HasValue)
                         {
                             this.scanDecoder.ResetInterval = this.resetInterval.Value;
@@ -1629,6 +1637,8 @@ private void ProcessStartOfScanMarker(BufferedReadStream stream, int remaining)
 
         if (this.scanDecoder is ArithmeticScanDecoder arithmeticScanDecoder)
         {
+            // An arithmetic JPEG can omit DAC entirely. The scan decoder adds defaults for undeclared entries.
+            this.arithmeticDecodingTables ??= new List<ArithmeticDecodingTable>(4);
             arithmeticScanDecoder.InitDecodingTables(this.arithmeticDecodingTables);
         }
 

tests/ImageSharp.Tests/Formats/Jpg/JpegDecoderTests.cs

@@ -10,6 +10,7 @@
 using SixLabors.ImageSharp.Tests.TestUtilities.ImageComparison;
 using SixLabors.ImageSharp.Tests.TestUtilities.ReferenceCodecs;
 using Xunit.Abstractions;
+using static SixLabors.ImageSharp.Tests.TestImages.Jpeg;
 
 // ReSharper disable InconsistentNaming
 namespace SixLabors.ImageSharp.Tests.Formats.Jpg;
@@ -300,6 +301,124 @@ public void Issue2136_DecodeWorks<TPixel>(TestImageProvider<TPixel> provider)
         image.CompareToOriginal(provider);
     }
 
+    /// <summary>
+    /// Uses the JPEG defaults when an arithmetic conditioning table is not declared.
+    /// </summary>
+    /// <param name="originalFile">The input with explicit conditioning tables.</param>
+    /// <param name="defaultFile">The input that omits default conditioning tables.</param>
+    [Theory]
+    [InlineData(MarkerDirectory + "arithmetic.jpg", MarkerDirectory + "arithmetic-defaults.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic-progressive.jpg", MarkerDirectory + "arithmetic-progressive-defaults.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic-progressive.jpg", MarkerDirectory + "arithmetic-progressive-default-ac.jpg")]
+    public void Decode_ArithmeticConditioningDefaults_MatchExplicitTables(string originalFile, string defaultFile)
+    {
+        // These stored files contain identical coded scans. Libjpeg produces identical pixels with or without
+        // the default DAC entries. Omission is valid JPEG syntax and does not require image-data recovery.
+        using Image<Rgba32> expected = Image.Load<Rgba32>(TestFile.Create(originalFile).Bytes);
+        byte[] data = TestFile.Create(defaultFile).Bytes;
+
+        foreach (SegmentIntegrityHandling integrity in Enum.GetValues<SegmentIntegrityHandling>())
+        {
+            DecoderOptions options = new() { SegmentIntegrityHandling = integrity };
+            using Image<Rgba32> actual = Image.Load<Rgba32>(options, data);
+
+            ImageComparer.Exact.VerifySimilarity(expected, actual);
+        }
+    }
+
+    /// <summary>
+    /// Checks restart recovery against clean stored files whose pixels also match in libjpeg.
+    /// </summary>
+    /// <param name="originalFile">The clean JPEG input.</param>
+    /// <param name="malformedFile">The input with extra markers before restart markers.</param>
+    [Theory]
+    [InlineData(MarkerDirectory + "baseline.jpg", MarkerDirectory + "baseline-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "baseline.jpg", MarkerDirectory + "baseline-tem.jpg")]
+    [InlineData(MarkerDirectory + "baseline.jpg", MarkerDirectory + "baseline-reserved-bf.jpg")]
+    [InlineData(MarkerDirectory + "baseline.jpg", MarkerDirectory + "baseline-reserved-chain.jpg")]
+    [InlineData(MarkerDirectory + "noninterleaved.jpg", MarkerDirectory + "noninterleaved-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "progressive.jpg", MarkerDirectory + "progressive-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic.jpg", MarkerDirectory + "arithmetic-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic-progressive.jpg", MarkerDirectory + "arithmetic-progressive-reserved-12.jpg")]
+    public void Decode_ReservedMarkerBeforeRestartMarkers_RecoversWithIgnoreImageData(string originalFile, string malformedFile)
+    {
+        // These stored files have complete restart intervals. Libjpeg warns about the extra markers,
+        // then finds the restart marker and produces the same pixels as the clean file.
+        // The arithmetic inputs declare their conditioning tables to isolate restart recovery.
+        DecoderOptions options = new() { SegmentIntegrityHandling = SegmentIntegrityHandling.IgnoreImageData };
+        using Image<Rgba32> expected = Image.Load<Rgba32>(TestFile.Create(originalFile).Bytes);
+        using Image<Rgba32> actual = Image.Load<Rgba32>(options, TestFile.Create(malformedFile).Bytes);
+
+        ImageComparer.Exact.VerifySimilarity(expected, actual);
+    }
+
+    /// <summary>
+    /// Rejects image-data errors when the selected policy permits only ancillary recovery or no recovery.
+    /// </summary>
+    /// <param name="fileName">The input with reserved markers before restart markers.</param>
+    [Theory]
+    [InlineData(MarkerDirectory + "baseline-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "noninterleaved-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "progressive-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic-reserved-12.jpg")]
+    [InlineData(MarkerDirectory + "arithmetic-progressive-reserved-12.jpg")]
+    public void Decode_ReservedMarkerBeforeRestartMarkers_RespectsIntegrity(string fileName)
+    {
+        byte[] data = TestFile.Create(fileName).Bytes;
+        DecoderOptions strict = new() { SegmentIntegrityHandling = SegmentIntegrityHandling.Strict };
+        DecoderOptions ancillary = new() { SegmentIntegrityHandling = SegmentIntegrityHandling.IgnoreAncillary };
+
+        Assert.Throws<InvalidImageContentException>(() =>
+        {
+            using Image<Rgba32> image = Image.Load<Rgba32>(strict, data);
+        });
+
+        Assert.Throws<InvalidImageContentException>(() =>
+        {
+            using Image<Rgba32> image = Image.Load<Rgba32>(ancillary, data);
+        });
+    }
+
+    /// <summary>
+    /// Accepts a standalone TEM marker between header segments under every integrity policy.
+    /// </summary>
+    [Fact]
+    public void Decode_TemMarkerBetweenHeaderSegments_IsSkipped()
+    {
+        byte[] data = TestFile.Create(MarkerDirectory + "header-tem.jpg").Bytes;
+        using Image<Rgba32> expected = Image.Load<Rgba32>(TestFile.Create(MarkerDirectory + "baseline.jpg").Bytes);
+
+        // TEM has no length field. The following bytes belong to the next marker, not to TEM.
+        foreach (SegmentIntegrityHandling integrity in Enum.GetValues<SegmentIntegrityHandling>())
+        {
+            DecoderOptions options = new() { SegmentIntegrityHandling = integrity };
+            Assert.Equal(new Size(32, 24), Image.Identify(options, data).Size);
+
+            using Image<Rgba32> actual = Image.Load<Rgba32>(options, data);
+            ImageComparer.Exact.VerifySimilarity(expected, actual);
+        }
+    }
+
+    /// <summary>
+    /// Keeps required header lengths valid after a standalone TEM marker.
+    /// </summary>
+    [Fact]
+    public void Decode_InvalidQuantizationSegmentAfterTem_RemainsFatal()
+    {
+        // TEM precedes DQT in this stored file. DQT declares more bytes than the file contains.
+        // No image data has been read, so there is no safe point for recovery.
+        byte[] data = TestFile.Create(MarkerDirectory + "header-tem-invalid-dqt-length.jpg").Bytes;
+
+        foreach (SegmentIntegrityHandling integrity in Enum.GetValues<SegmentIntegrityHandling>())
+        {
+            DecoderOptions options = new() { SegmentIntegrityHandling = integrity };
+            Assert.Throws<InvalidImageContentException>(() =>
+            {
+                using Image<Rgba32> image = Image.Load<Rgba32>(options, data);
+            });
+        }
+    }
+
     // https://github.com/SixLabors/ImageSharp/issues/2315
     // https://github.com/SixLabors/ImageSharp/issues/2334
     [Theory]