jellyfin-ha-src/MediaBrowser.Controller/QuickConnect/InMemoryQuickConnectStore.cs

using System;
using System.Collections.Concurrent;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using MediaBrowser.Controller.Authentication;
using MediaBrowser.Model.QuickConnect;

namespace MediaBrowser.Controller.QuickConnect;

/// <summary>
/// A process-local <see cref="IQuickConnectStore"/>. It is the single-instance default, and the
/// fallback a shared store degrades to while its backend is unreachable, so quick connect keeps
/// working for clients whose three legs happen to land on one instance.
/// </summary>
public sealed class InMemoryQuickConnectStore : IQuickConnectStore
{
    private readonly ConcurrentDictionary<string, Entry<QuickConnectResult>> _requests = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, Entry<AuthenticationResult>> _authorizations = new(StringComparer.Ordinal);

    /// <inheritdoc />
    public Task<QuickConnectResult?> GetRequestBySecretAsync(string secret, CancellationToken cancellationToken = default)
    {
        Expire();
        return Task.FromResult(_requests.TryGetValue(secret, out var entry) ? entry.Value : null);
    }

    /// <inheritdoc />
    public Task<QuickConnectResult?> GetRequestByCodeAsync(string code, CancellationToken cancellationToken = default)
    {
        Expire();
        return Task.FromResult(_requests.Values
            .Select(entry => entry.Value)
            .FirstOrDefault(request => string.Equals(request.Code, code, StringComparison.Ordinal)));
    }

    /// <inheritdoc />
    public Task SetRequestAsync(QuickConnectResult request, DateTime expiresUtc, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);

        Expire();
        _requests[request.Secret] = new Entry<QuickConnectResult>(expiresUtc, request);
        return Task.CompletedTask;
    }

    /// <inheritdoc />
    public Task SetAuthorizationAsync(string secret, AuthenticationResult authenticationResult, DateTime expiresUtc, CancellationToken cancellationToken = default)
    {
        Expire();
        _authorizations[secret] = new Entry<AuthenticationResult>(expiresUtc, authenticationResult);
        return Task.CompletedTask;
    }

    /// <inheritdoc />
    public Task<AuthenticationResult?> TryConsumeAuthorizationAsync(string secret, CancellationToken cancellationToken = default)
    {
        Expire();
        if (!_authorizations.TryRemove(secret, out var entry) || entry.ExpiresUtc <= DateTime.UtcNow)
        {
            return Task.FromResult<AuthenticationResult?>(null);
        }

        return Task.FromResult<AuthenticationResult?>(entry.Value);
    }

    private void Expire()
    {
        var now = DateTime.UtcNow;

        foreach (var (secret, entry) in _requests)
        {
            if (entry.ExpiresUtc <= now)
            {
                _requests.TryRemove(secret, out _);
            }
        }

        foreach (var (secret, entry) in _authorizations)
        {
            if (entry.ExpiresUtc <= now)
            {
                _authorizations.TryRemove(secret, out _);
            }
        }
    }

    private sealed record Entry<T>(DateTime ExpiresUtc, T Value);
}