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

using System;
using System.Collections.Generic;
using System.Net;
using System.Threading.Tasks;
using Duende.IdentityModel.OidcClient;
using Jellyfin.Plugin.SSO_Auth.Api;
using Jellyfin.Plugin.SSO_Auth.State;
using Moq;
using StackExchange.Redis;
using Xunit;

namespace Jellyfin.Plugin.SSO_Auth.Tests;

public class RedisSsoStateStoreTests
{
    [Fact]
    public async Task CreateAsync_WritesPrefixedKeyWithExpiry()
    {
        var db = new Mock<IDatabase>();
        RedisKey key = default;
        RedisValue value = default;
        TimeSpan? expiry = null;
        SetupStringSet(db, (k, v, e) =>
        {
            key = k;
            value = v;
            expiry = e;
        });

        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));
        await store.CreateAsync(NewState("abc123"), TimeSpan.FromSeconds(300));

        Assert.Equal("jellyfin:sso:oid:state:abc123", key.ToString());
        Assert.Equal(TimeSpan.FromSeconds(300), expiry);
        Assert.Contains("abc123", value.ToString(), StringComparison.Ordinal);
    }

    [Fact]
    public async Task UpdateAsync_RefreshesExpiry()
    {
        var db = new Mock<IDatabase>();
        TimeSpan? expiry = null;
        SetupStringSet(db, (k, v, e) => expiry = e);

        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));
        await store.UpdateAsync(NewState("abc123"), TimeSpan.FromSeconds(120));

        Assert.Equal(TimeSpan.FromSeconds(120), expiry);
    }

    [Fact]
    public async Task GetAsync_WhenKeyMissing_ReturnsNull()
    {
        var db = new Mock<IDatabase>();
        db.Setup(d => d.StringGetAsync(It.IsAny<RedisKey>(), It.IsAny<CommandFlags>()))
            .ReturnsAsync(RedisValue.Null);

        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));

        Assert.Null(await store.GetAsync("missing"));
    }

    [Fact]
    public async Task GetAsync_WithoutState_DoesNotTouchRedis()
    {
        var db = new Mock<IDatabase>(MockBehavior.Strict);
        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));

        Assert.Null(await store.GetAsync(null));
        Assert.Null(await store.GetAsync(string.Empty));
    }

    [Fact]
    public async Task Mutations_WrittenByTheCallback_AreVisibleToTheAuthLeg()
    {
        var db = new Mock<IDatabase>();
        string stored = null;
        SetupStringSet(db, (k, v, e) => stored = v.ToString());
        db.Setup(d => d.StringGetAsync(It.IsAny<RedisKey>(), It.IsAny<CommandFlags>()))
            .ReturnsAsync(() => stored is null ? RedisValue.Null : (RedisValue)stored);
        db.Setup(d => d.StringGetDeleteAsync(It.IsAny<RedisKey>(), It.IsAny<CommandFlags>()))
            .ReturnsAsync(() =>
            {
                var payload = stored;
                stored = null;
                return payload is null ? RedisValue.Null : (RedisValue)payload;
            });

        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));
        var expiry = TimeSpan.FromSeconds(300);

        // Leg 1: the start endpoint stores a freshly prepared state.
        await store.CreateAsync(NewState("abc123"), expiry);

        // Leg 2: the callback endpoint validates the response and writes its decision back.
        var callbackState = await store.GetAsync("abc123");
        Assert.NotNull(callbackState);
        Assert.False(callbackState.Valid);
        callbackState.Valid = true;
        callbackState.Username = "alice";
        callbackState.Admin = true;
        callbackState.AvatarURL = "https://idp.example/alice.png";
        await store.UpdateAsync(callbackState, expiry);

        // Leg 3: the auth endpoint redeems the state and sees the callback's decision.
        var authState = await store.ConsumeAsync("abc123");
        Assert.NotNull(authState);
        Assert.True(authState.Valid);
        Assert.Equal("alice", authState.Username);
        Assert.True(authState.Admin);
        Assert.Equal("https://idp.example/alice.png", authState.AvatarURL);
        Assert.Equal("code-verifier-secret", authState.State.CodeVerifier);

        // A replay of the same response finds nothing.
        Assert.Null(await store.ConsumeAsync("abc123"));
    }

    [Fact]
    public async Task ConsumeAsync_UsesAnAtomicGetAndDelete()
    {
        var db = new Mock<IDatabase>(MockBehavior.Strict);
        db.Setup(d => d.StringGetDeleteAsync(It.IsAny<RedisKey>(), It.IsAny<CommandFlags>()))
            .ReturnsAsync(RedisValue.Null);

        var store = new RedisSsoStateStore(CreateMultiplexer(db.Object));

        Assert.Null(await store.ConsumeAsync("abc123"));
        db.Verify(
            d => d.StringGetDeleteAsync(
                It.Is<RedisKey>(k => k.ToString() == "jellyfin:sso:oid:state:abc123"),
                It.IsAny<CommandFlags>()),
            Times.Once);
    }

    [Fact]
    public async Task ListAsync_ScansPrimariesOnly()
    {
        var db = new Mock<IDatabase>();
        db.Setup(d => d.StringGetAsync(It.IsAny<RedisKey>(), It.IsAny<CommandFlags>()))
            .ReturnsAsync((RedisValue)SsoStateSerializer.Serialize(NewState("abc123")));

        var primary = CreateServer(isReplica: false, keys: new[] { (RedisKey)"jellyfin:sso:oid:state:abc123" });
        var replica = CreateServer(isReplica: true, keys: new[] { (RedisKey)"jellyfin:sso:oid:state:abc123" });

        var primaryEndpoint = new DnsEndPoint("primary", 6379);
        var replicaEndpoint = new DnsEndPoint("replica", 6379);

        var mux = new Mock<IConnectionMultiplexer>();
        mux.Setup(m => m.GetDatabase(It.IsAny<int>(), It.IsAny<object>())).Returns(db.Object);
        mux.Setup(m => m.GetEndPoints(It.IsAny<bool>())).Returns(new EndPoint[] { primaryEndpoint, replicaEndpoint });
        mux.Setup(m => m.GetServer(primaryEndpoint, It.IsAny<object>())).Returns(primary);
        mux.Setup(m => m.GetServer(replicaEndpoint, It.IsAny<object>())).Returns(replica);

        var states = await new RedisSsoStateStore(mux.Object).ListAsync();

        Assert.Single(states);
        Assert.Equal("abc123", states[0].State.State);
    }

    private static TimedAuthorizeState NewState(string state)
    {
        return new TimedAuthorizeState(
            new AuthorizeState
            {
                State = state,
                StartUrl = "https://idp.example/authorize",
                RedirectUri = "https://jellyfin.example/sso/OID/redirect/authentik",
                CodeVerifier = "code-verifier-secret"
            },
            DateTime.UtcNow);
    }

    private static void SetupStringSet(Mock<IDatabase> db, Action<RedisKey, RedisValue, TimeSpan?> capture)
    {
        db.Setup(d => d.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                It.IsAny<When>(),
                It.IsAny<CommandFlags>()))
            .Callback<RedisKey, RedisValue, TimeSpan?, bool, When, CommandFlags>(
                (key, value, expiry, keepTtl, when, flags) => capture(key, value, expiry))
            .ReturnsAsync(true);
    }

    private static IServer CreateServer(bool isReplica, IReadOnlyList<RedisKey> keys)
    {
        var server = new Mock<IServer>();
        server.SetupGet(s => s.IsConnected).Returns(true);
        server.SetupGet(s => s.IsReplica).Returns(isReplica);
        server.Setup(s => s.KeysAsync(
                It.IsAny<int>(),
                It.IsAny<RedisValue>(),
                It.IsAny<int>(),
                It.IsAny<long>(),
                It.IsAny<int>(),
                It.IsAny<CommandFlags>()))
            .Returns(ToAsyncEnumerable(keys));
        return server.Object;
    }

    private static async IAsyncEnumerable<RedisKey> ToAsyncEnumerable(IReadOnlyList<RedisKey> keys)
    {
        foreach (var key in keys)
        {
            yield return key;
        }

        await Task.CompletedTask;
    }

    private static IConnectionMultiplexer CreateMultiplexer(IDatabase database)
    {
        var mux = new Mock<IConnectionMultiplexer>();
        mux.Setup(m => m.GetDatabase(It.IsAny<int>(), It.IsAny<object>())).Returns(database);
        return mux.Object;
    }
}