jellyfin-plugin-sso/SSO-Auth/State/InMemorySsoStateStore.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading.Tasks;
using Jellyfin.Plugin.SSO_Auth.Api;

namespace Jellyfin.Plugin.SSO_Auth.State;

/// <summary>
/// A process-local <see cref="ISsoStateStore"/> used when the server has no shared cache configured.
/// States are held serialized so that this store and <see cref="RedisSsoStateStore"/> behave identically:
/// a mutation is only visible to a later leg of the flow once it has been written back.
/// </summary>
public sealed class InMemorySsoStateStore : ISsoStateStore
{
    private readonly ConcurrentDictionary<string, Entry> _entries = new ConcurrentDictionary<string, Entry>(StringComparer.Ordinal);

    /// <inheritdoc />
    public Task CreateAsync(TimedAuthorizeState state, TimeSpan expiry)
    {
        return WriteAsync(state, expiry);
    }

    /// <inheritdoc />
    public Task UpdateAsync(TimedAuthorizeState state, TimeSpan expiry)
    {
        return WriteAsync(state, expiry);
    }

    /// <inheritdoc />
    public Task<TimedAuthorizeState> GetAsync(string state)
    {
        if (string.IsNullOrEmpty(state) || !_entries.TryGetValue(state, out var entry))
        {
            return Task.FromResult<TimedAuthorizeState>(null);
        }

        if (IsExpired(entry))
        {
            _entries.TryRemove(state, out _);
            return Task.FromResult<TimedAuthorizeState>(null);
        }

        return Task.FromResult(SsoStateSerializer.Deserialize(entry.Payload));
    }

    /// <inheritdoc />
    public Task<TimedAuthorizeState> ConsumeAsync(string state)
    {
        if (string.IsNullOrEmpty(state) || !_entries.TryRemove(state, out var entry) || IsExpired(entry))
        {
            return Task.FromResult<TimedAuthorizeState>(null);
        }

        return Task.FromResult(SsoStateSerializer.Deserialize(entry.Payload));
    }

    /// <inheritdoc />
    public Task<IReadOnlyList<TimedAuthorizeState>> ListAsync()
    {
        var states = new List<TimedAuthorizeState>();

        foreach (var pair in _entries)
        {
            if (IsExpired(pair.Value))
            {
                _entries.TryRemove(pair.Key, out _);
                continue;
            }

            var parsed = SsoStateSerializer.Deserialize(pair.Value.Payload);
            if (parsed is not null)
            {
                states.Add(parsed);
            }
        }

        return Task.FromResult<IReadOnlyList<TimedAuthorizeState>>(states);
    }

    private static bool IsExpired(Entry entry)
    {
        return entry.ExpiresAt <= DateTimeOffset.UtcNow;
    }

    private Task WriteAsync(TimedAuthorizeState state, TimeSpan expiry)
    {
        ArgumentNullException.ThrowIfNull(state);

        _entries[state.State.State] = new Entry(SsoStateSerializer.Serialize(state), DateTimeOffset.UtcNow.Add(expiry));
        return Task.CompletedTask;
    }

    private sealed record Entry(string Payload, DateTimeOffset ExpiresAt);
}