fix(connections): address review — stable heartbeat fake, Ensure race gate, no token leak, swap value validation

Author: HandyS11

src/RustPlusBot.Features.Connections/Listening/RustPlusSocketSource.cs

@@ -13,13 +13,29 @@ public IRustServerConnection Create(string ip, int port, ulong steamId, string p
     {
         if (!int.TryParse(playerToken, CultureInfo.InvariantCulture, out var token))
         {
-            throw new ArgumentException($"playerToken is not a valid numeric token: '{playerToken}'.",
-                nameof(playerToken));
+            LogInvalidToken(logger);
+            return new RejectedConnection();
         }
 
         return new RustPlusServerConnection(ip, port, steamId, token, logger);
     }
 
+    [LoggerMessage(Level = LogLevel.Warning,
+        Message = "Player token is not a valid numeric token; treating the credential as rejected.")]
+    private static partial void LogInvalidToken(ILogger logger);
+
+    /// <summary>Returned when the player token is unusable; reports the credential as rejected and does nothing else.</summary>
+    private sealed class RejectedConnection : IRustServerConnection
+    {
+        public Task<SocketConnectOutcome> ConnectAsync(TimeSpan timeout, CancellationToken cancellationToken) =>
+            Task.FromResult(SocketConnectOutcome.AuthRejected);
+
+        public Task<HeartbeatResult> GetInfoAsync(TimeSpan timeout, CancellationToken cancellationToken) =>
+            Task.FromResult(HeartbeatResult.AuthRejected);
+
+        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+    }
+
     /// <summary>
     /// One live Rust+ connection backed by <see cref="RustPlus"/>.
     /// <para>API notes (2.0.0-beta.1): constructor takes <see cref="RustPlusConnection"/> instead of

src/RustPlusBot.Features.Connections/Modules/ConnectionComponentModule.cs

@@ -28,7 +28,7 @@ public async Task SwapAsync(string serverIdRaw, string[] selectedValues)
         }
 
         if (!Guid.TryParse(serverIdRaw, out var serverId)
-            || selectedValues.Length == 0
+            || selectedValues.Length != 1
             || !Guid.TryParse(selectedValues[0], out var credentialId))
         {
             await RespondAsync("That selection wasn't valid.", ephemeral: true).ConfigureAwait(false);

src/RustPlusBot.Features.Connections/Supervisor/ConnectionSupervisor.cs

@@ -32,6 +32,7 @@ internal sealed partial class ConnectionSupervisor(
     ILogger<ConnectionSupervisor> logger) : IConnectionSupervisor, IAsyncDisposable
 {
     private readonly ConcurrentDictionary<(ulong Guild, Guid Server), Handle> _connections = new();
+    private readonly SemaphoreSlim _gate = new(1, 1);
     private readonly ConnectionOptions _options = options.Value;
     private readonly CancellationTokenSource _shutdown = new();
 
@@ -40,6 +41,7 @@ public async ValueTask DisposeAsync()
     {
         await StopAllAsync().ConfigureAwait(false);
         _shutdown.Dispose();
+        _gate.Dispose();
     }
 
     /// <inheritdoc />
@@ -63,26 +65,53 @@ public async Task StartAllAsync(CancellationToken cancellationToken = default)
     public async Task EnsureConnectionAsync(ulong guildId, Guid serverId, CancellationToken cancellationToken = default)
     {
         var key = (guildId, serverId);
-        await StopConnectionAsync(key).ConfigureAwait(false);
-        if (_shutdown.IsCancellationRequested)
+        await _gate.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
         {
-            return;
-        }
+            await StopConnectionAsync(key).ConfigureAwait(false);
+            if (_shutdown.IsCancellationRequested)
+            {
+                return;
+            }
 
-        var cts = CancellationTokenSource.CreateLinkedTokenSource(_shutdown.Token);
-        _connections[key] = new Handle(cts, Task.Run(() => RunAsync(key, cts.Token), CancellationToken.None));
+            var cts = CancellationTokenSource.CreateLinkedTokenSource(_shutdown.Token);
+            _connections[key] = new Handle(cts, Task.Run(() => RunAsync(key, cts.Token), CancellationToken.None));
+        }
+        finally
+        {
+            _gate.Release();
+        }
     }
 
     /// <inheritdoc />
-    public Task StopAsync(ulong guildId, Guid serverId) => StopConnectionAsync((guildId, serverId));
+    public async Task StopAsync(ulong guildId, Guid serverId)
+    {
+        await _gate.WaitAsync().ConfigureAwait(false);
+        try
+        {
+            await StopConnectionAsync((guildId, serverId)).ConfigureAwait(false);
+        }
+        finally
+        {
+            _gate.Release();
+        }
+    }
 
     /// <inheritdoc />
     public async Task StopAllAsync()
     {
         await _shutdown.CancelAsync().ConfigureAwait(false);
-        foreach (var key in _connections.Keys.ToList())
+        await _gate.WaitAsync().ConfigureAwait(false);
+        try
         {
-            await StopConnectionAsync(key).ConfigureAwait(false);
+            foreach (var key in _connections.Keys.ToList())
+            {
+                await StopConnectionAsync(key).ConfigureAwait(false);
+            }
+        }
+        finally
+        {
+            _gate.Release();
         }
     }
 

tests/RustPlusBot.Features.Connections.Tests/ConnectionSupervisorTests.cs

@@ -37,11 +37,21 @@ private static Harness CreateHarness(FakeRustSocketSource source)
         services.AddSingleton(dm);
         services.AddSingleton<IEventBus, InMemoryEventBus>();
 
-        var (seed, connection) = TestDb.Create();
-        seed.Dispose();
-        services.AddSingleton(connection);
-        services.AddScoped(sp => new BotDbContext(
-            new DbContextOptionsBuilder<BotDbContext>().UseSqlite(sp.GetRequiredService<SqliteConnection>()).Options));
+        // Each scope opens its OWN connection to a shared-cache in-memory database, so the background
+        // supervisor loop and the test's polling never run concurrent commands on a single SqliteConnection
+        // (which throws "active statements" misuse errors). One kept-open connection keeps the in-memory DB alive.
+        var connectionString = $"DataSource=connsup-{Guid.NewGuid():N};Mode=Memory;Cache=Shared";
+        var keepAlive = new SqliteConnection(connectionString);
+        keepAlive.Open();
+        using (var seed = new BotDbContext(
+                   new DbContextOptionsBuilder<BotDbContext>().UseSqlite(connectionString).Options))
+        {
+            seed.Database.Migrate();
+        }
+
+        services.AddSingleton(keepAlive);
+        services.AddScoped(_ => new BotDbContext(
+            new DbContextOptionsBuilder<BotDbContext>().UseSqlite(connectionString).Options));
         services.AddScoped<IConnectionStore, ConnectionStore>();
         services
             .AddScoped<RustPlusBot.Persistence.Servers.IServerService, RustPlusBot.Persistence.Servers.ServerService>();
@@ -103,7 +113,7 @@ private static Harness CreateHarness(FakeRustSocketSource source)
         Guid serverId,
         Func<ConnectionState, bool> predicate)
     {
-        var deadline = DateTimeOffset.UtcNow.AddSeconds(20);
+        var deadline = DateTimeOffset.UtcNow.AddSeconds(30);
         while (DateTimeOffset.UtcNow < deadline)
         {
             using var scope = provider.CreateScope();
@@ -224,7 +234,7 @@ public async Task StartAll_StartsAConnectionPerConnectableServer()
 
         await h.Supervisor.StartAllAsync();
 
-        var deadline = DateTimeOffset.UtcNow.AddSeconds(20);
+        var deadline = DateTimeOffset.UtcNow.AddSeconds(30);
         while (source.CreateCount < 1 && DateTimeOffset.UtcNow < deadline)
         {
             await Task.Delay(15);

tests/RustPlusBot.Features.Connections.Tests/Fakes/FakeRustSocketSource.cs

@@ -12,14 +12,17 @@ namespace RustPlusBot.Features.Connections.Tests.Fakes;
 /// source creates, consumed in enqueue order. This is intentional: the supervisor tests drive one
 /// server at a time and enqueue a deterministic sequence that must flow across reconnections of that
 /// server. Tests that need concurrent multi-server interleaving would require per-connection queues.
+/// Once the heartbeat queue is exhausted, the fake HOLDS the last dequeued heartbeat (mimicking a
+/// healthy server that keeps reporting the same state), rather than falling back to Ok(0).
 /// </remarks>
 internal sealed class FakeRustSocketSource : IRustSocketSource
 {
     private readonly ConcurrentQueue<SocketConnectOutcome> _connectOutcomes = new();
     private readonly ConcurrentQueue<HeartbeatResult> _heartbeats = new();
-
     private int _createCount;
 
+    private HeartbeatResult _lastHeartbeat = HeartbeatResult.Ok(0);
+
     /// <summary>Number of times <see cref="Create"/> has been called. Safe to read from any thread.</summary>
     public int CreateCount => Volatile.Read(ref _createCount);
 
@@ -35,21 +38,31 @@ public IRustServerConnection Create(string ip, int port, ulong steamId, string p
         LastIp = ip;
         LastSteamId = steamId;
         var outcome = _connectOutcomes.TryDequeue(out var next) ? next : SocketConnectOutcome.Connected;
-        return new FakeConnection(outcome, _heartbeats);
+        return new FakeConnection(outcome, this);
     }
 
     public void EnqueueConnect(SocketConnectOutcome outcome) => _connectOutcomes.Enqueue(outcome);
 
     public void EnqueueHeartbeat(HeartbeatResult result) => _heartbeats.Enqueue(result);
 
-    private sealed class FakeConnection(SocketConnectOutcome outcome, ConcurrentQueue<HeartbeatResult> heartbeats)
+    internal HeartbeatResult NextHeartbeat()
+    {
+        if (_heartbeats.TryDequeue(out var next))
+        {
+            _lastHeartbeat = next;
+        }
+
+        return _lastHeartbeat;
+    }
+
+    private sealed class FakeConnection(SocketConnectOutcome outcome, FakeRustSocketSource source)
         : IRustServerConnection
     {
         public Task<SocketConnectOutcome> ConnectAsync(TimeSpan timeout, CancellationToken cancellationToken) =>
             Task.FromResult(outcome);
 
         public Task<HeartbeatResult> GetInfoAsync(TimeSpan timeout, CancellationToken cancellationToken) =>
-            Task.FromResult(heartbeats.TryDequeue(out var next) ? next : HeartbeatResult.Ok(0));
+            Task.FromResult(source.NextHeartbeat());
 
         public ValueTask DisposeAsync() => ValueTask.CompletedTask;
     }

tests/RustPlusBot.Features.Connections.Tests/RustPlusSocketSourceTests.cs

@@ -15,4 +15,15 @@ public async Task Create_ReturnsADisposableConnection()
 
         Assert.NotNull(connection);
     }
+
+    [Fact]
+    public async Task Create_NonNumericToken_YieldsAuthRejectedConnection()
+    {
+        var source = new RustPlusSocketSource(NullLogger<RustPlusSocketSource>.Instance);
+
+        var connection = source.Create("127.0.0.1", 28015, 100UL, "not-a-number");
+        await using var _ = connection;
+
+        Assert.Equal(SocketConnectOutcome.AuthRejected, await connection.ConnectAsync(TimeSpan.Zero, default));
+    }
 }