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, an authorization is
/// claimed with a Lua check-and-set and consumed with <c>GETDEL</c>, so only one instance can ever mint
/// a given secret's access token and only one can ever hand it out.
/// </summary>
/// <remarks>
/// A pending request survives an unreachable Redis through a process-local fallback, because a second
/// copy of it is harmless. An authorization has none: a second copy of it is a second access token, and
/// a write whose response timed out may well have been applied, so a transport failure on that path is
/// surfaced rather than degraded.
/// </remarks>
public sealed class RedisQuickConnectStore : IQuickConnectStore
{
    private const string KeyPrefix = "jellyfin:quickconnect:";

    /// <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 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)
    {
        RedisValue raw;
        try
        {
            raw = await _db.StringGetAsync(RequestKey(secret)).ConfigureAwait(false);
        }
        catch (Exception ex) when (IsTransportFailure(ex))
        {
            LogDegraded(ex);
            return await _fallback.GetRequestBySecretAsync(secret, cancellationToken).ConfigureAwait(false);
        }

        // A miss is an answer rather than a transport failure, so the fallback is not consulted for it.
        return raw.HasValue ? JsonSerializer.Deserialize<QuickConnectResult>(raw.ToString(), JsonDefaults.Options) : null;
    }

    /// <inheritdoc />
    public async Task<QuickConnectResult?> GetRequestByCodeAsync(string code, CancellationToken cancellationToken = default)
    {
        RedisValue secret;
        try
        {
            secret = await _db.StringGetAsync(CodeKey(code)).ConfigureAwait(false);
        }
        catch (Exception ex) when (IsTransportFailure(ex))
        {
            LogDegraded(ex);
            return await _fallback.GetRequestByCodeAsync(code, cancellationToken).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;
        }

        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) when (IsTransportFailure(ex))
        {
            LogDegraded(ex);
            await _fallback.SetRequestAsync(request, expiresUtc, cancellationToken).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 _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 _db.StringSetAsync(AuthorizationKey(secret), json, ttl).ConfigureAwait(false);
    }

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

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

    // Deliberately excludes a malformed stored value, which is a fault of its own rather than a reason
    // to answer from this instance.
    private static bool IsTransportFailure(Exception exception) => exception is RedisException or TimeoutException;

    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 void LogDegraded(Exception exception)
        => _logger.LogWarning(exception, "Quick connect request state could not be shared through Redis; falling back to this instance only.");
}