ersatztv/ErsatzTV.Infrastructure.Tests/Streaming/Graphics/RemoteImageDecodeLimitTests.cs

using System.Buffers.Binary;
using ErsatzTV.Infrastructure.Streaming.Graphics;
using NUnit.Framework;
using Shouldly;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.PixelFormats;
using Image = SixLabors.ImageSharp.Image;

namespace ErsatzTV.Infrastructure.Tests.Streaming.Graphics;

/// <summary>
///     The byte cap in <c>HttpRemoteImageFetcher</c> does not bound decoding: a decompression bomb
///     is tiny on the wire and enormous in memory. These pin the header-first check that does.
///     (ersatztv#511)
/// </summary>
[TestFixture]
public class RemoteImageDecodeLimitTests
{
    private static readonly Uri ImageUri = new("https://example.com/logo.png");

    [Test]
    public async Task Should_Decode_A_Normal_Image()
    {
        await using MemoryStream stream = await RealPng(64, 32);

        using Image image = await ImageElementBase.DecodeRemoteImage(stream, ImageUri, CancellationToken.None);

        image.Width.ShouldBe(64);
        image.Height.ShouldBe(32);
    }

    // the bomb: a few dozen bytes on the wire, ~3.6 GB if decoded. it sails through the byte cap,
    // the content-type check and the Content-Length reject -- only the header dimensions catch it.
    [Test]
    public async Task Should_Reject_An_Image_Whose_Declared_Dimensions_Are_A_Decompression_Bomb()
    {
        await using MemoryStream stream = PngHeaderDeclaring(30000, 30000);

        stream.Length.ShouldBeLessThan(100, "the point is that this is tiny on the wire");

        InvalidOperationException ex = await Should.ThrowAsync<InvalidOperationException>(
            () => ImageElementBase.DecodeRemoteImage(stream, ImageUri, CancellationToken.None));

        ex.Message.ShouldContain("pixel limit");
    }

    [Test]
    public async Task Should_Accept_Dimensions_Exactly_At_The_Limit()
    {
        // 10000 x 5000 = 50,000,000 -- exactly MaxRemotePixels, so it must NOT be rejected. the
        // decode then fails on the truncated body, which proves the check let it through.
        await using MemoryStream stream = PngHeaderDeclaring(10000, 5000);

        Exception ex = await Should.ThrowAsync<Exception>(
            () => ImageElementBase.DecodeRemoteImage(stream, ImageUri, CancellationToken.None));

        ex.Message.ShouldNotContain("pixel limit");
    }

    // PNG chunk CRC-32 (IEEE, reflected). hand-rolled because the repo does not reference
    // System.IO.Hashing, and ImageSharp validates the CRC of critical chunks like IHDR.
    private static uint Crc32(ReadOnlySpan<byte> data)
    {
        uint crc = 0xFFFFFFFF;
        foreach (byte b in data)
        {
            crc ^= b;
            for (var i = 0; i < 8; i++)
            {
                crc = (crc & 1) != 0 ? (crc >> 1) ^ 0xEDB88320 : crc >> 1;
            }
        }

        return crc ^ 0xFFFFFFFF;
    }

    /// <summary>A real, decodable PNG.</summary>
    private static async Task<MemoryStream> RealPng(int width, int height)
    {
        using var image = new Image<Rgba32>(width, height);
        var stream = new MemoryStream();
        await image.SaveAsync(stream, new PngEncoder());
        stream.Position = 0;
        return stream;
    }

    /// <summary>
    ///     A PNG signature plus a single valid IHDR chunk declaring <paramref name="width" /> x
    ///     <paramref name="height" /> and nothing else — enough for Identify, far too little to
    ///     decode. This is what a decompression bomb looks like at the point we have to reject it.
    /// </summary>
    private static MemoryStream PngHeaderDeclaring(int width, int height)
    {
        var stream = new MemoryStream();
        stream.Write([0x89, (byte)'P', (byte)'N', (byte)'G', 0x0D, 0x0A, 0x1A, 0x0A]);

        var ihdr = new byte[17];
        "IHDR"u8.CopyTo(ihdr);
        BinaryPrimitives.WriteInt32BigEndian(ihdr.AsSpan(4), width);
        BinaryPrimitives.WriteInt32BigEndian(ihdr.AsSpan(8), height);
        ihdr[12] = 8; // bit depth
        ihdr[13] = 6; // color type: truecolor + alpha
        ihdr[14] = 0; // compression
        ihdr[15] = 0; // filter
        ihdr[16] = 0; // interlace

        var length = new byte[4];
        BinaryPrimitives.WriteInt32BigEndian(length, 13);
        stream.Write(length);
        stream.Write(ihdr);

        var crc = new byte[4];
        BinaryPrimitives.WriteUInt32BigEndian(crc, Crc32(ihdr));
        stream.Write(crc);

        stream.Position = 0;
        return stream;
    }
}