Disconnect existing client connections when another client instance for same user is detected

SampleMultiplayerClient/MultiplayerClient.cs

@@ -7,6 +7,7 @@
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Client;
+using osu.Game.Online;
 using osu.Game.Online.API;
 using osu.Game.Online.Multiplayer;
 using osu.Game.Online.Rooms;
@@ -45,6 +46,7 @@ public MultiplayerClient(HubConnection connection, int userId)
             connection.On<MultiplayerPlaylistItem>(nameof(IMultiplayerClient.PlaylistItemChanged), ((IMultiplayerClient)this).PlaylistItemChanged);
             connection.On<long>(nameof(IMultiplayerClient.PlaylistItemRemoved), ((IMultiplayerClient)this).PlaylistItemRemoved);
             connection.On<int, long, string>(nameof(IMultiplayerClient.Invited), ((IMultiplayerClient)this).Invited);
+            connection.On(nameof(IStatefulUserHubClient.DisconnectRequested), ((IStatefulUserHubClient)this).DisconnectRequested);
         }
 
         public MultiplayerUserState State { get; private set; }
@@ -250,5 +252,11 @@ public Task PlaylistItemChanged(MultiplayerPlaylistItem item)
             Console.WriteLine($"Playlist item changed (id: {item.ID} beatmap: {item.BeatmapID}, ruleset: {item.RulesetID})");
             return Task.CompletedTask;
         }
+
+        public async Task DisconnectRequested()
+        {
+            Console.WriteLine("Disconnect requested");
+            await LeaveRoom();
+        }
     }
 }

SampleSpectatorClient/SpectatorClient.cs

@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Client;
+using osu.Game.Online;
 using osu.Game.Online.Spectator;
 
 namespace SampleSpectatorClient
@@ -26,6 +27,7 @@ public SpectatorClient(HubConnection connection)
             connection.On<int, FrameDataBundle>(nameof(ISpectatorClient.UserSentFrames), ((ISpectatorClient)this).UserSentFrames);
             connection.On<int, SpectatorState>(nameof(ISpectatorClient.UserFinishedPlaying), ((ISpectatorClient)this).UserFinishedPlaying);
             connection.On<int, long>(nameof(ISpectatorClient.UserScoreProcessed), ((ISpectatorClient)this).UserScoreProcessed);
+            connection.On(nameof(IStatefulUserHubClient.DisconnectRequested), ((IStatefulUserHubClient)this).DisconnectRequested);
         }
 
         Task ISpectatorClient.UserBeganPlaying(int userId, SpectatorState state)
@@ -69,5 +71,11 @@ Task ISpectatorClient.UserScoreProcessed(int userId, long scoreId)
         public Task EndPlaying(SpectatorState state) => connection.SendAsync(nameof(ISpectatorServer.EndPlaySession), state);
 
         public Task WatchUser(int userId) => connection.SendAsync(nameof(ISpectatorServer.StartWatchingUser), userId);
+
+        public Task DisconnectRequested()
+        {
+            Console.WriteLine($"{connection.ConnectionId} Disconnect requested");
+            return Task.CompletedTask;
+        }
     }
 }

osu.Server.Spectator.Tests/Multiplayer/DelegatingMultiplayerClient.cs

@@ -141,5 +141,11 @@ public virtual async Task PlaylistItemChanged(MultiplayerPlaylistItem item)
         {
             return (Task)GetType().GetMethod(method, BindingFlags.Instance | BindingFlags.Public)!.Invoke(this, args)!;
         }
+
+        public async Task DisconnectRequested()
+        {
+            foreach (var c in Clients)
+                await c.DisconnectRequested();
+        }
     }
 }

osu.Server.Spectator.Tests/StatefulUserHubTest.cs

@@ -9,7 +9,9 @@
 using Microsoft.Extensions.Caching.Memory;
 using Microsoft.Extensions.Options;
 using Moq;
+using osu.Game.Online;
 using osu.Server.Spectator.Entities;
+using osu.Server.Spectator.Extensions;
 using osu.Server.Spectator.Hubs;
 using Xunit;
 
@@ -104,7 +106,7 @@ public async Task SameUserOldConnectionDoesntDestroyNewState()
         private void setNewConnectionId(string? connectionId = null) =>
             mockContext.Setup(context => context.ConnectionId).Returns(connectionId ?? Guid.NewGuid().ToString());
 
-        private class TestStatefulHub : StatefulUserHub<object, ClientState>
+        private class TestStatefulHub : StatefulUserHub<IStatefulUserHubClient, ClientState>
         {
             public TestStatefulHub(IDistributedCache cache, EntityStore<ClientState> userStates)
                 : base(cache, userStates)
@@ -114,7 +116,7 @@ public TestStatefulHub(IDistributedCache cache, EntityStore<ClientState> userSta
             public async Task CreateUserState()
             {
                 using (var state = await GetOrCreateLocalUserState())
-                    state.Item = new ClientState(Context.ConnectionId, CurrentContextUserId);
+                    state.Item = new ClientState(Context.ConnectionId, Context.GetUserId());
             }
         }
     }

osu.Server.Spectator/ConcurrentConnectionLimiter.cs

@@ -0,0 +1,121 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.SignalR;
+using Microsoft.Extensions.DependencyInjection;
+using osu.Framework.Extensions.TypeExtensions;
+using osu.Framework.Logging;
+using osu.Game.Online;
+using osu.Game.Online.Multiplayer;
+using osu.Server.Spectator.Entities;
+using osu.Server.Spectator.Extensions;
+using osu.Server.Spectator.Hubs;
+
+namespace osu.Server.Spectator
+{
+    public class ConcurrentConnectionLimiter : IHubFilter
+    {
+        private readonly EntityStore<ConnectionState> connectionStates;
+
+        private readonly IServiceProvider serviceProvider;
+
+        private static readonly IEnumerable<Type> stateful_user_hubs
+            = typeof(IStatefulUserHub).Assembly.GetTypes().Where(type => typeof(IStatefulUserHub).IsAssignableFrom(type) && typeof(Hub).IsAssignableFrom(type) && !type.IsInterface && !type.IsAbstract).ToArray();
+
+        public ConcurrentConnectionLimiter(
+            EntityStore<ConnectionState> connectionStates,
+            IServiceProvider serviceProvider)
+        {
+            this.connectionStates = connectionStates;
+            this.serviceProvider = serviceProvider;
+        }
+
+        public async Task OnConnectedAsync(HubLifetimeContext context, Func<HubLifetimeContext, Task> next)
+        {
+            try
+            {
+                var userId = context.Context.GetUserId();
+
+                using (var userState = await connectionStates.GetForUse(userId, true))
+                {
+                    if (context.Context.GetTokenId() == userState.Item?.TokenId)
+                    {
+                        log(context, "subsequent connection from same client instance, registering");
+                        userState.Item.RegisterConnectionId(context);
+                        return;
+                    }
+
+                    if (userState.Item != null)
+                    {
+                        log(context, "connection from new client instance, dropping existing state");
+
+                        foreach (var hub in stateful_user_hubs)
+                        {
+                            var hubContextType = typeof(IHubContext<>).MakeGenericType(hub);
+                            var hubContext = serviceProvider.GetRequiredService(hubContextType) as IHubContext;
+                            hubContext?.Clients.Client(userState.Item.ConnectionIds[hub])
+                                      .SendCoreAsync(nameof(IStatefulUserHubClient.DisconnectRequested), Array.Empty<object>());
+                        }
+
+                        log(context, "existing state dropped");
+                    }
+                    else
+                        log(context, "connection from first client instance");
+
+                    userState.Item = new ConnectionState(context);
+                }
+            }
+            finally
+            {
+                await next(context);
+            }
+        }
+
+        private static void log(HubLifetimeContext context, string message)
+            => Logger.Log($"[user:{context.Context.GetUserId()}] [connection:{context.Context.ConnectionId}] [hub:{context.Hub.GetType().ReadableName()}] {message}");
+
+        public async ValueTask<object?> InvokeMethodAsync(HubInvocationContext invocationContext, Func<HubInvocationContext, ValueTask<object?>> next)
+        {
+            // TODO: allow things to execute for hubs that aren't exclusive (like metadata or whatever)
+
+            var userId = invocationContext.Context.GetUserId();
+
+            using (var userState = await connectionStates.GetForUse(userId))
+            {
+                if (invocationContext.Context.GetTokenId() != userState.Item?.TokenId)
+                    throw new InvalidStateException("State is not valid for this connection");
+            }
+
+            return await next(invocationContext);
+        }
+
+        public async Task OnDisconnectedAsync(HubLifetimeContext context, Exception? exception, Func<HubLifetimeContext, Exception?, Task> next)
+        {
+            try
+            {
+                if (exception != null)
+                    // network disconnection. wait for user to return.
+                    return;
+
+                var userId = context.Context.GetUserId();
+
+                using (var userState = await connectionStates.GetForUse(userId, true))
+                {
+                    if (userState.Item?.TokenId == context.Context.GetTokenId())
+                    {
+                        log(context, "disconnected");
+                        userState.Destroy();
+                    }
+                }
+            }
+            finally
+            {
+                await next(context, exception);
+            }
+        }
+    }
+}

osu.Server.Spectator/Entities/ConnectionState.cs

@@ -0,0 +1,38 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System;
+using System.Collections.Generic;
+using Microsoft.AspNetCore.SignalR;
+using osu.Server.Spectator.Extensions;
+
+namespace osu.Server.Spectator.Entities
+{
+    public class ConnectionState
+    {
+        /// <summary>
+        /// The unique ID of the JWT the user is using to authenticate.
+        /// This is used to control user uniqueness.
+        /// </summary>
+        public readonly string TokenId;
+
+        /// <summary>
+        /// The connection IDs of the user.
+        /// </summary>
+        /// <remarks>
+        /// In SignalR, connection IDs are unique per user, <em>and</em> per hub instance.
+        /// Therefore, to keep track of all of them, a dictionary is necessary.
+        /// </remarks>
+        public readonly Dictionary<Type, string> ConnectionIds = new Dictionary<Type, string>();
+
+        public ConnectionState(HubLifetimeContext context)
+        {
+            TokenId = context.Context.GetTokenId();
+
+            RegisterConnectionId(context);
+        }
+
+        public void RegisterConnectionId(HubLifetimeContext context)
+            => ConnectionIds.Add(context.Hub.GetType(), context.Context.ConnectionId);
+    }
+}

osu.Server.Spectator/Extensions/HubCallerContextExtensions.cs

@@ -0,0 +1,34 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System;
+using Microsoft.AspNetCore.SignalR;
+
+namespace osu.Server.Spectator.Extensions
+{
+    public static class HubCallerContextExtensions
+    {
+        /// <summary>
+        /// Returns the osu! user id for the supplied <paramref name="context"/>.
+        /// </summary>
+        public static int GetUserId(this HubCallerContext context)
+        {
+            if (context.UserIdentifier == null)
+                throw new InvalidOperationException($"Attempted to get user id with null {nameof(context.UserIdentifier)}");
+
+            return int.Parse(context.UserIdentifier);
+        }
+
+        /// <summary>
+        /// Returns the ID of the authorisation token (more accurately, the <c>jti</c> claim)
+        /// for the supplied <paramref name="context"/>.
+        /// This is used for the purpose of identifying individual client instances
+        /// and preventing multiple concurrent sessions from being active.
+        /// </summary>
+        public static string GetTokenId(this HubCallerContext context)
+        {
+            return context.User?.FindFirst(claim => claim.Type == "jti")?.Value
+                   ?? throw new InvalidOperationException("Could not retrieve JWT ID claim from token");
+        }
+    }
+}

osu.Server.Spectator/Extensions/ServiceCollectionExtensions.cs

@@ -20,6 +20,7 @@ public static IServiceCollection AddHubEntities(this IServiceCollection serviceC
             return serviceCollection.AddSingleton<EntityStore<SpectatorClientState>>()
                                     .AddSingleton<EntityStore<MultiplayerClientState>>()
                                     .AddSingleton<EntityStore<ServerMultiplayerRoom>>()
+                                    .AddSingleton<EntityStore<ConnectionState>>()
                                     .AddSingleton<GracefulShutdownManager>()
                                     .AddSingleton<MetadataBroadcaster>()
                                     .AddSingleton<IScoreStorage, S3ScoreStorage>()

osu.Server.Spectator/Hubs/IStatefulUserHub.cs

@@ -0,0 +1,13 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+namespace osu.Server.Spectator.Hubs
+{
+    /// <summary>
+    /// Marker interface for <see cref="StatefulUserHub{TClient,TUserState}"/>.
+    /// Allows bypassing generic constraints.
+    /// </summary>
+    public interface IStatefulUserHub
+    {
+    }
+}

osu.Server.Spectator/Hubs/LoggingHub.cs

@@ -28,20 +28,6 @@ public LoggingHub()
             logger = Logger.GetLogger(Name);
         }
 
-        /// <summary>
-        /// The osu! user id for the currently processing context.
-        /// </summary>
-        protected int CurrentContextUserId
-        {
-            get
-            {
-                if (Context.UserIdentifier == null)
-                    throw new InvalidOperationException($"Attempted to get user id with null {nameof(Context.UserIdentifier)}");
-
-                return int.Parse(Context.UserIdentifier);
-            }
-        }
-
         public override async Task OnConnectedAsync()
         {
             Log("Connected");