Skip to content

Commit

Permalink
Fix 925 (#929)
Browse files Browse the repository at this point in the history
* Prevent overflow

* Cleanup huffman table

* Search for RST markers.

* Fix Benchmarks project
  • Loading branch information
JimBobSquarePants authored Jun 14, 2019
1 parent 6fcfe09 commit 7a484f7
Show file tree
Hide file tree
Showing 7 changed files with 103 additions and 78 deletions.
92 changes: 63 additions & 29 deletions src/ImageSharp/Formats/Jpeg/Components/Decoder/HuffmanScanBuffer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ internal struct HuffmanScanBuffer
private ulong data;

// The number of valid bits left to read in the buffer.
private int remain;
private int remainingBits;

// Whether there is no more good data to pull from the stream for the current mcu.
private bool badData;
Expand All @@ -26,10 +26,9 @@ public HuffmanScanBuffer(DoubleBufferedStreamReader stream)
{
this.stream = stream;
this.data = 0ul;
this.remain = 0;
this.remainingBits = 0;
this.Marker = JpegConstants.Markers.XFF;
this.MarkerPosition = 0;
this.BadMarker = false;
this.badData = false;
this.NoData = false;
}
Expand All @@ -44,11 +43,6 @@ public HuffmanScanBuffer(DoubleBufferedStreamReader stream)
/// </summary>
public long MarkerPosition { get; private set; }

/// <summary>
/// Gets a value indicating whether a bad marker has been detected, I.E. One that is not between RST0 and RST7
/// </summary>
public bool BadMarker { get; private set; }

/// <summary>
/// Gets a value indicating whether to continue reading the input stream.
/// </summary>
Expand All @@ -57,7 +51,7 @@ public HuffmanScanBuffer(DoubleBufferedStreamReader stream)
[MethodImpl(InliningOptions.ShortMethod)]
public void CheckBits()
{
if (this.remain < 16)
if (this.remainingBits < 16)
{
this.FillBuffer();
}
Expand All @@ -67,27 +61,31 @@ public void CheckBits()
public void Reset()
{
this.data = 0ul;
this.remain = 0;
this.remainingBits = 0;
this.Marker = JpegConstants.Markers.XFF;
this.MarkerPosition = 0;
this.BadMarker = false;
this.badData = false;
this.NoData = false;
}

/// <summary>
/// Whether a RST marker has been detected, I.E. One that is between RST0 and RST7
/// </summary>
[MethodImpl(InliningOptions.ShortMethod)]
public bool HasRestart()
{
byte m = this.Marker;
return m >= JpegConstants.Markers.RST0 && m <= JpegConstants.Markers.RST7;
}
public bool HasRestartMarker() => HasRestart(this.Marker);

/// <summary>
/// Whether a bad marker has been detected, I.E. One that is not between RST0 and RST7
/// </summary>
[MethodImpl(InliningOptions.ShortMethod)]
public bool HasBadMarker() => this.Marker != JpegConstants.Markers.XFF && !this.HasRestartMarker();

[MethodImpl(InliningOptions.ShortMethod)]
public void FillBuffer()
{
// Attempt to load at least the minimum number of required bits into the buffer.
// We fail to do so only if we hit a marker or reach the end of the input stream.
this.remain += 48;
this.remainingBits += 48;
this.data = (this.data << 48) | this.GetBytes();
}

Expand All @@ -101,17 +99,17 @@ public unsafe int DecodeHuffman(ref HuffmanTable h)

if (size == JpegConstants.Huffman.SlowBits)
{
ulong x = this.data << (JpegConstants.Huffman.RegisterSize - this.remain);
ulong x = this.data << (JpegConstants.Huffman.RegisterSize - this.remainingBits);
while (x > h.MaxCode[size])
{
size++;
}

v = (int)(x >> (JpegConstants.Huffman.RegisterSize - size));
symbol = h.Values[h.ValOffset[size] + v];
symbol = h.Values[(h.ValOffset[size] + v) & 0xFF];
}

this.remain -= size;
this.remainingBits -= size;

return symbol;
}
Expand All @@ -124,10 +122,14 @@ public int Receive(int nbits)
}

[MethodImpl(InliningOptions.ShortMethod)]
public int GetBits(int nbits) => (int)ExtractBits(this.data, this.remain -= nbits, nbits);
private static bool HasRestart(byte marker)
=> marker >= JpegConstants.Markers.RST0 && marker <= JpegConstants.Markers.RST7;

[MethodImpl(InliningOptions.ShortMethod)]
public int PeekBits(int nbits) => (int)ExtractBits(this.data, this.remain - nbits, nbits);
public int GetBits(int nbits) => (int)ExtractBits(this.data, this.remainingBits -= nbits, nbits);

[MethodImpl(InliningOptions.ShortMethod)]
public int PeekBits(int nbits) => (int)ExtractBits(this.data, this.remainingBits - nbits, nbits);

[MethodImpl(InliningOptions.ShortMethod)]
private static ulong ExtractBits(ulong value, int offset, int size) => (value >> offset) & (ulong)((1 << size) - 1);
Expand All @@ -149,22 +151,18 @@ private ulong GetBytes()
int c = this.ReadStream();
while (c == JpegConstants.Markers.XFF)
{
// Loop here to discard any padding FF's on terminating marker,
// Loop here to discard any padding FF bytes on terminating marker,
// so that we can save a valid marker value.
c = this.ReadStream();
}

// We accept multiple FF's followed by a 0 as meaning a single FF data byte.
// We accept multiple FF bytes followed by a 0 as meaning a single FF data byte.
// This data pattern is not valid according to the standard.
if (c != 0)
{
this.Marker = (byte)c;
this.badData = true;
if (!this.HasRestart())
{
this.MarkerPosition = this.stream.Position - 2;
this.BadMarker = true;
}
this.MarkerPosition = this.stream.Position - 2;
}
}

Expand All @@ -174,6 +172,42 @@ private ulong GetBytes()
return temp;
}

[MethodImpl(InliningOptions.ShortMethod)]
public bool FindNextMarker()
{
while (true)
{
int b = this.stream.ReadByte();
if (b == -1)
{
return false;
}

// Found a marker.
if (b == JpegConstants.Markers.XFF)
{
while (b == JpegConstants.Markers.XFF)
{
// Loop here to discard any padding FF bytes on terminating marker.
b = this.stream.ReadByte();
if (b == -1)
{
return false;
}
}

// Found a valid marker. Exit loop
if (b != 0)
{
this.Marker = (byte)b;
this.badData = true;
this.MarkerPosition = this.stream.Position - 2;
return true;
}
}
}
}

[MethodImpl(InliningOptions.ShortMethod)]
private int ReadStream()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ public void ParseEntropyCodedData()
this.ParseProgressiveData();
}

if (this.scanBuffer.BadMarker)
if (this.scanBuffer.HasBadMarker())
{
this.stream.Position = this.scanBuffer.MarkerPosition;
}
Expand Down Expand Up @@ -684,15 +684,23 @@ private bool HandleRestart()
{
if (this.restartInterval > 0 && (--this.todo) == 0)
{
if (this.scanBuffer.Marker == JpegConstants.Markers.XFF)
{
if (!this.scanBuffer.FindNextMarker())
{
return false;
}
}

this.todo = this.restartInterval;

if (this.scanBuffer.HasRestart())
if (this.scanBuffer.HasRestartMarker())
{
this.Reset();
return true;
}

if (this.scanBuffer.Marker != JpegConstants.Markers.XFF)
if (this.scanBuffer.HasBadMarker())
{
this.stream.Position = this.scanBuffer.MarkerPosition;
this.Reset();
Expand Down
48 changes: 21 additions & 27 deletions src/ImageSharp/Formats/Jpeg/Components/Decoder/HuffmanTable.cs
Original file line number Diff line number Diff line change
@@ -1,20 +1,18 @@
// Copyright (c) Six Labors and contributors.
// Copyright (c) Six Labors and contributors.
// Licensed under the Apache License, Version 2.0.

using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using SixLabors.Memory;

namespace SixLabors.ImageSharp.Formats.Jpeg.Components.Decoder
{
/// <summary>
/// Represents a Huffman coding table containing basic coding data plus tables for accellerated computation.
/// Represents a Huffman coding table containing basic coding data plus tables for accelerated computation.
/// </summary>
[StructLayout(LayoutKind.Sequential)]
internal unsafe struct HuffmanTable
{
private readonly MemoryAllocator memoryAllocator;
private bool isConfigured;

/// <summary>
Expand Down Expand Up @@ -60,13 +58,11 @@ internal unsafe struct HuffmanTable
/// <summary>
/// Initializes a new instance of the <see cref="HuffmanTable"/> struct.
/// </summary>
/// <param name="memoryAllocator">The <see cref="MemoryAllocator"/> to use for buffer allocations.</param>
/// <param name="codeLengths">The code lengths</param>
/// <param name="values">The huffman values</param>
public HuffmanTable(MemoryAllocator memoryAllocator, ReadOnlySpan<byte> codeLengths, ReadOnlySpan<byte> values)
public HuffmanTable(ReadOnlySpan<byte> codeLengths, ReadOnlySpan<byte> values)
{
this.isConfigured = false;
this.memoryAllocator = memoryAllocator;
Unsafe.CopyBlockUnaligned(ref this.Sizes[0], ref MemoryMarshal.GetReference(codeLengths), (uint)codeLengths.Length);
Unsafe.CopyBlockUnaligned(ref this.Values[0], ref MemoryMarshal.GetReference(values), (uint)values.Length);
}
Expand All @@ -81,33 +77,31 @@ public void Configure()
return;
}

int p, si;
Span<char> huffsize = stackalloc char[257];
Span<uint> huffcode = stackalloc uint[257];
uint code;
Span<char> huffSize = stackalloc char[257];
Span<uint> huffCode = stackalloc uint[257];

// Figure C.1: make table of Huffman code length for each symbol
p = 0;
int p = 0;
for (int l = 1; l <= 16; l++)
{
int i = this.Sizes[l];
while (i-- != 0)
{
huffsize[p++] = (char)l;
huffSize[p++] = (char)l;
}
}

huffsize[p] = (char)0;
huffSize[p] = (char)0;

// Figure C.2: generate the codes themselves
code = 0;
si = huffsize[0];
uint code = 0;
int si = huffSize[0];
p = 0;
while (huffsize[p] != 0)
while (huffSize[p] != 0)
{
while (huffsize[p] == si)
while (huffSize[p] == si)
{
huffcode[p++] = code;
huffCode[p++] = code;
code++;
}

Expand All @@ -121,10 +115,10 @@ public void Configure()
{
if (this.Sizes[l] != 0)
{
int offset = p - (int)huffcode[p];
int offset = p - (int)huffCode[p];
this.ValOffset[l] = offset;
p += this.Sizes[l];
this.MaxCode[l] = huffcode[p - 1]; // Maximum code of length l
this.MaxCode[l] = huffCode[p - 1]; // Maximum code of length l
this.MaxCode[l] <<= 64 - l; // Left justify
this.MaxCode[l] |= (1ul << (64 - l)) - 1;
}
Expand All @@ -150,19 +144,19 @@ public void Configure()
{
for (int i = 1; i <= this.Sizes[length]; i++, p++)
{
// length = current code's length, p = its index in huffcode[] & huffval[].
// length = current code's length, p = its index in huffCode[] & Values[].
// Generate left-justified code followed by all possible bit sequences
int lookbits = (int)(huffcode[p] << (JpegConstants.Huffman.LookupBits - length));
int lookBits = (int)(huffCode[p] << (JpegConstants.Huffman.LookupBits - length));
for (int ctr = 1 << (JpegConstants.Huffman.LookupBits - length); ctr > 0; ctr--)
{
this.LookaheadSize[lookbits] = (byte)length;
this.LookaheadValue[lookbits] = this.Values[p];
lookbits++;
this.LookaheadSize[lookBits] = (byte)length;
this.LookaheadValue[lookBits] = this.Values[p];
lookBits++;
}
}
}

this.isConfigured = true;
}
}
}
}
6 changes: 3 additions & 3 deletions src/ImageSharp/Formats/Jpeg/JpegDecoderCore.cs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -952,7 +952,7 @@ private void ProcessStartOfScanMarker()
/// <param name="values">The values</param>
[MethodImpl(InliningOptions.ShortMethod)]
private void BuildHuffmanTable(HuffmanTable[] tables, int index, ReadOnlySpan<byte> codeLengths, ReadOnlySpan<byte> values)
=> tables[index] = new HuffmanTable(this.configuration.MemoryAllocator, codeLengths, values);
=> tables[index] = new HuffmanTable(codeLengths, values);

/// <summary>
/// Reads a <see cref="ushort"/> from the stream advancing it by two bytes
Expand Down Expand Up @@ -992,4 +992,4 @@ private Image<TPixel> PostProcessIntoImage<TPixel>()
return image;
}
}
}
}
11 changes: 2 additions & 9 deletions tests/ImageSharp.Benchmarks/ImageSharp.Benchmarks.csproj
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,15 @@
<OutputType>Exe</OutputType>
<RootNamespace>SixLabors.ImageSharp.Benchmarks</RootNamespace>
<TargetFrameworks>netcoreapp2.1;net472</TargetFrameworks>
</PropertyGroup>

<PropertyGroup Condition="'$(TargetFramework)' == 'net461'">
<RuntimeIdentifier>win7-x64</RuntimeIdentifier>
<Prefer32Bit>false</Prefer32Bit>
<IsPackable>false</IsPackable>
<GenerateProgramFile>false</GenerateProgramFile>
</PropertyGroup>

<ItemGroup>
<Compile Include="..\ImageSharp.Tests\TestImages.cs" Link="Tests\TestImages.cs" />
<Compile Include="..\ImageSharp.Tests\TestUtilities\TestEnvironment.cs" Link="Tests\TestEnvironment.cs" />
</ItemGroup>

<ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp2.1'">
<Compile Remove="Program.cs" />
</ItemGroup>

<ItemGroup>
<PackageReference Include="BenchmarkDotNet" />
<PackageReference Include="Colourful" />
Expand Down
Loading

0 comments on commit 7a484f7

Please sign in to comment.