ersatztv/ErsatzTV.Infrastructure.Tests/Images/RemoteImageValidatorTests.cs

using System.Buffers.Binary;
using ErsatzTV.Core.Images;
using ErsatzTV.Core.Interfaces.Images;
using ErsatzTV.Infrastructure.Images;
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.Images;

[TestFixture]
public class RemoteImageValidatorTests
{
    private static readonly Uri Uri = new("https://example.com/logo.png");

    // decode cases exercise the static method (used by the render path)
    [Test]
    public async Task Should_Decode_A_Normal_Image()
    {
        await using MemoryStream stream = await RealPng(64, 32);
        using Image image = await RemoteImageValidator.DecodeAndValidate(stream, Uri, CancellationToken.None);
        image.Width.ShouldBe(64);
        image.Height.ShouldBe(32);
    }

    [Test]
    public async Task Should_Reject_A_Decompression_Bomb_By_Declared_Dimensions()
    {
        await using MemoryStream stream = PngHeaderDeclaring(30000, 30000);
        InvalidOperationException ex = await Should.ThrowAsync<InvalidOperationException>(
            () => RemoteImageValidator.DecodeAndValidate(stream, Uri, CancellationToken.None));
        ex.Message.ShouldContain("pixel limit");
    }

    [Test]
    public async Task Should_Reject_An_Apng_Whose_Header_Under_Reports_Its_Frames()
    {
        await using MemoryStream stream = Apng(64, 64, RemoteImageDecodeBudget.MaxRemoteFrames + 100);
        InvalidOperationException ex = await Should.ThrowAsync<InvalidOperationException>(
            () => RemoteImageValidator.DecodeAndValidate(stream, Uri, CancellationToken.None));
        ex.Message.ShouldContain("frame limit");
    }

    [Test]
    public async Task Should_Decode_An_Apng_That_A_Default_Identify_Cannot_Read()
    {
        await using MemoryStream stream = Apng(288, 288, 60);
        stream.Position = 0;
        await Should.ThrowAsync<Exception>(() => Image.IdentifyAsync(stream));
        stream.Position = 0;
        using Image image = await RemoteImageValidator.DecodeAndValidate(stream, Uri, CancellationToken.None);
        image.Frames.Count.ShouldBe(60);
    }

    // the Core interface Validate() is the save/upload contract: throws on invalid, returns on valid,
    // never surfaces an ImageSharp type
    [Test]
    public async Task Validate_Returns_On_A_Good_Image()
    {
        IRemoteImageValidator validator = new RemoteImageValidator();
        await using MemoryStream stream = await RealPng(64, 32);
        await Should.NotThrowAsync(() => validator.Validate(stream, Uri, CancellationToken.None));
    }

    [Test]
    public async Task Validate_Throws_On_A_Bomb()
    {
        IRemoteImageValidator validator = new RemoteImageValidator();
        await using MemoryStream stream = PngHeaderDeclaring(30000, 30000);
        await Should.ThrowAsync<InvalidOperationException>(
            () => validator.Validate(stream, Uri, CancellationToken.None));
    }

    /// <summary>A real multi-frame APNG. Small on the wire, many frames — the shape that matters.</summary>
    private static MemoryStream Apng(int width, int height, int frames)
    {
        using var image = new Image<Rgba32>(width, height);
        for (var i = 1; i < frames; i++)
        {
            image.Frames.CreateFrame();
        }

        var stream = new MemoryStream();
        image.Save(stream, new PngEncoder { ColorType = PngColorType.RgbWithAlpha });
        stream.Position = 0;
        return stream;
    }

    // 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;
    }
}