jellyfin-plugin-sso/tests/SSO-Auth.Tests/InMemorySsoStateStoreTests.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using Duende.IdentityModel.OidcClient;
using Jellyfin.Plugin.SSO_Auth.Api;
using Jellyfin.Plugin.SSO_Auth.State;
using Xunit;

namespace Jellyfin.Plugin.SSO_Auth.Tests;

public class InMemorySsoStateStoreTests
{
    private static readonly TimeSpan Expiry = TimeSpan.FromSeconds(300);

    [Fact]
    public async Task GetAsync_WhenUnknown_ReturnsNull()
    {
        Assert.Null(await new InMemorySsoStateStore().GetAsync("nope"));
    }

    [Fact]
    public async Task Mutations_AreOnlyVisibleOnceWrittenBack()
    {
        var store = new InMemorySsoStateStore();
        await store.CreateAsync(NewState("abc123"), Expiry);

        var dropped = await store.GetAsync("abc123");
        dropped.Valid = true;
        dropped.Username = "mallory";

        // Not written back, so the mutation must not be observable.
        var reread = await store.GetAsync("abc123");
        Assert.False(reread.Valid);
        Assert.Null(reread.Username);

        reread.Valid = true;
        reread.Username = "alice";
        reread.Admin = true;
        reread.Folders = new List<string> { "movies" };
        await store.UpdateAsync(reread, Expiry);

        var redeemed = await store.ConsumeAsync("abc123");
        Assert.True(redeemed.Valid);
        Assert.Equal("alice", redeemed.Username);
        Assert.True(redeemed.Admin);
        Assert.Equal(new[] { "movies" }, redeemed.Folders);
    }

    [Fact]
    public async Task ConsumeAsync_RedeemsOnlyOnce()
    {
        var store = new InMemorySsoStateStore();
        await store.CreateAsync(NewState("abc123"), Expiry);

        Assert.NotNull(await store.ConsumeAsync("abc123"));
        Assert.Null(await store.ConsumeAsync("abc123"));
        Assert.Null(await store.GetAsync("abc123"));
    }

    [Fact]
    public async Task ExpiredState_IsNotReturned()
    {
        var store = new InMemorySsoStateStore();
        await store.CreateAsync(NewState("abc123"), TimeSpan.FromMilliseconds(-1));

        Assert.Null(await store.GetAsync("abc123"));
        Assert.Empty(await store.ListAsync());
    }

    [Fact]
    public async Task ListAsync_ReturnsLiveStates()
    {
        var store = new InMemorySsoStateStore();
        await store.CreateAsync(NewState("abc123"), Expiry);
        await store.CreateAsync(NewState("def456"), Expiry);

        var states = await store.ListAsync();

        Assert.Equal(2, states.Count);
    }

    [Fact]
    public async Task Writing_SweepsExpiredEntriesOutOfTheStore()
    {
        var store = new InMemorySsoStateStore();
        await store.CreateAsync(NewState("stale1"), TimeSpan.FromMilliseconds(-1));
        await store.CreateAsync(NewState("stale2"), TimeSpan.FromMilliseconds(-1));
        await store.CreateAsync(NewState("stale3"), TimeSpan.FromMilliseconds(-1));

        await store.CreateAsync(NewState("live"), Expiry);

        Assert.Equal(1, store.Count);
    }

    [Fact]
    public async Task Writing_BoundsTheStoreAtItsCapacity()
    {
        var store = new InMemorySsoStateStore();

        for (var i = 0; i < InMemorySsoStateStore.Capacity + 64; i++)
        {
            await store.CreateAsync(NewState($"state{i}"), Expiry);
        }

        Assert.Equal(InMemorySsoStateStore.Capacity, store.Count);
    }

    [Fact]
    public async Task Writing_AtCapacity_IsSafeUnderConcurrency()
    {
        var store = new InMemorySsoStateStore();

        for (var i = 0; i < InMemorySsoStateStore.Capacity; i++)
        {
            await store.CreateAsync(NewState($"seed{i}"), Expiry);
        }

        const int Writers = 8;

        var writers = Enumerable.Range(0, Writers).Select(writer => Task.Run(async () =>
        {
            for (var i = 0; i < 2000; i++)
            {
                var state = NewState($"w{writer}-{i}");
                await store.CreateAsync(state, Expiry);
                await store.UpdateAsync(state, Expiry);
            }
        }));

        await Task.WhenAll(writers);

        // Each in-flight writer can miss or repeat one eviction, so the cap holds to within the writer count.
        Assert.InRange(store.Count, InMemorySsoStateStore.Capacity - Writers, InMemorySsoStateStore.Capacity + Writers);
    }

    private static TimedAuthorizeState NewState(string state)
    {
        return new TimedAuthorizeState(
            new AuthorizeState { State = state, CodeVerifier = "code-verifier-secret" },
            DateTime.UtcNow);
    }
}