jellyfin-ha-src/tests/Jellyfin.Server.Tests/QuickConnect/QuickConnectReplicaTests.cs

using System;
using System.Collections.Generic;
using System.Globalization;
using System.Threading;
using System.Threading.Tasks;
using Emby.Server.Implementations.QuickConnect;
using Jellyfin.Data.Queries;
using Jellyfin.Database.Implementations;
using Jellyfin.Database.Implementations.DbConfiguration;
using Jellyfin.Database.Implementations.Entities;
using Jellyfin.Database.Implementations.Entities.Security;
using Jellyfin.Database.Implementations.Locking;
using Jellyfin.Database.Providers.PostgreSQL;
using Jellyfin.Server.Implementations.Devices;
using Jellyfin.Server.Tests.HighAvailability;
using Jellyfin.Server.Tests.Migrations;
using MediaBrowser.Common.Extensions;
using MediaBrowser.Controller.Authentication;
using MediaBrowser.Controller.Configuration;
using MediaBrowser.Controller.Devices;
using MediaBrowser.Controller.Library;
using MediaBrowser.Controller.Net;
using MediaBrowser.Controller.QuickConnect;
using MediaBrowser.Controller.Session;
using MediaBrowser.Model.Configuration;
using MediaBrowser.Model.Dto;
using MediaBrowser.Model.QuickConnect;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using Moq;
using Npgsql;
using StackExchange.Redis;
using Xunit;

namespace Jellyfin.Server.Tests.QuickConnect;

/// <summary>
/// Three independently constructed <see cref="QuickConnectManager"/> instances over one PostgreSQL
/// database and one Redis are the in-process stand-in for three replicas without sticky sessions: the
/// initiate, authorize and exchange legs of one flow each land on a different one.
/// </summary>
[Trait("Category", "RequiresDocker")]
public sealed class QuickConnectReplicaTests : IAsyncLifetime
{
    private static readonly AuthorizationInfo _authorizationInfo = new AuthorizationInfo
    {
        Device = "Living Room TV",
        DeviceId = "device-1",
        Client = "Jellyfin Web",
        Version = "1.0.0"
    };

    private readonly List<IConnectionMultiplexer> _connections = new();

    private PostgreSqlTestServer _postgres = null!;
    private RedisTestServer _redis = null!;

    /// <inheritdoc/>
    public async ValueTask InitializeAsync()
    {
        _postgres = await PostgreSqlTestServer.StartAsync().ConfigureAwait(false);
        _redis = await RedisTestServer.StartAsync().ConfigureAwait(false);
    }

    /// <inheritdoc/>
    public async ValueTask DisposeAsync()
    {
        foreach (var connection in _connections)
        {
            await connection.DisposeAsync().ConfigureAwait(false);
        }

        await _redis.DisposeAsync().ConfigureAwait(false);
        await _postgres.DisposeAsync().ConfigureAwait(false);
    }

    /// <summary>
    /// The three legs of a quick connect flow land on three different replicas, and the token the third
    /// one hands out is the one the second one minted into the shared database.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task InitiateAuthorizeExchange_AcrossThreeReplicas_Succeeds()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _postgres.CreateDatabaseAsync("quickconnect_replica_flow", cancellationToken);

        await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        var user = await CreateSchemaWithUserAsync(dataSource, cancellationToken);

        var replicaA = await CreateReplicaAsync(dataSource, user);
        var replicaB = await CreateReplicaAsync(dataSource, user);
        var replicaC = await CreateReplicaAsync(dataSource, user);

        var initiated = await replicaA.Manager.TryConnect(_authorizationInfo);

        // The code is shown to the user on whichever replica serves the dashboard.
        Assert.True(await replicaB.Manager.AuthorizeRequest(user.Id, initiated.Code));

        var polled = await replicaC.Manager.CheckRequestStatus(initiated.Secret);
        Assert.True(polled.Authenticated);
        Assert.Equal(initiated.Code, polled.Code);
        Assert.Equal(_authorizationInfo.DeviceId, polled.DeviceId);

        var exchanged = await replicaC.Manager.GetAuthorizedRequest(initiated.Secret);

        Assert.False(string.IsNullOrEmpty(exchanged.AccessToken));
        Assert.Equal(user.Id, exchanged.User.Id);

        var devices = await replicaA.Devices.GetDevices(new DeviceQuery { AccessToken = exchanged.AccessToken });
        Assert.Equal(user.Id, Assert.Single(devices.Items).UserId);
    }

    /// <summary>
    /// Exchanging a secret does not spend it: a client that retries, or whose retry lands on another
    /// replica, gets the same access token back rather than a 404, and the device is minted once.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task Exchange_RepeatedOnTwoReplicas_ReturnsTheSameToken()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _postgres.CreateDatabaseAsync("quickconnect_replica_reexchange", cancellationToken);

        await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        var user = await CreateSchemaWithUserAsync(dataSource, cancellationToken);

        var replicaA = await CreateReplicaAsync(dataSource, user);
        var replicaB = await CreateReplicaAsync(dataSource, user);
        var replicaC = await CreateReplicaAsync(dataSource, user);

        for (var attempt = 0; attempt < 10; attempt++)
        {
            var authorizationInfo = AuthorizationInfoFor(attempt);
            var initiated = await replicaA.Manager.TryConnect(authorizationInfo);
            await replicaB.Manager.AuthorizeRequest(user.Id, initiated.Code);

            var exchanged = await Task.WhenAll(
                Task.Run(() => ExchangeAsync(replicaA.Manager, initiated.Secret), cancellationToken),
                Task.Run(() => ExchangeAsync(replicaC.Manager, initiated.Secret), cancellationToken));

            Assert.All(exchanged, outcome => Assert.NotNull(outcome));
            Assert.Equal(exchanged[0]!.AccessToken, exchanged[1]!.AccessToken);

            // Still there afterwards, on a replica that has not exchanged it yet.
            var later = await replicaB.Manager.GetAuthorizedRequest(initiated.Secret);
            Assert.Equal(exchanged[0]!.AccessToken, later.AccessToken);

            var devices = await replicaA.Devices.GetDevices(new DeviceQuery { DeviceId = authorizationInfo.DeviceId });
            Assert.Equal(later.AccessToken, Assert.Single(devices.Items).AccessToken);
        }
    }

    /// <summary>
    /// Two replicas authorizing one code at the same time mint one access token between them. A second
    /// one would be live, attached to the same device and reachable by nobody.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task Authorize_RacedOnTwoReplicas_MintsOneAccessToken()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _postgres.CreateDatabaseAsync("quickconnect_replica_authorize_race", cancellationToken);

        await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        var user = await CreateSchemaWithUserAsync(dataSource, cancellationToken);

        var replicaA = await CreateReplicaAsync(dataSource, user);
        var replicaB = await CreateReplicaAsync(dataSource, user);
        var replicaC = await CreateReplicaAsync(dataSource, user);

        for (var attempt = 0; attempt < 20; attempt++)
        {
            var authorizationInfo = AuthorizationInfoFor(attempt);
            var initiated = await replicaA.Manager.TryConnect(authorizationInfo);

            var outcomes = await Task.WhenAll(
                Task.Run(() => AuthorizeAsync(replicaB.Manager, user.Id, initiated.Code), cancellationToken),
                Task.Run(() => AuthorizeAsync(replicaC.Manager, user.Id, initiated.Code), cancellationToken));

            Assert.Single(outcomes, authorized => authorized);

            var devices = await replicaA.Devices.GetDevices(new DeviceQuery { DeviceId = authorizationInfo.DeviceId });
            var device = Assert.Single(devices.Items);

            var exchanged = await replicaA.Manager.GetAuthorizedRequest(initiated.Secret);
            Assert.Equal(device.AccessToken, exchanged.AccessToken);
        }
    }

    /// <summary>
    /// An expired request is rejected on a replica that never saw it created, rather than resolving to a
    /// stale authorization.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task ExpiredRequest_IsRejectedOnEveryReplica()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _postgres.CreateDatabaseAsync("quickconnect_replica_expiry", cancellationToken);

        await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        var user = await CreateSchemaWithUserAsync(dataSource, cancellationToken);

        var replicaA = await CreateReplicaAsync(dataSource, user);
        var replicaB = await CreateReplicaAsync(dataSource, user);

        var initiated = await replicaA.Manager.TryConnect(_authorizationInfo);
        Assert.NotNull(await replicaB.Manager.CheckRequestStatus(initiated.Secret));

        // Shorten the stored expiry instead of waiting out the ten minute timeout.
        await replicaA.Store.SetRequestAsync(initiated, DateTime.UtcNow.AddSeconds(1), cancellationToken);
        await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken);

        await Assert.ThrowsAsync<ResourceNotFoundException>(() => replicaB.Manager.CheckRequestStatus(initiated.Secret));
        await Assert.ThrowsAsync<ResourceNotFoundException>(() => replicaB.Manager.AuthorizeRequest(user.Id, initiated.Code));
    }

    /// <summary>
    /// An authorization that was never exchanged expires too, so a code authorized and then abandoned
    /// cannot be redeemed later from another replica.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task ExpiredAuthorization_IsRejectedOnEveryReplica()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _postgres.CreateDatabaseAsync("quickconnect_replica_auth_expiry", cancellationToken);

        await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).Build();
        var user = await CreateSchemaWithUserAsync(dataSource, cancellationToken);

        var replicaA = await CreateReplicaAsync(dataSource, user);
        var replicaB = await CreateReplicaAsync(dataSource, user);

        var initiated = await replicaA.Manager.TryConnect(_authorizationInfo);
        await replicaA.Manager.AuthorizeRequest(user.Id, initiated.Code);

        var stored = await replicaA.Store.GetRequestBySecretAsync(initiated.Secret, cancellationToken);
        Assert.True(stored?.Authenticated);

        await replicaA.Store.SetAuthorizationAsync(
            initiated.Secret,
            new AuthenticationResult { AccessToken = "stale" },
            DateTime.UtcNow.AddSeconds(1),
            cancellationToken);
        await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken);

        await Assert.ThrowsAsync<ResourceNotFoundException>(() => replicaB.Manager.GetAuthorizedRequest(initiated.Secret));
    }

    private static AuthorizationInfo AuthorizationInfoFor(int attempt) => new AuthorizationInfo
    {
        Device = _authorizationInfo.Device,
        DeviceId = string.Create(CultureInfo.InvariantCulture, $"device-{attempt}"),
        Client = _authorizationInfo.Client,
        Version = _authorizationInfo.Version
    };

    private static async Task<bool> AuthorizeAsync(IQuickConnect manager, Guid userId, string code)
    {
        try
        {
            return await manager.AuthorizeRequest(userId, code).ConfigureAwait(false);
        }
        catch (ConflictException)
        {
            return false;
        }
    }

    private static async Task<AuthenticationResult?> ExchangeAsync(IQuickConnect manager, string secret)
    {
        try
        {
            return await manager.GetAuthorizedRequest(secret).ConfigureAwait(false);
        }
        catch (ResourceNotFoundException)
        {
            return null;
        }
    }

    private static async Task<User> CreateSchemaWithUserAsync(NpgsqlDataSource dataSource, CancellationToken cancellationToken)
    {
        var context = CreateContext(dataSource);
        await using (context.ConfigureAwait(false))
        {
            await context.Database.EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);

            var user = new User("quickconnect-user", "provider", "provider");
            context.Users.Add(user);
            await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

            return user;
        }
    }

    private static JellyfinDbContext CreateContext(NpgsqlDataSource dataSource)
    {
        var optionsBuilder = new DbContextOptionsBuilder<JellyfinDbContext>();
        var provider = new PostgreSqlDatabaseProvider(dataSource);
        provider.Initialise(optionsBuilder, new DatabaseConfigurationOptions { DatabaseType = "PostgreSQL" });
        return new JellyfinDbContext(
            optionsBuilder.Options,
            NullLogger<JellyfinDbContext>.Instance,
            provider,
            new NoLockBehavior(NullLogger<NoLockBehavior>.Instance));
    }

    private async Task<Replica> CreateReplicaAsync(NpgsqlDataSource dataSource, User user)
    {
        var connection = await _redis.ConnectAsync().ConfigureAwait(false);
        _connections.Add(connection);

        var userManager = new Mock<IUserManager>();
        userManager.Setup(manager => manager.GetUserById(user.Id)).Returns(user);
        var deviceManager = new DeviceManager(new DataSourceContextFactory(dataSource), userManager.Object);

        var configManager = new Mock<IServerConfigurationManager>();
        configManager.Setup(manager => manager.Configuration).Returns(new ServerConfiguration { QuickConnectAvailable = true });

        // Stands in for SessionManager.AuthenticateDirect: the token has to be minted into the shared
        // database, because the replica that exchanges the secret is not the one that authorized it.
        var sessionManager = new Mock<ISessionManager>();
        sessionManager
            .Setup(manager => manager.AuthenticateDirect(It.IsAny<AuthenticationRequest>()))
            .Returns<AuthenticationRequest>(async request =>
            {
                var device = await deviceManager.CreateDevice(
                    new Device(request.UserId, request.App, request.AppVersion, request.DeviceName, request.DeviceId)).ConfigureAwait(false);

                return new AuthenticationResult
                {
                    AccessToken = device.AccessToken,
                    ServerId = "server-1",
                    User = new UserDto { Id = user.Id, Name = user.Username, ServerId = "server-1" },
                    SessionInfo = new SessionInfoDto
                    {
                        Id = device.Id.ToString(CultureInfo.InvariantCulture),
                        UserId = user.Id,
                        UserName = user.Username,
                        Client = request.App,
                        DeviceId = request.DeviceId,
                        DeviceName = request.DeviceName,
                        ApplicationVersion = request.AppVersion
                    }
                };
            });

        var store = new RedisQuickConnectStore(connection, NullLogger<RedisQuickConnectStore>.Instance);
        var manager = new QuickConnectManager(
            configManager.Object,
            NullLogger<QuickConnectManager>.Instance,
            sessionManager.Object,
            store);

        return new Replica(manager, store, deviceManager);
    }

    private sealed record Replica(IQuickConnect Manager, IQuickConnectStore Store, IDeviceManager Devices);

    private sealed class DataSourceContextFactory : IDbContextFactory<JellyfinDbContext>
    {
        private readonly NpgsqlDataSource _dataSource;

        public DataSourceContextFactory(NpgsqlDataSource dataSource)
        {
            _dataSource = dataSource;
        }

        public JellyfinDbContext CreateDbContext() => CreateContext(_dataSource);
    }
}