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.Common.Extensions;
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
/// claimed with a Lua check-and-set, so only one instance can ever mint a given secret's access token.
/// </summary>
/// <remarks>
/// There is no local fallback: a call Redis did not answer is inconclusive, and reporting it as a miss
/// would tell a polling client its secret is invalid. Quick connect is unavailable for as long as Redis
/// is, which password login is not.
/// </remarks>
public sealed class RedisQuickConnectStore : IQuickConnectStore
{
    private const string KeyPrefix = "jellyfin:quickconnect:";

    /// <summary>
    /// Lua script writing the two keys a request is resolvable by in one step, so it can never be
    /// reachable by its secret while the code the user is reading off the screen resolves to nothing.
    /// </summary>
    private const string SetRequestScript = @"
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[3])
redis.call('SET', KEYS[2], ARGV[2], 'PX', ARGV[3])
return 1";

    /// <summary>
    /// Lua script for the atomic claim of the sole right to authorize a request: the request has to
    /// exist and not already be authorized, and the claim marker is taken with <c>SET NX</c>, so of two
    /// instances racing on one code exactly one goes on to mint an access token.
    /// </summary>
    private const string ClaimAuthorizationScript = @"
local raw = redis.call('GET', KEYS[1])
if not raw then return 0 end
if cjson.decode(raw)['Authenticated'] then return 0 end
if redis.call('SET', KEYS[2], '1', 'NX', 'PX', ARGV[1]) then return 1 end
return 0";

    private readonly IDatabase _db;
    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();
        _logger = logger;
    }

    /// <inheritdoc />
    public async Task<QuickConnectResult?> GetRequestBySecretAsync(string secret, CancellationToken cancellationToken = default)
    {
        var raw = await CallAsync(() => _db.StringGetAsync(RequestKey(secret))).ConfigureAwait(false);

        // Deserialization is outside the guard: a malformed stored value is a fault of its own, not Redis
        // being unavailable.
        return raw.HasValue ? JsonSerializer.Deserialize<QuickConnectResult>(raw.ToString(), JsonDefaults.Options) : null;
    }

    /// <inheritdoc />
    public async Task<QuickConnectResult?> GetRequestByCodeAsync(string code, CancellationToken cancellationToken = default)
    {
        var secret = await CallAsync(() => _db.StringGetAsync(CodeKey(code))).ConfigureAwait(false);

        return secret.HasValue
            ? await GetRequestBySecretAsync(secret.ToString(), cancellationToken).ConfigureAwait(false)
            : null;
    }

    /// <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;
        }

        var json = JsonSerializer.Serialize(request, JsonDefaults.Options);
        await CallAsync(() => _db.ScriptEvaluateAsync(
            SetRequestScript,
            keys: new RedisKey[] { RequestKey(request.Secret), CodeKey(request.Code) },
            values: new RedisValue[] { json, request.Secret, (long)ttl.TotalMilliseconds })).ConfigureAwait(false);
    }

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

        var claimed = (long?)await CallAsync(() => _db.ScriptEvaluateAsync(
            ClaimAuthorizationScript,
            keys: new RedisKey[] { RequestKey(secret), ClaimKey(secret) },
            values: new RedisValue[] { (long)ttl.TotalMilliseconds })).ConfigureAwait(false);

        return claimed == 1;
    }

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

        var json = JsonSerializer.Serialize(authenticationResult, JsonDefaults.Options);
        await CallAsync(() => _db.StringSetAsync(AuthorizationKey(secret), json, ttl)).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public async Task<AuthenticationResult?> GetAuthorizationAsync(string secret, CancellationToken cancellationToken = default)
    {
        var raw = await CallAsync(() => _db.StringGetAsync(AuthorizationKey(secret))).ConfigureAwait(false);

        return raw.HasValue
            ? JsonSerializer.Deserialize<AuthenticationResult>(raw.ToString(), JsonDefaults.Options)
            : null;
    }

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

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

    private static string ClaimKey(string secret) => KeyPrefix + "claim:" + secret;

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

    private async Task<T> CallAsync<T>(Func<Task<T>> call)
    {
        try
        {
            return await call().ConfigureAwait(false);
        }
        catch (Exception exception) when (exception is RedisException or RedisCommandException or TimeoutException)
        {
            _logger.LogError(exception, "Quick connect state could not be reached in Redis.");
            throw new ServiceUnavailableException("Quick connect is temporarily unavailable.", exception);
        }
    }
}