jellyfin-ha-src/Emby.Server.Implementations/QuickConnect/RedisQuickConnectStore.cs

using System;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Jellyfin.Extensions.Json;
using MediaBrowser.Controller.Authentication;
using MediaBrowser.Controller.QuickConnect;
using MediaBrowser.Model.QuickConnect;
using Microsoft.Extensions.Logging;
using StackExchange.Redis;

namespace Emby.Server.Implementations.QuickConnect;

/// <summary>
/// A Redis-backed <see cref="IQuickConnectStore"/> that lets the initiate, authorize and exchange legs
/// of a quick connect flow land on different instances. Expiry is the key TTL, and an authorization is
/// consumed with <c>GETDEL</c> so only one instance can ever hand out a given secret's access token.
/// </summary>
public sealed class RedisQuickConnectStore : IQuickConnectStore
{
    private const string KeyPrefix = "jellyfin:quickconnect:";

    private readonly IDatabase _db;
    private readonly InMemoryQuickConnectStore _fallback;
    private readonly ILogger<RedisQuickConnectStore> _logger;

    /// <summary>
    /// Initializes a new instance of the <see cref="RedisQuickConnectStore"/> class.
    /// </summary>
    /// <param name="redis">The Redis connection multiplexer.</param>
    /// <param name="logger">The logger.</param>
    public RedisQuickConnectStore(IConnectionMultiplexer redis, ILogger<RedisQuickConnectStore> logger)
    {
        ArgumentNullException.ThrowIfNull(redis);

        _db = redis.GetDatabase();
        _fallback = new InMemoryQuickConnectStore();
        _logger = logger;
    }

    /// <inheritdoc />
    public async Task<QuickConnectResult?> GetRequestBySecretAsync(string secret, CancellationToken cancellationToken = default)
    {
        try
        {
            var raw = await _db.StringGetAsync(RequestKey(secret)).ConfigureAwait(false);
            if (raw.HasValue)
            {
                return JsonSerializer.Deserialize<QuickConnectResult>(raw.ToString(), JsonDefaults.Options);
            }
        }
        catch (Exception ex)
        {
            LogDegraded(ex);
        }

        return await _fallback.GetRequestBySecretAsync(secret, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public async Task<QuickConnectResult?> GetRequestByCodeAsync(string code, CancellationToken cancellationToken = default)
    {
        try
        {
            var secret = await _db.StringGetAsync(CodeKey(code)).ConfigureAwait(false);
            if (secret.HasValue)
            {
                return await GetRequestBySecretAsync(secret.ToString(), cancellationToken).ConfigureAwait(false);
            }
        }
        catch (Exception ex)
        {
            LogDegraded(ex);
        }

        return await _fallback.GetRequestByCodeAsync(code, cancellationToken).ConfigureAwait(false);
    }

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

        var ttl = expiresUtc - DateTime.UtcNow;
        if (ttl <= TimeSpan.Zero)
        {
            return;
        }

        try
        {
            var json = JsonSerializer.Serialize(request, JsonDefaults.Options);
            await _db.StringSetAsync(RequestKey(request.Secret), json, ttl).ConfigureAwait(false);
            await _db.StringSetAsync(CodeKey(request.Code), request.Secret, ttl).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            LogDegraded(ex);
            await _fallback.SetRequestAsync(request, expiresUtc, cancellationToken).ConfigureAwait(false);
        }
    }

    /// <inheritdoc />
    public async Task SetAuthorizationAsync(string secret, AuthenticationResult authenticationResult, DateTime expiresUtc, CancellationToken cancellationToken = default)
    {
        var ttl = expiresUtc - DateTime.UtcNow;
        if (ttl <= TimeSpan.Zero)
        {
            return;
        }

        try
        {
            var json = JsonSerializer.Serialize(authenticationResult, JsonDefaults.Options);
            await _db.StringSetAsync(AuthorizationKey(secret), json, ttl).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            LogDegraded(ex);
            await _fallback.SetAuthorizationAsync(secret, authenticationResult, expiresUtc, cancellationToken).ConfigureAwait(false);
        }
    }

    /// <inheritdoc />
    public async Task<AuthenticationResult?> TryConsumeAuthorizationAsync(string secret, CancellationToken cancellationToken = default)
    {
        try
        {
            var raw = await _db.StringGetDeleteAsync(AuthorizationKey(secret)).ConfigureAwait(false);
            if (raw.HasValue)
            {
                return JsonSerializer.Deserialize<AuthenticationResult>(raw.ToString(), JsonDefaults.Options);
            }
        }
        catch (Exception ex)
        {
            LogDegraded(ex);
        }

        return await _fallback.TryConsumeAuthorizationAsync(secret, cancellationToken).ConfigureAwait(false);
    }

    private static string RequestKey(string secret) => KeyPrefix + "request:" + secret;

    private static string CodeKey(string code) => KeyPrefix + "code:" + code;

    private static string AuthorizationKey(string secret) => KeyPrefix + "auth:" + secret;

    private void LogDegraded(Exception exception)
        => _logger.LogWarning(exception, "Quick connect state could not be shared through Redis; falling back to this instance only.");
}