av-parser/tests/AvParser.Infrastructure.Tests/Collecting/MediaFetcherTests.cs

using System.Net;
using System.Text;
using AvParser.Core.Collecting;
using AvParser.Core.Proxies;
using AvParser.Infrastructure.Collecting;
using AvParser.Infrastructure.Media;
using AvParser.Infrastructure.Proxies;
using AvParser.Infrastructure.Storage;
using Microsoft.Extensions.Logging.Abstractions;

namespace AvParser.Infrastructure.Tests.Collecting;

/// <summary>
/// A factory that never uses a proxy but mirrors the real handler's configuration.
/// </summary>
/// <remarks>
/// Matching <c>AllowAutoRedirect = false</c> and <c>AutomaticDecompression = All</c> matters: with
/// redirects handled by the BCL these tests would exercise its logic rather than the fetcher's.
/// </remarks>
internal sealed class DirectHttpClientFactory : IProxiedHttpClientFactory
{
    public HttpClient Create(ProxyEndpoint? endpoint, TimeSpan? timeout = null) =>
        Create(endpoint, HttpClientTimeouts.Default);

    public HttpClient Create(ProxyEndpoint? endpoint, HttpClientTimeouts timeouts) =>
        new(
            new SocketsHttpHandler
            {
                AllowAutoRedirect = false,
                ConnectTimeout = timeouts.Connect,
                AutomaticDecompression = DecompressionMethods.All,
            },
            disposeHandler: true
        )
        {
            Timeout = Timeout.InfiniteTimeSpan,
        };

    public Task<(HttpClient Client, ProxyLease? Lease)> CreateFromPoolAsync(
        TimeSpan? timeout = null,
        CancellationToken cancellationToken = default
    ) => Task.FromResult<(HttpClient, ProxyLease?)>((Create(null, timeout), null));

    public Task<LeasedHttpClient> LeaseAsync(
        HttpClientTimeouts timeouts,
        bool requireProxy,
        CancellationToken cancellationToken = default
    ) =>
        requireProxy
            ? throw new ProxyUnavailableException()
            : Task.FromResult(new LeasedHttpClient(Create(null, timeouts), null));
}

public sealed class MediaFetcherTests : IAsyncLifetime
{
    private readonly string _root = Path.Combine(Path.GetTempPath(), "AvParserTests", Guid.NewGuid().ToString("N"));

    private LoopbackServer _server = null!;
    private AppPaths _paths = null!;
    private BlobStore _blobs = null!;
    private HostThrottle _throttle = null!;
    private MediaFetcher _fetcher = null!;

    public ValueTask InitializeAsync()
    {
        _paths = new AppPaths(_root);
        _paths.EnsureCreated();

        _server = new LoopbackServer();
        _blobs = new BlobStore(_paths, NullLogger<BlobStore>.Instance);
        _throttle = Throttle();
        _fetcher = new MediaFetcher(
            new DirectHttpClientFactory(),
            _blobs,
            _throttle,
            NullLogger<MediaFetcher>.Instance
        );

        return ValueTask.CompletedTask;
    }

    public async ValueTask DisposeAsync()
    {
        await _server.DisposeAsync();
        _throttle.Dispose();

        if (Directory.Exists(_root))
        {
            try
            {
                Directory.Delete(_root, recursive: true);
            }
            catch (IOException)
            {
                // Not worth failing a green test over.
            }
        }
    }

    private HostThrottle Throttle(int concurrent = 4, int intervalMs = 0) =>
        new(concurrent, TimeSpan.FromMilliseconds(intervalMs), NullLogger<MediaFetcherTests>.Instance);

    /// <summary>Options with the floor lowered, since the sample files are deliberately tiny.</summary>
    private static FetchOptions Options(params (string Key, object Value)[] _) =>
        new() { MinItemBytes = 1, Timeouts = HttpClientTimeouts.Default };

    private Task<FetchResult> FetchAsync(Uri url, FetchOptions? options = null) =>
        _fetcher.FetchAsync(
            new MediaCandidate(url) { SourceId = "test" },
            options ?? Options(),
            TestContext.Current.CancellationToken
        );

    /// <summary>Waits for the server's own bookkeeping to catch up with the client.</summary>
    private async Task<List<RecordedRequest>> WaitForRequestsAsync(int count, Func<RecordedRequest, bool> match)
    {
        for (var attempt = 0; attempt < 100; attempt++)
        {
            var matched = _server.Requests.Where(match).OrderBy(r => r.ArrivedUtc).ToList();

            if (matched.Count >= count)
            {
                return matched;
            }

            await Task.Delay(20, TestContext.Current.CancellationToken);
        }

        var found = _server.Requests.Where(match).OrderBy(r => r.ArrivedUtc).ToList();
        found.Count.ShouldBe(count);

        return found;
    }

    private int StagedFileCount() =>
        Directory.Exists(_paths.MediaTempDirectory) ? Directory.EnumerateFiles(_paths.MediaTempDirectory).Count() : 0;

    [Fact]
    public async Task A_valid_image_is_downloaded_and_hashed()
    {
        var url = _server.MapBody("/a.png", Samples.Png(), "image/png");

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(SeenOutcome.Stored);
        result.Blob!.Kind.ShouldBe(MediaKind.Png);
        result.Blob.Sha256.Length.ShouldBe(64);
        result.TempPath.ShouldNotBeNull();
        File.Exists(result.TempPath).ShouldBeTrue();
    }

    [Fact]
    public async Task The_extension_comes_from_the_signature_not_from_the_content_type()
    {
        // The origin says PNG and serves a GIF. Believing the header would write a .png that no
        // viewer opens, and would let a host dictate the extension of a file on the user's disk.
        var url = _server.MapBody("/liar.png", Samples.Gif(), "image/png");

        var result = await FetchAsync(url);

        result.Blob!.Kind.ShouldBe(MediaKind.Gif);
        result.Blob.Extension.ShouldBe(".gif");
        result.ContentType.ShouldBe("image/png");
    }

    [Fact]
    public async Task An_error_page_behind_a_200_is_refused()
    {
        var url = _server.MapBody("/gone.jpg", Samples.Html(), "image/jpeg");

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(SeenOutcome.NotMedia);
        result.TempPath.ShouldBeNull();
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_kind_the_user_excluded_is_refused()
    {
        var url = _server.MapBody("/clip.mp4", Samples.Mp4());
        var options = Options() with { AllowedKinds = new HashSet<MediaKind> { MediaKind.Png } };

        (await FetchAsync(url, options)).Outcome.ShouldBe(SeenOutcome.UnsupportedType);
    }

    [Theory]
    [InlineData(403, SeenOutcome.Failed)]
    [InlineData(404, SeenOutcome.Gone)]
    [InlineData(410, SeenOutcome.Gone)]
    [InlineData(500, SeenOutcome.Failed)]
    public async Task Error_statuses_are_reported_without_staging_anything(int status, SeenOutcome expected)
    {
        var url = _server.Map($"/e{status}", new Reply { Status = status, Body = "no"u8.ToArray() });

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(expected);
        result.HttpStatus.ShouldBe(status);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_rate_limit_is_honoured_rather_than_worked_around()
    {
        var url = _server.Map(
            "/busy",
            new Reply
            {
                Status = 429,
                Headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase) { ["Retry-After"] = "3" },
            }
        );

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(SeenOutcome.RateLimited);
        result.RetryAfter.ShouldBe(TimeSpan.FromSeconds(3));

        // The host is now cooling: the collector slows down instead of switching proxy.
        _throttle.CooldownRemaining(url).ShouldBeGreaterThan(TimeSpan.Zero);
    }

    [Fact]
    public async Task A_service_unavailable_without_a_hint_still_backs_off()
    {
        var url = _server.Map("/down", new Reply { Status = 503 });

        (await FetchAsync(url)).Outcome.ShouldBe(SeenOutcome.RateLimited);
        _throttle.CooldownRemaining(url).ShouldBeGreaterThan(TimeSpan.Zero);
    }

    [Fact]
    public async Task Headers_that_never_arrive_time_out()
    {
        var url = _server.Map("/slowhead", new Reply { Body = Samples.Png(), HeaderDelay = TimeSpan.FromSeconds(3) });

        var options = Options() with
        {
            Timeouts = HttpClientTimeouts.Default with { Headers = TimeSpan.FromMilliseconds(300) },
        };

        (await FetchAsync(url, options)).Outcome.ShouldBe(SeenOutcome.Timeout);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_body_that_stops_arriving_is_abandoned_and_leaves_nothing_behind()
    {
        var url = _server.Map(
            "/drip",
            new Reply
            {
                Body = Encoding.ASCII.GetBytes(new string('x', 4096)),
                BodyChunkSize = 8,
                BodyDelay = TimeSpan.FromSeconds(2),
            }
        );

        var options = Options() with
        {
            Timeouts = HttpClientTimeouts.Default with { Idle = TimeSpan.FromMilliseconds(300) },
        };

        var result = await FetchAsync(url, options);

        result.Outcome.ShouldBe(SeenOutcome.Timeout);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_body_shorter_than_declared_is_a_truncated_transfer_not_a_small_file()
    {
        // The single most important rejection: a partial image promoted into the blob store would
        // be indistinguishable from a real one for ever afterwards.
        var body = Samples.Png();
        var url = _server.Map(
            "/short",
            new Reply
            {
                Body = body,
                DeclaredLength = body.Length + 5000,
                TruncateAfter = body.Length,
            }
        );

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(SeenOutcome.Failed);
        result.ErrorCode.ShouldBe("Truncated");
        result.TempPath.ShouldBeNull();
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_declared_length_over_the_cap_is_refused_before_the_body_is_read()
    {
        var body = Encoding.ASCII.GetBytes(new string('x', 8192));
        var url = _server.Map("/huge", new Reply { Body = body });
        var options = Options() with { MaxItemBytes = 1024 };

        var result = await FetchAsync(url, options);

        result.Outcome.ShouldBe(SeenOutcome.TooLarge);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_chunked_body_that_grows_past_the_cap_is_cut_off()
    {
        // No Content-Length to check against, so the cap has to be enforced as bytes arrive.
        var body = Encoding.ASCII.GetBytes(new string('x', 8192));
        var url = _server.Map(
            "/huge-chunked",
            new Reply
            {
                Body = body,
                Chunked = true,
                BodyChunkSize = 256,
            }
        );
        var options = Options() with { MaxItemBytes = 1024 };

        (await FetchAsync(url, options)).Outcome.ShouldBe(SeenOutcome.TooLarge);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_well_formed_chunked_body_is_accepted()
    {
        var url = _server.Map(
            "/chunked.png",
            new Reply
            {
                Body = Samples.Png(),
                Chunked = true,
                BodyChunkSize = 7,
            }
        );

        var result = await FetchAsync(url);

        result.Outcome.ShouldBe(SeenOutcome.Stored);
        result.Blob!.Length.ShouldBe(Samples.Png().Length);
    }

    [Fact]
    public async Task The_same_image_served_gzipped_and_plain_is_one_blob()
    {
        // The handler decompresses transparently, so the hash must be over the decoded bytes.
        var plain = _server.MapBody("/plain.png", Samples.Png());
        var zipped = _server.Map("/zipped.png", new Reply { Body = Samples.Png(), Gzip = true });

        var first = await FetchAsync(plain);
        var second = await FetchAsync(zipped);

        first.Blob!.Sha256.ShouldBe(second.Blob!.Sha256);
        first.Blob.Length.ShouldBe(second.Blob.Length);
    }

    [Fact]
    public async Task A_tracking_pixel_is_below_the_floor()
    {
        var url = _server.MapBody("/pixel.gif", Samples.Gif());

        // The default floor, rather than the lowered one the other tests use.
        var result = await FetchAsync(url, new FetchOptions());

        result.Outcome.ShouldBe(SeenOutcome.TooSmall);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task A_redirect_chain_is_followed_and_the_final_address_recorded()
    {
        var final = _server.MapBody("/final.png", Samples.Png());
        _server.Map("/hop2", Redirect(302, "/final.png"));
        var start = _server.Map("/hop1", Redirect(301, "/hop2"));

        var result = await FetchAsync(start);

        result.Outcome.ShouldBe(SeenOutcome.Stored);
        result.FinalUrl.ShouldBe(final);
    }

    [Fact]
    public async Task A_relative_location_is_resolved_against_the_current_address()
    {
        _server.MapBody("/other.png", Samples.Png());
        var start = _server.Map("/rel", Redirect(302, "/other.png"));

        (await FetchAsync(start)).Outcome.ShouldBe(SeenOutcome.Stored);
    }

    [Fact]
    public async Task A_redirect_loop_is_broken()
    {
        _server.Map("/loopB", Redirect(302, "/loopA"));
        var start = _server.Map("/loopA", Redirect(302, "/loopB"));

        var result = await FetchAsync(start);

        result.ErrorCode.ShouldBe("RedirectLoop");
    }

    [Fact]
    public async Task A_chain_longer_than_the_cap_gives_up()
    {
        // Distinct hops, so this is the hop budget rather than the loop detector doing the work.
        for (var hop = 0; hop < 10; hop++)
        {
            _server.Map($"/chain{hop}", Redirect(302, $"/chain{hop + 1}"));
        }

        var result = await FetchAsync(new Uri(_server.BaseAddress, "/chain0"), Options() with { MaxRedirects = 3 });

        result.ErrorCode.ShouldBe("TooManyRedirects");
    }

    [Fact]
    public async Task A_redirect_to_a_non_http_scheme_is_refused()
    {
        var start = _server.Map("/evil", Redirect(302, "data:image/png;base64,AAAA"));

        (await FetchAsync(start)).ErrorCode.ShouldBe("BadRedirect");
    }

    [Fact]
    public async Task A_known_placeholder_is_recognised_by_its_hash()
    {
        var url = _server.MapBody("/dead.png", Samples.Png());
        var probe = await FetchAsync(url);
        File.Delete(probe.TempPath!);

        var options = Options() with
        {
            Tombstones = new HashSet<string>(StringComparer.Ordinal) { probe.Blob!.Sha256 },
        };

        var result = await FetchAsync(url, options);

        result.Outcome.ShouldBe(SeenOutcome.Placeholder);
        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task Animation_is_detected_end_to_end()
    {
        var still = _server.MapBody("/still.gif", Samples.Gif(frames: 1));
        var moving = _server.MapBody("/moving.gif", Samples.Gif(frames: 3));

        (await FetchAsync(still)).Blob!.IsAnimated.ShouldBeFalse();
        (await FetchAsync(moving)).Blob!.IsAnimated.ShouldBeTrue();
    }

    [Fact]
    public async Task Dimensions_are_recorded_when_the_header_carries_them()
    {
        var url = _server.MapBody("/big.png", Samples.Png(1024, 768));

        var blob = (await FetchAsync(url)).Blob!;

        blob.Width.ShouldBe(1024);
        blob.Height.ShouldBe(768);
    }

    [Fact]
    public async Task A_referer_is_sent_when_the_candidate_carries_one()
    {
        var url = _server.MapBody("/ref.png", Samples.Png());
        var referer = new Uri("https://example.test/page");

        await _fetcher.FetchAsync(
            new MediaCandidate(url) { SourceId = "test", Referer = referer },
            Options(),
            TestContext.Current.CancellationToken
        );

        var request = _server.Requests.Single(r => r.Path == "/ref.png");
        request.Headers["Referer"].ShouldBe(referer.AbsoluteUri);
        request.Headers["User-Agent"].ShouldContain("AvParser");
    }

    [Fact]
    public async Task One_request_at_a_time_per_host_when_the_cap_says_so()
    {
        _throttle.Dispose();
        _throttle = Throttle(concurrent: 1);
        _fetcher = new MediaFetcher(
            new DirectHttpClientFactory(),
            _blobs,
            _throttle,
            NullLogger<MediaFetcher>.Instance
        );

        // Real media padded out, so neither request is refused after 32 bytes and cut short —
        // the point of the test is two full responses that could have overlapped.
        var body = (byte[])[.. Samples.Png(), .. new byte[2048]];
        var slow = new Reply
        {
            Body = body,
            BodyChunkSize = 256,
            BodyDelay = TimeSpan.FromMilliseconds(40),
        };

        _server.Map("/one", slow);
        _server.Map("/two", slow);

        await Task.WhenAll(
            FetchAsync(new Uri(_server.BaseAddress, "/one")),
            FetchAsync(new Uri(_server.BaseAddress, "/two"))
        );

        // The server records a request from its own task, which can lag the client by a moment.
        var requests = await WaitForRequestsAsync(2, r => r.Path is "/one" or "/two");

        // Serialised, the second request cannot arrive until the first body has been written:
        // nine 256-byte writes at 40 ms apiece. Run concurrently the gap would be near zero.
        // Compared against arrival rather than completion because the server records completion
        // from its own task, which can lag the client by a few milliseconds.
        var gap = requests[1].ArrivedUtc - requests[0].ArrivedUtc;
        gap.ShouldBeGreaterThan(TimeSpan.FromMilliseconds(200));
    }

    [Fact]
    public async Task Cancelling_mid_download_leaves_no_staged_file()
    {
        // Real media, or the fetcher rightly rejects it after 32 bytes and there is nothing left
        // to cancel.
        var url = _server.Map(
            "/cancel",
            new Reply
            {
                Body = [.. Samples.Png(), .. new byte[65536]],
                BodyChunkSize = 128,
                BodyDelay = TimeSpan.FromMilliseconds(30),
            }
        );

        using var cancellation = new CancellationTokenSource();
        var fetch = _fetcher.FetchAsync(new MediaCandidate(url) { SourceId = "test" }, Options(), cancellation.Token);

        await Task.Delay(200, TestContext.Current.CancellationToken);
        await cancellation.CancelAsync();

        await Should.ThrowAsync<OperationCanceledException>(async () => await fetch);

        StagedFileCount().ShouldBe(0);
    }

    [Fact]
    public async Task Demanding_a_proxy_that_is_not_there_fails_the_item_rather_than_going_direct()
    {
        // Otherwise the request leaves from the user's own address at precisely the moment they
        // asked for it not to.
        var url = _server.MapBody("/direct.png", Samples.Png());
        var options = Options() with { RequireProxy = true };

        var result = await FetchAsync(url, options);

        result.ErrorCode.ShouldBe("NoProxy");
        _server.Requests.ShouldNotContain(r => r.Path == "/direct.png");
    }

    private static Reply Redirect(int status, string location) =>
        new()
        {
            Status = status,
            Headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase) { ["Location"] = location },
        };
}