ersatztv/ErsatzTV.Infrastructure/Streaming/Graphics/Image/ImageElementBase.cs

using System.Runtime.InteropServices;
using ErsatzTV.Core.Domain;
using ErsatzTV.Core.Interfaces.Streaming;
using ErsatzTV.FFmpeg.State;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats;
using SixLabors.ImageSharp.Formats.Gif;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.Formats.Webp;
using SixLabors.ImageSharp.Metadata;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using SkiaSharp;
using Image = SixLabors.ImageSharp.Image;

namespace ErsatzTV.Infrastructure.Streaming.Graphics;

public abstract class ImageElementBase(IRemoteImageFetcher remoteImageFetcher) : GraphicsElement, IDisposable
{
    /// <summary>
    ///     Ceiling on TOTAL decoded pixels for a remote image — width x height x frames, as one
    ///     product. Checking dimensions and frame count independently does not bound the decode:
    ///     2500x2500 x 600 frames is 60 KiB on the wire, passes a 50 MP dimension check and a 600
    ///     frame check, and costs ~14 GiB to decode. Only the product catches that.
    ///     8K is ~33 MP, so a single large still fits comfortably.
    /// </summary>
    internal const long MaxRemoteDecodedPixels = 50_000_000;

    /// <summary>
    ///     Frame ceiling for a remote animation, kept alongside the product budget as a cheap,
    ///     legible guard against absurd frame counts of tiny frames.
    /// </summary>
    internal const int MaxRemoteFrames = 600;

    /// <summary>
    ///     Ceiling on total pixels RETAINED after scaling — frames x scaled width x scaled height.
    ///     Independent of the source budget above: a 100x100 source is trivial to decode but, at 600
    ///     frames scaled to 1920x1080, retains ~5 GB of <see cref="SKBitmap" />. At 4 bytes per
    ///     pixel this bounds retention at ~800 MB, which still allows ~96 full-frame 1080p frames
    ///     (~3s at 30fps) or 600 frames of a 577x577 logo.
    /// </summary>
    internal const long MaxRemoteScaledPixels = 200_000_000;

    private readonly List<double> _frameDelays = [];
    private readonly List<SKBitmap> _scaledFrames = [];
    private double _animatedDurationSeconds;
    private ushort _repeatCount;

    private Image _sourceImage;

    protected SKPointI Location { get; private set; }

    public virtual void Dispose()
    {
        GC.SuppressFinalize(this);
        _sourceImage?.Dispose();
        _scaledFrames?.ForEach(f => f.Dispose());
    }

    protected async Task LoadImage(
        Resolution squarePixelFrameSize,
        Resolution frameSize,
        string image,
        WatermarkLocation location,
        bool scale,
        double? scaleWidthPercent,
        double? horizontalMarginPercent,
        double? verticalMarginPercent,
        bool placeWithinSourceContent,
        CancellationToken cancellationToken)
    {
        bool isRemoteUri = Uri.TryCreate(image, UriKind.Absolute, out Uri uriResult)
                           && (uriResult.Scheme == Uri.UriSchemeHttp || uriResult.Scheme == Uri.UriSchemeHttps);

        if (isRemoteUri)
        {
            await using Stream imageStream = await remoteImageFetcher.Fetch(uriResult, cancellationToken);
            _sourceImage = await DecodeRemoteImage(imageStream, uriResult, cancellationToken);
        }
        else
        {
            _sourceImage = await Image.LoadAsync(image!, cancellationToken);
        }

        int scaledWidth = _sourceImage.Width;
        int scaledHeight = _sourceImage.Height;
        if (scale)
        {
            scaledWidth = (int)Math.Round((scaleWidthPercent ?? 100) / 100.0 * frameSize.Width);
            double aspectRatio = (double)_sourceImage.Height / _sourceImage.Width;
            scaledHeight = (int)(scaledWidth * aspectRatio);
        }

        if (isRemoteUri)
        {
            EnsureScaledFramesAffordable(_sourceImage.Frames.Count, scaledWidth, scaledHeight, uriResult);
        }

        (int horizontalMargin, int verticalMargin) = placeWithinSourceContent
            ? SourceContentMargins(
                squarePixelFrameSize,
                frameSize,
                horizontalMarginPercent ?? 0,
                verticalMarginPercent ?? 0)
            : NormalMargins(frameSize, horizontalMarginPercent ?? 0, verticalMarginPercent ?? 0);

        Location = CalculatePosition(
            location,
            frameSize.Width,
            frameSize.Height,
            scaledWidth,
            scaledHeight,
            horizontalMargin,
            verticalMargin);

        if (_sourceImage.Metadata.DecodedImageFormat == GifFormat.Instance)
        {
            _repeatCount = _sourceImage.Metadata.GetFormatMetadata(GifFormat.Instance).RepeatCount;
        }

        _animatedDurationSeconds = 0;

        for (var i = 0; i < _sourceImage.Frames.Count; i++)
        {
            Image frame = _sourceImage.Frames.CloneFrame(i);
            frame.Mutate(ctx => ctx.Resize(scaledWidth, scaledHeight));
            _scaledFrames.Add(ToSkiaBitmap(frame));

            double frameDelay = GetFrameDelaySeconds(_sourceImage, i);
            _animatedDurationSeconds += frameDelay;
            _frameDelays.Add(frameDelay);
        }
    }

    /// <summary>
    ///     Decodes a remote image only after the header says decoding it is affordable.
    /// </summary>
    /// <remarks>
    ///     The fetcher's byte cap does NOT bound this: a decompression bomb is small on the wire and
    ///     huge in memory. A 4 KB PNG can declare 30000x30000 (~3.6 GB), and a 60 KiB GIF can
    ///     declare 2500x2500 across 600 frames (~14 GiB). The budget is therefore on the PRODUCT of
    ///     dimensions and frames, read from the header before the decoder allocates.
    ///     Local images are deliberately not checked — they are files an operator put on disk, not
    ///     bytes an arbitrary host returned. (ersatztv#511)
    /// </remarks>
    internal static async Task<Image> DecodeRemoteImage(Stream stream, Uri uri, CancellationToken cancellationToken)
    {
        if (!stream.CanSeek)
        {
            // Identify consumes the stream, so the decode below needs to rewind it. Fail with the
            // real reason rather than letting Position throw NotSupportedException, which the
            // caller's blanket catch would report as a generic initialization failure.
            throw new InvalidOperationException(
                $"Remote image {uri} was returned on a non-seekable stream; IRemoteImageFetcher must "
                + "return a fully buffered, seekable stream");
        }

        ImageInfo info = await Image.IdentifyAsync(stream, cancellationToken);

        // DIMENSIONS from the header are trustworthy; the FRAME COUNT is not, and is deliberately
        // not used as a budget input. Measured on ImageSharp 3.1.12: an APNG reports
        // FrameMetadataCollection.Count == 0 while the decoder happily produces 600 frames, so a
        // header-derived frame budget is enforced on a number the decoder does not honor — a
        // 134 KiB file decodes to ~36 GiB. (Second adversarial re-review; ersatztv#511.)
        EnsureDimensionsAffordable(info.Width, info.Height, uri);

        int affordableFrames = AffordableFrames(info.Width, info.Height);

        stream.Position = 0;

        // MaxFrames is enforced BY THE DECODER, so it holds whatever the header claimed. Measured:
        // MaxFrames = N yields N-1 frames (for N >= 2), so asking for affordableFrames + 2 decodes
        // at most affordableFrames + 1 — one more than allowed, which is exactly what lets the
        // post-decode check distinguish "exactly at the limit" from "over it" without truncating a
        // legitimate animation by a frame.
        var decoderOptions = new DecoderOptions { MaxFrames = (uint)(affordableFrames + 2) };

        Image image = await Image.LoadAsync(decoderOptions, stream, cancellationToken);

        try
        {
            // re-verify against REALITY rather than against the header. this is the check that
            // actually holds; everything above it only avoids decoding when we can tell in advance.
            EnsureDecodeAffordable(image.Width, image.Height, image.Frames.Count, uri);
            return image;
        }
        catch
        {
            image.Dispose();
            throw;
        }
    }

    /// <summary>Rejects a single frame that cannot fit the decode budget on its own.</summary>
    internal static void EnsureDimensionsAffordable(int width, int height, Uri uri)
    {
        long pixels = (long)width * height;
        if (pixels > MaxRemoteDecodedPixels)
        {
            throw new InvalidOperationException(
                $"Remote image {uri} is {width}x{height} ({pixels} pixels), over the "
                + $"{MaxRemoteDecodedPixels} pixel limit");
        }
    }

    /// <summary>
    ///     How many frames of this size the decode budget affords. Used to cap the DECODER, so the
    ///     bound does not depend on the header's frame count being honest.
    /// </summary>
    internal static int AffordableFrames(int width, int height)
    {
        long perFrame = Math.Max((long)width * height, 1);
        return (int)Math.Clamp(MaxRemoteDecodedPixels / perFrame, 1, MaxRemoteFrames);
    }

    /// <summary>
    ///     The decode-budget policy, kept free of I/O so the arithmetic can be tested at every
    ///     boundary without materializing multi-gigabyte images. Call this with the number of frames
    ///     the decoder ACTUALLY produced — never with a header-reported count, which can be zero for
    ///     an animation the decoder then expands to hundreds of frames.
    /// </summary>
    internal static void EnsureDecodeAffordable(int width, int height, int frameCount, Uri uri)
    {
        int frames = Math.Max(frameCount, 1);

        if (frames > MaxRemoteFrames)
        {
            throw new InvalidOperationException(
                $"Remote image {uri} has {frames} frames, over the {MaxRemoteFrames} frame limit");
        }

        // THE PRODUCT is the real bound. Checking dimensions and frames separately lets a 60 KiB
        // 2500x2500 x600 GIF through at a ~14 GiB decode cost. (Found by adversarial re-review of
        // the first fix for this, which checked them independently.)
        long totalPixels = (long)width * height * frames;
        if (totalPixels > MaxRemoteDecodedPixels)
        {
            throw new InvalidOperationException(
                $"Remote image {uri} decodes to {width}x{height} x{frames} frames "
                + $"({totalPixels} pixels), over the {MaxRemoteDecodedPixels} pixel limit");
        }
    }

    /// <summary>
    ///     Bounds what is RETAINED after scaling. Separate from the source budget because the two
    ///     are independent: a cheap-to-decode 100x100 source scaled to 1920x1080 across 600 frames
    ///     retains ~5 GB. Only applied to remote images, matching the rest of this guard.
    /// </summary>
    internal static void EnsureScaledFramesAffordable(int frameCount, int scaledWidth, int scaledHeight, Uri uri)
    {
        long retainedPixels = (long)Math.Max(frameCount, 1) * scaledWidth * scaledHeight;
        if (retainedPixels > MaxRemoteScaledPixels)
        {
            throw new InvalidOperationException(
                $"Remote image {uri} scales to {frameCount} frames of {scaledWidth}x{scaledHeight} "
                + $"({retainedPixels} pixels), over the {MaxRemoteScaledPixels} pixel limit");
        }
    }

    protected static SKBitmap ToSkiaBitmap(Image image)
    {
        using Image<Rgba32> rgbaImage = image.CloneAs<Rgba32>();

        int width = rgbaImage.Width;
        int height = rgbaImage.Height;

        var info = new SKImageInfo(width, height, SKColorType.Rgba8888, SKAlphaType.Unpremul);
        var skBitmap = new SKBitmap(info);
        if (!skBitmap.TryAllocPixels(info))
        {
            skBitmap.Dispose();
            throw new InvalidOperationException("Failed to allocate pixels for SKBitmap.");
        }

        var pixelArray = new Rgba32[width * height];
        rgbaImage.CopyPixelDataTo(pixelArray);

        var bytes = new byte[pixelArray.Length * 4];
        MemoryMarshal.AsBytes(pixelArray.AsSpan()).CopyTo(bytes);

        IntPtr dstPtr = skBitmap.GetPixels(out _);
        Marshal.Copy(bytes, 0, dstPtr, bytes.Length);

        return skBitmap;
    }

    protected static double GetFrameDelaySeconds(Image image, int frameIndex)
    {
        IImageFormat format = image.Metadata.DecodedImageFormat;
        ImageFrameMetadata frameMeta = image.Frames[frameIndex].Metadata;

        if (format == GifFormat.Instance)
        {
            // GIF frame delay is in hundredths of a second
            GifFrameMetadata gifMeta = frameMeta.GetFormatMetadata(GifFormat.Instance);
            return gifMeta.FrameDelay / 100.0;
        }

        if (format == PngFormat.Instance)
        {
            // PNG animated frame delay is in seconds (as double)
            PngFrameMetadata pngMeta = frameMeta.GetFormatMetadata(PngFormat.Instance);
            return pngMeta.FrameDelay.ToDouble();
        }

        if (format == WebpFormat.Instance)
        {
            // WEBP animated frame delay is in milliseconds
            WebpFrameMetadata webpMeta = frameMeta.GetFormatMetadata(WebpFormat.Instance);
            return webpMeta.FrameDelay / 1000.0;
        }

        // Default: assume 1/60th second (~16.67 ms) if unknown
        return 1.0 / 60.0;
    }

    protected SKBitmap GetFrameForTimestamp(TimeSpan timestamp)
    {
        if (_scaledFrames.Count <= 1)
        {
            return _scaledFrames[0];
        }

        if (_repeatCount > 0 && timestamp.TotalSeconds >= _animatedDurationSeconds * _repeatCount)
        {
            return _scaledFrames.Last();
        }

        double currentTime = timestamp.TotalSeconds % _animatedDurationSeconds;

        double frameTime = 0;
        for (var i = 0; i < _sourceImage.Frames.Count; i++)
        {
            frameTime += _frameDelays[i];
            if (currentTime <= frameTime)
            {
                return _scaledFrames[i];
            }
        }

        return _scaledFrames.Last();
    }
}