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

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
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
{
    /// <summary>
    /// The most states this store keeps. The endpoint that creates them is unauthenticated, so the store
    /// is capped as well as swept; past the cap the states closest to expiring are dropped first.
    /// Every state carries the same expiry, so that is the oldest state: the cap trades unbounded growth
    /// for a login DoS, where enough anonymous starts evict the in-flight logins of real users.
    /// </summary>
    public const int Capacity = 1024;

    private readonly ConcurrentDictionary<string, Entry> _entries = new ConcurrentDictionary<string, Entry>(StringComparer.Ordinal);

    /// <summary>
    /// Gets the number of states currently held.
    /// </summary>
    public int Count => _entries.Count;

    /// <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));
        Evict();
        return Task.CompletedTask;
    }

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

        // ConcurrentDictionary.ToArray() is a locked snapshot; LINQ buffering throws if a write races it.
        var snapshot = _entries.ToArray();
        var excess = snapshot.Length - Capacity;
        if (excess <= 0)
        {
            return;
        }

        foreach (var pair in snapshot.OrderBy(p => p.Value.ExpiresAt).Take(excess))
        {
            _entries.TryRemove(pair.Key, out _);
        }
    }

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