From 4f76e3b49c60c14857879ed9e4eecacb98edfbbf Mon Sep 17 00:00:00 2001
From: James Jackson-South <james_south@hotmail.com>
Date: Tue, 4 Jun 2019 15:41:45 +1000
Subject: [PATCH] Prevent zigzag overflow. Fix #922

---
 .../Formats/Jpeg/Components/ZigZag.cs         | 45 +++++++++++--------
 .../Formats/Jpg/ZigZagTests.cs                | 44 ++++++++++++++++++
 2 files changed, 70 insertions(+), 19 deletions(-)
 create mode 100644 tests/ImageSharp.Tests/Formats/Jpg/ZigZagTests.cs

diff --git a/src/ImageSharp/Formats/Jpeg/Components/ZigZag.cs b/src/ImageSharp/Formats/Jpeg/Components/ZigZag.cs
index a3701f2c1..250e98c02 100644
--- a/src/ImageSharp/Formats/Jpeg/Components/ZigZag.cs
+++ b/src/ImageSharp/Formats/Jpeg/Components/ZigZag.cs
@@ -1,4 +1,4 @@
-// Copyright (c) Six Labors and contributors.
+// Copyright (c) Six Labors and contributors.
 // Licensed under the Apache License, Version 2.0.
 
 using System;
@@ -19,30 +19,37 @@ namespace SixLabors.ImageSharp.Formats.Jpeg.Components
         /// <summary>
         /// Copy of <see cref="Unzig"/> in a value type
         /// </summary>
-        public fixed byte Data[64];
+        public fixed byte Data[64 + 16];
 
         /// <summary>
+        /// <para>
         /// Unzig maps from the zigzag ordering to the natural ordering. For example,
         /// unzig[3] is the column and row of the fourth element in zigzag order. The
         /// value is 16, which means first column (16%8 == 0) and third row (16/8 == 2).
+        /// </para>
+        /// <para>
+        /// When reading corrupted data, the Huffman decoders could attempt
+        /// to reference an entry beyond the end of this array (if the decoded
+        /// zero run length reaches past the end of the block).  To prevent
+        /// wild stores without adding an inner-loop test, we put some extra
+        /// "63"s after the real entries.  This will cause the extra coefficient
+        /// to be stored in location 63 of the block, not somewhere random.
+        /// The worst case would be a run-length of 15, which means we need 16
+        /// fake entries.
+        /// </para>
         /// </summary>
         private static readonly byte[] Unzig =
         {
-            0,
-            1, 8,
-            16, 9, 2,
-            3, 10, 17, 24,
-            32, 25, 18, 11, 4,
-            5, 12, 19, 26, 33, 40,
-            48, 41, 34, 27, 20, 13, 6,
-            7, 14, 21, 28, 35, 42, 49, 56,
-            57, 50, 43, 36, 29, 22, 15,
-            23, 30, 37, 44, 51, 58,
-            59, 52, 45, 38, 31,
-            39, 46, 53, 60,
-            61, 54, 47,
-            55, 62,
-            63
+            0,  1,  8, 16,  9,  2,  3, 10,
+            17, 24, 32, 25, 18, 11,  4,  5,
+            12, 19, 26, 33, 40, 48, 41, 34,
+            27, 20, 13,  6,  7, 14, 21, 28,
+            35, 42, 49, 56, 57, 50, 43, 36,
+            29, 22, 15, 23, 30, 37, 44, 51,
+            58, 59, 52, 45, 38, 31, 39, 46,
+            53, 60, 61, 54, 47, 55, 62, 63,
+            63, 63, 63, 63, 63, 63, 63, 63, // Extra entries for safety in decoder
+            63, 63, 63, 63, 63, 63, 63, 63
         };
 
         /// <summary>
@@ -68,7 +75,7 @@ namespace SixLabors.ImageSharp.Formats.Jpeg.Components
         {
             ZigZag result = default;
             byte* unzigPtr = result.Data;
-            Marshal.Copy(Unzig, 0, (IntPtr)unzigPtr, 64);
+            Marshal.Copy(Unzig, 0, (IntPtr)unzigPtr, 64 + 16);
             return result;
         }
 
@@ -87,4 +94,4 @@ namespace SixLabors.ImageSharp.Formats.Jpeg.Components
             return result;
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/ImageSharp.Tests/Formats/Jpg/ZigZagTests.cs b/tests/ImageSharp.Tests/Formats/Jpg/ZigZagTests.cs
new file mode 100644
index 000000000..e59584991
--- /dev/null
+++ b/tests/ImageSharp.Tests/Formats/Jpg/ZigZagTests.cs
@@ -0,0 +1,44 @@
+// Copyright (c) Six Labors and contributors.
+// Licensed under the Apache License, Version 2.0.
+
+using SixLabors.ImageSharp.Formats.Jpeg.Components;
+using Xunit;
+
+namespace SixLabors.ImageSharp.Tests.Formats.Jpg
+{
+    public class ZigZagTests
+    {
+        [Fact]
+        public void ZigZagCanHandleAllPossibleCoefficients()
+        {
+            // Mimic the behaviour of the huffman scan decoder using all possible byte values
+            short[] block = new short[64];
+            var zigzag = ZigZag.CreateUnzigTable();
+
+            for (int h = 0; h < 255; h++)
+            {
+                for (int i = 1; i < 64; i++)
+                {
+                    int s = h;
+                    int r = s >> 4;
+                    s &= 15;
+
+                    if (s != 0)
+                    {
+                        i += r;
+                        block[zigzag[i++]] = (short)s;
+                    }
+                    else
+                    {
+                        if (r == 0)
+                        {
+                            break;
+                        }
+
+                        i += 16;
+                    }
+                }
+            }
+        }
+    }
+}