Experimental: Add lock on game objects

src/Impostor.Api/Games/IGame.cs

@@ -1,10 +1,12 @@
 using System.Collections.Generic;
 using System.Net;
 using System.Threading.Tasks;
+using Impostor.Api.Events;
 using Impostor.Api.Innersloth;
 using Impostor.Api.Innersloth.GameOptions;
 using Impostor.Api.Net;
 using Impostor.Api.Net.Inner;
+using Impostor.Api.Utils;
 
 namespace Impostor.Api.Games
 {
@@ -52,6 +54,11 @@ public interface IGame
         T? FindObjectByNetId<T>(uint netId)
             where T : IInnerNetObject;
 
+        /// <summary>Lock an <see cref="IGame"/> to prevent concurrent modification.</summary>
+        /// <returns>A <see cref="AsyncLock.Releaser" /> that should be disposed after modifying the game.</returns>
+        /// <remarks>This method should be called when modifying a game outside of an <see cref="IGameEvent" /> handler.</remarks>
+        ValueTask<AsyncLock.Releaser> LockAsync();
+
         /// <summary>
         ///     Adds an <see cref="IPAddress" /> to the ban list of this game.
         ///     Prevents all future joins from this <see cref="IPAddress" />.

src/Impostor.Api/Utils/AsyncLock.cs

@@ -0,0 +1,71 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Impostor.Api.Utils;
+
+/// <summary>
+/// Represents an asynchronous lock mechanism.
+/// </summary>
+public sealed class AsyncLock
+{
+    private readonly SemaphoreSlim _semaphore = new(1, 1);
+
+    /// <summary>
+    /// Asynchronously acquires a lock.
+    /// </summary>
+    /// <returns>
+    /// A <see cref="ValueTask{Releaser}"/> that represents the asynchronous lock operation.
+    /// The result of the task is a <see cref="Releaser"/> that should be disposed to release the lock.
+    /// </returns>
+    public ValueTask<Releaser> LockAsync()
+    {
+        return LockAsync(CancellationToken.None);
+    }
+
+    /// <inheritdoc cref="LockAsync()"/>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to observe while waiting for the lock.</param>
+    /// <exception cref="T:System.OperationCanceledException"><paramref name="cancellationToken" /> was canceled.</exception>
+    public ValueTask<Releaser> LockAsync(CancellationToken cancellationToken)
+    {
+        var wait = _semaphore.WaitAsync(cancellationToken);
+
+        if (wait.IsCompleted)
+        {
+            return ValueTask.FromResult(new Releaser(this));
+        }
+
+        return new ValueTask<Releaser>(wait.ContinueWith(
+            static (_, state) => new Releaser((AsyncLock)state!),
+            this,
+            cancellationToken,
+            TaskContinuationOptions.ExecuteSynchronously,
+            TaskScheduler.Default
+        ));
+    }
+
+    /// <summary>
+    /// Represents a disposable object that, when disposed, releases the associated lock.
+    /// </summary>
+    public struct Releaser : IDisposable
+    {
+        private AsyncLock? _lock;
+
+        internal Releaser(AsyncLock @lock)
+        {
+            _lock = @lock;
+        }
+
+        /// <summary>
+        /// Releases the lock held by this <see cref="Releaser"/> instance.
+        /// </summary>
+        public void Dispose()
+        {
+            if (_lock != null)
+            {
+                _lock._semaphore.Release();
+                _lock = null;
+            }
+        }
+    }
+}

src/Impostor.Server/Net/State/Game.Data.cs

@@ -63,6 +63,8 @@ internal partial class Game
 
         public async ValueTask<bool> HandleGameDataAsync(IMessageReader parent, ClientPlayer sender, bool toPlayer)
         {
+            using var guard = await LockAsync();
+
             // Find target player.
             ClientPlayer? target = null;
 

src/Impostor.Server/Net/State/Game.Incoming.cs

@@ -1,5 +1,3 @@
-using System;
-using System.Threading;
 using System.Threading.Tasks;
 using Impostor.Api.Games;
 using Impostor.Api.Innersloth;
@@ -13,10 +11,9 @@ namespace Impostor.Server.Net.State
 {
     internal partial class Game
     {
-        private readonly SemaphoreSlim _clientAddLock = new SemaphoreSlim(1, 1);
-
         public async ValueTask HandleStartGame(IMessageReader message)
         {
+            using var guard = await LockAsync();
             GameState = GameStates.Starting;
 
             using var packet = MessageWriter.Get(MessageType.Reliable);
@@ -28,6 +25,7 @@ public async ValueTask HandleStartGame(IMessageReader message)
 
         public async ValueTask HandleEndGame(IMessageReader message, GameOverReason gameOverReason)
         {
+            using var guard = await LockAsync();
             GameState = GameStates.Ended;
 
             // Broadcast end of the game.
@@ -48,6 +46,7 @@ public async ValueTask HandleEndGame(IMessageReader message, GameOverReason game
 
         public async ValueTask HandleAlterGame(IMessageReader message, IClientPlayer sender, bool isPublic)
         {
+            using var guard = await LockAsync();
             IsPublic = isPublic;
 
             using var packet = MessageWriter.Get(MessageType.Reliable);
@@ -74,6 +73,7 @@ public async ValueTask HandleRemovePlayer(int playerId, DisconnectReason reason)
 
         public async ValueTask HandleKickPlayer(int playerId, bool isBan)
         {
+            using var guard = await LockAsync();
             _logger.LogInformation("{0} - Player {1} has left.", Code, playerId);
 
             using var message = MessageWriter.Get(MessageType.Reliable);
@@ -96,26 +96,8 @@ public async ValueTask HandleKickPlayer(int playerId, bool isBan)
 
         public async ValueTask<GameJoinResult> AddClientAsync(ClientBase client)
         {
-            var hasLock = false;
-
-            try
-            {
-                hasLock = await _clientAddLock.WaitAsync(TimeSpan.FromMinutes(1));
-
-                if (hasLock)
-                {
-                    return await AddClientSafeAsync(client);
-                }
-            }
-            finally
-            {
-                if (hasLock)
-                {
-                    _clientAddLock.Release();
-                }
-            }
-
-            return GameJoinResult.FromError(GameJoinError.InvalidClient);
+            using var guard = await LockAsync();
+            return await AddClientSafeAsync(client);
         }
 
         private async ValueTask HandleJoinGameNew(ClientPlayer sender, bool isNew)

src/Impostor.Server/Net/State/Game.cs

@@ -13,6 +13,7 @@
 using Impostor.Api.Net;
 using Impostor.Api.Net.Manager;
 using Impostor.Api.Net.Messages.S2C;
+using Impostor.Api.Utils;
 using Impostor.Server.Events;
 using Impostor.Server.Net.Manager;
 using Microsoft.Extensions.Logging;
@@ -32,6 +33,7 @@ internal partial class Game
         private readonly ICompatibilityManager _compatibilityManager;
         private readonly CompatibilityConfig _compatibilityConfig;
         private readonly TimeoutConfig _timeoutConfig;
+        private readonly AsyncLock _lock = new();
 
         public Game(
             ILogger<Game> logger,
@@ -113,6 +115,11 @@ public bool TryGetPlayer(int id, [MaybeNullWhen(false)] out ClientPlayer player)
             return _players.TryGetValue(clientId, out var clientPlayer) ? clientPlayer : null;
         }
 
+        public ValueTask<AsyncLock.Releaser> LockAsync()
+        {
+            return _lock.LockAsync();
+        }
+
         internal async ValueTask StartedAsync()
         {
             if (GameState == GameStates.Starting)