Merge pull request #2110 from SixLabors/bp/png-iccp

brianpopow · web-flow · commit 5634228f5922 · 2022-05-12T09:42:18.000+02:00
Preserve color profile when encoding PNG images
diff --git a/src/ImageSharp/Formats/Png/PngDecoderCore.cs b/src/ImageSharp/Formats/Png/PngDecoderCore.cs
@@ -19,6 +19,7 @@
 using SixLabors.ImageSharp.Memory;
 using SixLabors.ImageSharp.Metadata;
 using SixLabors.ImageSharp.Metadata.Profiles.Exif;
+using SixLabors.ImageSharp.Metadata.Profiles.Icc;
 using SixLabors.ImageSharp.Metadata.Profiles.Xmp;
 using SixLabors.ImageSharp.PixelFormats;
 
@@ -205,6 +206,9 @@ public Image<TPixel> Decode<TPixel>(BufferedReadStream stream, CancellationToken
                                     this.MergeOrSetExifProfile(metadata, new ExifProfile(exifData), replaceExistingKeys: true);
                                 }
 
+                                break;
+                            case PngChunkType.EmbeddedColorProfile:
+                                this.ReadColorProfileChunk(metadata, chunk.Data.GetSpan());
                                 break;
                             case PngChunkType.End:
                                 goto EOF;
@@ -1174,6 +1178,76 @@ private bool TryReadLegacyExifTextChunk(ImageMetadata metadata, string data)
             return true;
         }
 
+        /// <summary>
+        /// Reads the color profile chunk. The data is stored similar to the zTXt chunk.
+        /// </summary>
+        /// <param name="metadata">The metadata.</param>
+        /// <param name="data">The bytes containing the profile.</param>
+        private void ReadColorProfileChunk(ImageMetadata metadata, ReadOnlySpan<byte> data)
+        {
+            int zeroIndex = data.IndexOf((byte)0);
+            if (zeroIndex is < PngConstants.MinTextKeywordLength or > PngConstants.MaxTextKeywordLength)
+            {
+                return;
+            }
+
+            byte compressionMethod = data[zeroIndex + 1];
+            if (compressionMethod != 0)
+            {
+                // Only compression method 0 is supported (zlib datastream with deflate compression).
+                return;
+            }
+
+            ReadOnlySpan<byte> keywordBytes = data.Slice(0, zeroIndex);
+            if (!this.TryReadTextKeyword(keywordBytes, out string name))
+            {
+                return;
+            }
+
+            ReadOnlySpan<byte> compressedData = data.Slice(zeroIndex + 2);
+
+            if (this.TryUncompressZlibData(compressedData, out byte[] iccpProfileBytes))
+            {
+                metadata.IccProfile = new IccProfile(iccpProfileBytes);
+            }
+        }
+
+        /// <summary>
+        /// Tries to un-compress zlib compressed data.
+        /// </summary>
+        /// <param name="compressedData">The compressed data.</param>
+        /// <param name="uncompressedBytesArray">The uncompressed bytes array.</param>
+        /// <returns>True, if de-compressing was successful.</returns>
+        private unsafe bool TryUncompressZlibData(ReadOnlySpan<byte> compressedData, out byte[] uncompressedBytesArray)
+        {
+            fixed (byte* compressedDataBase = compressedData)
+            {
+                using (IMemoryOwner<byte> destBuffer = this.memoryAllocator.Allocate<byte>(this.Configuration.StreamProcessingBufferSize))
+                using (var memoryStreamOutput = new MemoryStream(compressedData.Length))
+                using (var memoryStreamInput = new UnmanagedMemoryStream(compressedDataBase, compressedData.Length))
+                using (var bufferedStream = new BufferedReadStream(this.Configuration, memoryStreamInput))
+                using (var inflateStream = new ZlibInflateStream(bufferedStream))
+                {
+                    Span<byte> destUncompressedData = destBuffer.GetSpan();
+                    if (!inflateStream.AllocateNewBytes(compressedData.Length, false))
+                    {
+                        uncompressedBytesArray = Array.Empty<byte>();
+                        return false;
+                    }
+
+                    int bytesRead = inflateStream.CompressedStream.Read(destUncompressedData, 0, destUncompressedData.Length);
+                    while (bytesRead != 0)
+                    {
+                        memoryStreamOutput.Write(destUncompressedData.Slice(0, bytesRead));
+                        bytesRead = inflateStream.CompressedStream.Read(destUncompressedData, 0, destUncompressedData.Length);
+                    }
+
+                    uncompressedBytesArray = memoryStreamOutput.ToArray();
+                    return true;
+                }
+            }
+        }
+
         /// <summary>
         /// Compares two ReadOnlySpan&lt;char&gt;s in a case-insensitive method.
         /// This is only needed because older frameworks are missing the extension method.
@@ -1306,7 +1380,7 @@ private void ReadInternationalTextChunk(ImageMetadata metadata, ReadOnlySpan<byt
             }
             else if (this.IsXmpTextData(keywordBytes))
             {
-                XmpProfile xmpProfile = new XmpProfile(data.Slice(dataStartIdx).ToArray());
+                var xmpProfile = new XmpProfile(data.Slice(dataStartIdx).ToArray());
                 metadata.XmpProfile = xmpProfile;
             }
             else
@@ -1325,29 +1399,14 @@ private void ReadInternationalTextChunk(ImageMetadata metadata, ReadOnlySpan<byt
         /// <returns>The <see cref="bool"/>.</returns>
         private bool TryUncompressTextData(ReadOnlySpan<byte> compressedData, Encoding encoding, out string value)
         {
-            using (var memoryStream = new MemoryStream(compressedData.ToArray()))
-            using (var bufferedStream = new BufferedReadStream(this.Configuration, memoryStream))
-            using (var inflateStream = new ZlibInflateStream(bufferedStream))
+            if (this.TryUncompressZlibData(compressedData, out byte[] uncompressedData))
             {
-                if (!inflateStream.AllocateNewBytes(compressedData.Length, false))
-                {
-                    value = null;
-                    return false;
-                }
-
-                var uncompressedBytes = new List<byte>();
-
-                // Note: this uses a buffer which is only 4 bytes long to read the stream, maybe allocating a larger buffer makes sense here.
-                int bytesRead = inflateStream.CompressedStream.Read(this.buffer, 0, this.buffer.Length);
-                while (bytesRead != 0)
-                {
-                    uncompressedBytes.AddRange(this.buffer.AsSpan(0, bytesRead).ToArray());
-                    bytesRead = inflateStream.CompressedStream.Read(this.buffer, 0, this.buffer.Length);
-                }
-
-                value = encoding.GetString(uncompressedBytes.ToArray());
+                value = encoding.GetString(uncompressedData);
                 return true;
             }
+
+            value = null;
+            return false;
         }
 
         /// <summary>
diff --git a/src/ImageSharp/Formats/Png/PngEncoderCore.cs b/src/ImageSharp/Formats/Png/PngEncoderCore.cs
@@ -87,6 +87,11 @@ internal sealed class PngEncoderCore : IImageEncoderInternals, IDisposable
         /// </summary>
         private IMemoryOwner<byte> currentScanline;
 
+        /// <summary>
+        /// The color profile name.
+        /// </summary>
+        private const string ColorProfileName = "ICC Profile";
+
         /// <summary>
         /// Initializes a new instance of the <see cref="PngEncoderCore" /> class.
         /// </summary>
@@ -134,6 +139,7 @@ public void Encode<TPixel>(Image<TPixel> image, Stream stream, CancellationToken
 
             this.WriteHeaderChunk(stream);
             this.WriteGammaChunk(stream);
+            this.WriteColorProfileChunk(stream, metadata);
             this.WritePaletteChunk(stream, quantized);
             this.WriteTransparencyChunk(stream, pngMetadata);
             this.WritePhysicalChunk(stream, metadata);
@@ -656,7 +662,7 @@ private void WriteExifChunk(Stream stream, ImageMetadata meta)
         }
 
         /// <summary>
-        /// Writes an iTXT chunk, containing the XMP metdata to the stream, if such profile is present in the metadata.
+        /// Writes an iTXT chunk, containing the XMP metadata to the stream, if such profile is present in the metadata.
         /// </summary>
         /// <param name="stream">The <see cref="Stream"/> containing image data.</param>
         /// <param name="meta">The image metadata.</param>
@@ -673,7 +679,7 @@ private void WriteXmpChunk(Stream stream, ImageMetadata meta)
                 return;
             }
 
-            var xmpData = meta.XmpProfile.Data;
+            byte[] xmpData = meta.XmpProfile.Data;
 
             if (xmpData.Length == 0)
             {
@@ -687,19 +693,49 @@ private void WriteXmpChunk(Stream stream, ImageMetadata meta)
                 PngConstants.XmpKeyword.CopyTo(payload);
                 int bytesWritten = PngConstants.XmpKeyword.Length;
 
-                // Write the iTxt header (all zeros in this case)
-                payload[bytesWritten++] = 0;
-                payload[bytesWritten++] = 0;
-                payload[bytesWritten++] = 0;
-                payload[bytesWritten++] = 0;
-                payload[bytesWritten++] = 0;
+                // Write the iTxt header (all zeros in this case).
+                Span<byte> iTxtHeader = payload.Slice(bytesWritten);
+                iTxtHeader[4] = 0;
+                iTxtHeader[3] = 0;
+                iTxtHeader[2] = 0;
+                iTxtHeader[1] = 0;
+                iTxtHeader[0] = 0;
+                bytesWritten += 5;
 
-                // And the XMP data itself
+                // And the XMP data itself.
                 xmpData.CopyTo(payload.Slice(bytesWritten));
                 this.WriteChunk(stream, PngChunkType.InternationalText, payload);
             }
         }
 
+        /// <summary>
+        /// Writes the color profile chunk.
+        /// </summary>
+        /// <param name="stream">The stream to write to.</param>
+        /// <param name="metaData">The image meta data.</param>
+        private void WriteColorProfileChunk(Stream stream, ImageMetadata metaData)
+        {
+            if (metaData.IccProfile is null)
+            {
+                return;
+            }
+
+            byte[] iccProfileBytes = metaData.IccProfile.ToByteArray();
+
+            byte[] compressedData = this.GetZlibCompressedBytes(iccProfileBytes);
+            int payloadLength = ColorProfileName.Length + compressedData.Length + 2;
+            using (IMemoryOwner<byte> owner = this.memoryAllocator.Allocate<byte>(payloadLength))
+            {
+                Span<byte> outputBytes = owner.GetSpan();
+                PngConstants.Encoding.GetBytes(ColorProfileName).CopyTo(outputBytes);
+                int bytesWritten = ColorProfileName.Length;
+                outputBytes[bytesWritten++] = 0; // Null separator.
+                outputBytes[bytesWritten++] = 0; // Compression.
+                compressedData.CopyTo(outputBytes.Slice(bytesWritten));
+                this.WriteChunk(stream, PngChunkType.EmbeddedColorProfile, outputBytes);
+            }
+        }
+
         /// <summary>
         /// Writes a text chunk to the stream. Can be either a tTXt, iTXt or zTXt chunk,
         /// depending whether the text contains any latin characters or should be compressed.
@@ -727,13 +763,12 @@ private void WriteTextChunks(Stream stream, PngMetadata meta)
                     }
                 }
 
-                if (hasUnicodeCharacters || (!string.IsNullOrWhiteSpace(textData.LanguageTag) ||
-                                             !string.IsNullOrWhiteSpace(textData.TranslatedKeyword)))
+                if (hasUnicodeCharacters || (!string.IsNullOrWhiteSpace(textData.LanguageTag) || !string.IsNullOrWhiteSpace(textData.TranslatedKeyword)))
                 {
                     // Write iTXt chunk.
                     byte[] keywordBytes = PngConstants.Encoding.GetBytes(textData.Keyword);
                     byte[] textBytes = textData.Value.Length > this.options.TextCompressionThreshold
-                        ? this.GetCompressedTextBytes(PngConstants.TranslatedEncoding.GetBytes(textData.Value))
+                        ? this.GetZlibCompressedBytes(PngConstants.TranslatedEncoding.GetBytes(textData.Value))
                         : PngConstants.TranslatedEncoding.GetBytes(textData.Value);
 
                     byte[] translatedKeyword = PngConstants.TranslatedEncoding.GetBytes(textData.TranslatedKeyword);
@@ -772,18 +807,17 @@ private void WriteTextChunks(Stream stream, PngMetadata meta)
                     if (textData.Value.Length > this.options.TextCompressionThreshold)
                     {
                         // Write zTXt chunk.
-                        byte[] compressedData =
-                            this.GetCompressedTextBytes(PngConstants.Encoding.GetBytes(textData.Value));
+                        byte[] compressedData = this.GetZlibCompressedBytes(PngConstants.Encoding.GetBytes(textData.Value));
                         int payloadLength = textData.Keyword.Length + compressedData.Length + 2;
                         using (IMemoryOwner<byte> owner = this.memoryAllocator.Allocate<byte>(payloadLength))
                         {
                             Span<byte> outputBytes = owner.GetSpan();
                             PngConstants.Encoding.GetBytes(textData.Keyword).CopyTo(outputBytes);
                             int bytesWritten = textData.Keyword.Length;
-                            outputBytes[bytesWritten++] = 0;
-                            outputBytes[bytesWritten++] = 0;
+                            outputBytes[bytesWritten++] = 0; // Null separator.
+                            outputBytes[bytesWritten++] = 0; // Compression.
                             compressedData.CopyTo(outputBytes.Slice(bytesWritten));
-                            this.WriteChunk(stream, PngChunkType.CompressedText, outputBytes.ToArray());
+                            this.WriteChunk(stream, PngChunkType.CompressedText, outputBytes);
                         }
                     }
                     else
@@ -796,9 +830,8 @@ private void WriteTextChunks(Stream stream, PngMetadata meta)
                             PngConstants.Encoding.GetBytes(textData.Keyword).CopyTo(outputBytes);
                             int bytesWritten = textData.Keyword.Length;
                             outputBytes[bytesWritten++] = 0;
-                            PngConstants.Encoding.GetBytes(textData.Value)
-                                .CopyTo(outputBytes.Slice(bytesWritten));
-                            this.WriteChunk(stream, PngChunkType.Text, outputBytes.ToArray());
+                            PngConstants.Encoding.GetBytes(textData.Value).CopyTo(outputBytes.Slice(bytesWritten));
+                            this.WriteChunk(stream, PngChunkType.Text, outputBytes);
                         }
                     }
                 }
@@ -808,15 +841,15 @@ private void WriteTextChunks(Stream stream, PngMetadata meta)
         /// <summary>
         /// Compresses a given text using Zlib compression.
         /// </summary>
-        /// <param name="textBytes">The text bytes to compress.</param>
-        /// <returns>The compressed text byte array.</returns>
-        private byte[] GetCompressedTextBytes(byte[] textBytes)
+        /// <param name="dataBytes">The bytes to compress.</param>
+        /// <returns>The compressed byte array.</returns>
+        private byte[] GetZlibCompressedBytes(byte[] dataBytes)
         {
             using (var memoryStream = new MemoryStream())
             {
                 using (var deflateStream = new ZlibDeflateStream(this.memoryAllocator, memoryStream, this.options.CompressionLevel))
                 {
-                    deflateStream.Write(textBytes);
+                    deflateStream.Write(dataBytes);
                 }
 
                 return memoryStream.ToArray();
diff --git a/tests/ImageSharp.Tests/Formats/Png/PngEncoderTests.Chunks.cs b/tests/ImageSharp.Tests/Formats/Png/PngEncoderTests.Chunks.cs
@@ -302,6 +302,7 @@ public void ExcludeFilter_WithNone_DoesNotExcludeChunks()
             {
                 PngChunkType.Header,
                 PngChunkType.Gamma,
+                PngChunkType.EmbeddedColorProfile,
                 PngChunkType.Palette,
                 PngChunkType.InternationalText,
                 PngChunkType.Text,
diff --git a/tests/ImageSharp.Tests/Formats/Png/PngMetadataTests.cs b/tests/ImageSharp.Tests/Formats/Png/PngMetadataTests.cs
@@ -17,7 +17,7 @@ namespace SixLabors.ImageSharp.Tests.Formats.Png
     public class PngMetadataTests
     {
         public static readonly TheoryData<string, int, int, PixelResolutionUnit> RatioFiles =
-            new TheoryData<string, int, int, PixelResolutionUnit>
+            new()
             {
                 { TestImages.Png.Splash, 11810, 11810, PixelResolutionUnit.PixelsPerMeter },
                 { TestImages.Png.Ratio1x4, 1, 4, PixelResolutionUnit.AspectRatio },
@@ -222,6 +222,33 @@ public void Decode_VerifyRatio(string imagePath, int xResolution, int yResolutio
             }
         }
 
+        [Theory]
+        [WithFile(TestImages.Png.PngWithMetadata, PixelTypes.Rgba32)]
+        public void Encode_PreservesColorProfile<TPixel>(TestImageProvider<TPixel> provider)
+            where TPixel : unmanaged, IPixel<TPixel>
+        {
+            using (Image<TPixel> input = provider.GetImage(new PngDecoder()))
+            {
+                ImageSharp.Metadata.Profiles.Icc.IccProfile expectedProfile = input.Metadata.IccProfile;
+                byte[] expectedProfileBytes = expectedProfile.ToByteArray();
+
+                using (var memStream = new MemoryStream())
+                {
+                    input.Save(memStream, new PngEncoder());
+
+                    memStream.Position = 0;
+                    using (var output = Image.Load<Rgba32>(memStream))
+                    {
+                        ImageSharp.Metadata.Profiles.Icc.IccProfile actualProfile = output.Metadata.IccProfile;
+                        byte[] actualProfileBytes = actualProfile.ToByteArray();
+
+                        Assert.NotNull(actualProfile);
+                        Assert.Equal(expectedProfileBytes, actualProfileBytes);
+                    }
+                }
+            }
+        }
+
         [Theory]
         [MemberData(nameof(RatioFiles))]
         public void Identify_VerifyRatio(string imagePath, int xResolution, int yResolution, PixelResolutionUnit resolutionUnit)

Original file line number	Diff line number	Diff line change
`@@ -302,6 +302,7 @@ public void ExcludeFilter_WithNone_DoesNotExcludeChunks()`
`302`	`302`	`{`
`303`	`303`	`PngChunkType.Header,`
`304`	`304`	`PngChunkType.Gamma,`
	`305`	`+ PngChunkType.EmbeddedColorProfile,`
`305`	`306`	`PngChunkType.Palette,`
`306`	`307`	`PngChunkType.InternationalText,`
`307`	`308`	`PngChunkType.Text,`