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 the budget, so it must NOT be rejected. the decode
        // then fails on the truncated body, which proves the check let it through. NOTE this test
        // would also pass with the guard deleted entirely; deletion is covered by the bomb test
        // above, and the boundary arithmetic by EnsureDecodeAffordable's own tests.
        await using MemoryStream stream = PngHeaderDeclaring(10000, 5000);

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

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

    // --- the budget policy itself, tested as arithmetic so no multi-GB image is ever allocated ---

    // THE bomb the first fix missed: 2500x2500 x600 frames is ~60 KiB on the wire, passes a
    // dimensions-only check (6.25 MP) AND a frames-only check (exactly 600), and costs ~14 GiB to
    // decode. Only the PRODUCT catches it. (Found by adversarial re-review; ersatztv#511.)
    [Test]
    public void Should_Reject_Dimensions_And_Frames_That_Are_Affordable_Alone_But_Not_Together()
    {
        const int Width = 2500;
        const int Height = 2500;
        const int Frames = 600;

        // each guard, in isolation, says yes
        ((long)Width * Height).ShouldBeLessThanOrEqualTo(ImageElementBase.MaxRemoteDecodedPixels);
        Frames.ShouldBeLessThanOrEqualTo(ImageElementBase.MaxRemoteFrames);

        InvalidOperationException ex = Should.Throw<InvalidOperationException>(
            () => ImageElementBase.EnsureDecodeAffordable(Width, Height, Frames, ImageUri));

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

    // M1: the frame guard had no coverage at all in the first fix
    [Test]
    public void Should_Reject_Too_Many_Frames_Even_When_Each_Is_Tiny()
    {
        InvalidOperationException ex = Should.Throw<InvalidOperationException>(
            () => ImageElementBase.EnsureDecodeAffordable(8, 8, ImageElementBase.MaxRemoteFrames + 1, ImageUri));

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

    [Test]
    public void Should_Allow_A_Single_Large_Still_Within_Budget()
    {
        // 8K is ~33 MP -- must keep working
        Should.NotThrow(() => ImageElementBase.EnsureDecodeAffordable(7680, 4320, 1, ImageUri));
    }

    [Test]
    public void Should_Allow_A_Typical_Animated_Logo()
    {
        Should.NotThrow(() => ImageElementBase.EnsureDecodeAffordable(288, 288, 600, ImageUri));
    }

    [Test]
    public void Should_Afford_Fewer_Frames_As_Frames_Get_Larger()
    {
        // tiny frames are capped by the frame guard, not the pixel budget
        ImageElementBase.AffordableFrames(8, 8).ShouldBe(ImageElementBase.MaxRemoteFrames);

        // 1000x1000 -> 50M / 1M = 50 frames
        ImageElementBase.AffordableFrames(1000, 1000).ShouldBe(50);

        // a frame so large only one fits
        ImageElementBase.AffordableFrames(7000, 7000).ShouldBe(1);
    }

    [Test]
    public void Should_Allow_A_Product_Exactly_At_The_Budget()
    {
        // 10000 x 5000 x 1 == MaxRemoteDecodedPixels exactly
        Should.NotThrow(() => ImageElementBase.EnsureDecodeAffordable(10000, 5000, 1, ImageUri));
    }

    // the retained-frame budget is INDEPENDENT of the decode budget: this source is trivial to
    // decode (6 MP total) but retains ~5 GB of SKBitmap once every frame is scaled to 1080p
    [Test]
    public void Should_Reject_Cheap_Frames_That_Are_Expensive_Once_Scaled()
    {
        Should.NotThrow(() => ImageElementBase.EnsureDecodeAffordable(100, 100, 600, ImageUri));

        InvalidOperationException ex = Should.Throw<InvalidOperationException>(
            () => ImageElementBase.EnsureScaledFramesAffordable(600, 1920, 1080, ImageUri));

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

    [Test]
    public void Should_Allow_A_Scaled_Watermark_Sized_Animation()
    {
        // a 10%-width logo on a 1080p frame, animated
        Should.NotThrow(() => ImageElementBase.EnsureScaledFramesAffordable(600, 192, 108, ImageUri));
    }


    // --- B1 regression: the header's frame count is a LIE for APNG ---

    // ImageSharp 3.1.12 reports FrameMetadataCollection.Count == 0 for an APNG while the decoder
    // produces every frame. A budget derived from that header count is enforced on a number the
    // decoder does not honor — this exact payload shape, at 4000x4000, is ~134 KiB on the wire and
    // ~36 GiB decoded. The bound therefore has to be imposed ON THE DECODER (DecoderOptions.
    // MaxFrames) and re-verified against the real frame count. (ersatztv#511, second re-review.)
    [Test]
    public async Task Should_Reject_An_Animation_Whose_Header_Under_Reports_Its_Frames()
    {
        await using MemoryStream stream = Apng(64, 64, ImageElementBase.MaxRemoteFrames + 100);

        // the premise: the header really does under-report, so a header-derived budget waves it through
        stream.Position = 0;
        ImageInfo info = await Image.IdentifyAsync(stream);
        info.FrameMetadataCollection.Count.ShouldBe(0, "the APNG header under-reports; that is the whole point");
        Should.NotThrow(() => ImageElementBase.EnsureDecodeAffordable(64, 64, info.FrameMetadataCollection.Count, ImageUri));

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

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

    // ...and an animation within budget still decodes IN FULL -- the decoder cap must not silently
    // truncate legitimate content by a frame
    [Test]
    public async Task Should_Decode_An_Animation_Within_Budget_Without_Truncating_It()
    {
        await using MemoryStream stream = Apng(64, 64, 300);

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

        image.Frames.Count.ShouldBe(300);
    }

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