jellyfin-ha-src/tests/Jellyfin.Server.Tests/Devices/DeviceManagerReplicaTests.cs

using System;
using System.Threading;
using System.Threading.Tasks;
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.Migrations;
using MediaBrowser.Controller.Library;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using Moq;
using Npgsql;
using Xunit;

namespace Jellyfin.Server.Tests.Devices;

/// <summary>
/// Two independently constructed <see cref="DeviceManager"/> instances over one PostgreSQL database are the
/// in-process stand-in for two replicas sharing one database: what either of them writes, the other has to
/// see on its very next read.
/// </summary>
[Trait("Category", "RequiresDocker")]
public sealed class DeviceManagerReplicaTests : IAsyncLifetime
{
    private PostgreSqlTestServer _server = null!;

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

    /// <inheritdoc/>
    public async ValueTask DisposeAsync()
    {
        await _server.DisposeAsync().ConfigureAwait(false);
    }

    /// <summary>
    /// A client that logs in against one replica has to be authenticated by every other replica that was
    /// already running when the token was minted - a rolling update or a scale-up must not 401 it.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task TokenMintedOnOneReplica_IsValidOnAnother()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _server.CreateDatabaseAsync("device_replica_create", cancellationToken);

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

        // Both replicas start before the login, so neither has the device in hand when it is created.
        var replicaA = CreateManager(dataSource, user);
        var replicaB = CreateManager(dataSource, user);

        var device = await replicaA.CreateDevice(new Device(user.Id, "Jellyfin Web", "1.0.0", "Living Room TV", "device-1"));

        var seenByB = await replicaB.GetDevices(new DeviceQuery { AccessToken = device.AccessToken });

        Assert.Equal(device.Id, Assert.Single(seenByB.Items).Id);
        Assert.Equal(user.Id, seenByB.Items[0].UserId);
    }

    /// <summary>
    /// Revocation has to propagate at least as fast as creation: a token logged out on one replica must not
    /// still authenticate on another.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task TokenRevokedOnOneReplica_IsInvalidOnAnother()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _server.CreateDatabaseAsync("device_replica_revoke", cancellationToken);

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

        var replicaA = CreateManager(dataSource, user);
        var device = await replicaA.CreateDevice(new Device(user.Id, "Jellyfin Web", "1.0.0", "Living Room TV", "device-1"));

        // Replica B comes up after the login, so it starts out agreeing that the token is valid.
        var replicaB = CreateManager(dataSource, user);
        Assert.Single((await replicaB.GetDevices(new DeviceQuery { AccessToken = device.AccessToken })).Items);

        await replicaA.DeleteDevice(device);

        Assert.Empty((await replicaB.GetDevices(new DeviceQuery { AccessToken = device.AccessToken })).Items);
    }

    /// <summary>
    /// A device renamed on one replica has to be reported under its new name by the others, and the rename has
    /// to survive being read back through a replica that never saw the write.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task DeviceOptionsWrittenOnOneReplica_AreReadOnAnother()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _server.CreateDatabaseAsync("device_replica_options", cancellationToken);

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

        var replicaA = CreateManager(dataSource, user);
        var replicaB = CreateManager(dataSource, user);

        await replicaA.CreateDevice(new Device(user.Id, "Jellyfin Web", "1.0.0", "Living Room TV", "device-1"));
        await replicaA.UpdateDeviceOptions("device-1", "Kitchen TV");

        var options = await replicaB.GetDeviceOptions("device-1");
        Assert.Equal("Kitchen TV", options?.CustomName);

        var info = await replicaB.GetDevice("device-1");
        Assert.Equal("Kitchen TV", info?.CustomName);
    }

    /// <summary>
    /// Activity written by the replica serving the request has to be visible to the others, because the next
    /// request from the same client can land anywhere.
    /// </summary>
    /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
    [Fact]
    public async Task DeviceUpdatedOnOneReplica_IsReadOnAnother()
    {
        var cancellationToken = TestContext.Current.CancellationToken;
        var connectionString = await _server.CreateDatabaseAsync("device_replica_update", cancellationToken);

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

        var replicaA = CreateManager(dataSource, user);
        var replicaB = CreateManager(dataSource, user);

        var device = await replicaA.CreateDevice(new Device(user.Id, "Jellyfin Web", "1.0.0", "Living Room TV", "device-1"));
        device.AppVersion = "2.0.0";
        await replicaA.UpdateDevice(device);

        var seenByB = Assert.Single((await replicaB.GetDevices(new DeviceQuery { DeviceId = "device-1" })).Items);
        Assert.Equal("2.0.0", seenByB.AppVersion);
    }

    private static DeviceManager CreateManager(NpgsqlDataSource dataSource, User user)
    {
        var userManager = new Mock<IUserManager>();
        userManager.Setup(manager => manager.GetUserById(user.Id)).Returns(user);
        return new DeviceManager(new DataSourceContextFactory(dataSource), userManager.Object);
    }

    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("replica-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));
    }

    /// <summary>
    /// Hands every <see cref="DeviceManager"/> its own context over the one shared database, the way the
    /// pooled factory does in the server.
    /// </summary>
    private sealed class DataSourceContextFactory : IDbContextFactory<JellyfinDbContext>
    {
        private readonly NpgsqlDataSource _dataSource;

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

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