Merge pull request #328 from VelvetToroyashi/fix/gateway-op7-on-connect

Handle OP7 on connect
Remora · Apr 2, 2024 · 5219a50 · 5219a50
2 parents c5f0992 + c12b88c
commit 5219a50
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 0 deletions.
diff --git a/Backend/Remora.Discord.Gateway/DiscordGatewayClient.cs b/Backend/Remora.Discord.Gateway/DiscordGatewayClient.cs
@@ -566,6 +566,18 @@ private async Task<Result> RunConnectionIterationAsync(CancellationToken stopReq
 
                 if (receiveHello.Entity is not IPayload<IHello> hello)
                 {
+                    if (receiveHello.Entity is IPayload<IReconnect>)
+                    {
+                        // Discord may spit out a reconnect if the node is we're connecting while the gateway node is
+                        // shutting down, but before the node is labeled as unavailable.
+                        return new GatewayError
+                        (
+                            "The gateway requested a reconnect.",
+                            false,
+                            false
+                        );
+                    }
+
                     // Not receiving a hello is a non-recoverable error
                     return new GatewayError
                     (

diff --git a/Tests/Remora.Discord.Gateway.Tests/Tests/DiscordGatewayClientTests.cs b/Tests/Remora.Discord.Gateway.Tests/Tests/DiscordGatewayClientTests.cs
@@ -711,6 +711,76 @@ public async Task CanReconnectAfterExceptionAsync()
         ResultAssert.Successful(runResult);
     }
 
+    /// <summary>
+    /// Tests whether the client can reconnect after being sent a Reconnect instead of Hello.
+    /// </summary>
+    /// <returns>A <see cref="Task"/> representing the asynchronous unit test.</returns>
+    [Fact]
+    public async Task CanReconnectAfterReconnectInsteadOfHelloAsync()
+    {
+        var tokenSource = new CancellationTokenSource();
+        var transportMock = new MockedTransportServiceBuilder(_testOutput)
+            .WithTimeout(TimeSpan.FromSeconds(30))
+            .Sequence
+            (
+                s => s
+                    .ExpectConnection
+                    (
+                        new Uri($"wss://gateway.discord.gg/?v={(int)DiscordAPIVersion.V10}&encoding=json")
+                    )
+                    .Send(new Reconnect())
+                    .ExpectDisconnect()
+                    .ExpectConnection
+                    (
+                        new Uri($"wss://gateway.discord.gg/?v={(int)DiscordAPIVersion.V10}&encoding=json")
+                    )
+                    .Send(new Hello(TimeSpan.FromMilliseconds(200)))
+                    .Expect<Identify>
+                    (
+                        i =>
+                        {
+                            Assert.Equal(Constants.MockToken, i?.Token);
+                            return true;
+                        }
+                    )
+                    .Send
+                    (
+                        new Ready
+                        (
+                            8,
+                            Constants.BotUser,
+                            new List<IUnavailableGuild>(),
+                            Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
+                            default,
+                            new PartialApplication()
+                        )
+                    )
+            )
+            .Continuously
+            (
+                c => c
+                    .Expect<IHeartbeat>()
+                    .Send<HeartbeatAcknowledge>()
+            )
+            .Finish(tokenSource)
+            .Build();
+
+        var transportMockDescriptor = ServiceDescriptor.Singleton(typeof(IPayloadTransportService), transportMock);
+
+        var services = new ServiceCollection()
+            .AddDiscordGateway(_ => Constants.MockToken)
+            .Replace(transportMockDescriptor)
+            .Replace(CreateMockedGatewayAPI())
+            .AddSingleton<IResponderTypeRepository, ResponderService>()
+            .BuildServiceProvider(true);
+
+        var client = services.GetRequiredService<DiscordGatewayClient>();
+        var runResult = await client.RunAsync(tokenSource.Token);
+
+        ResultAssert.Successful(runResult);
+    }
+
     private static ServiceDescriptor CreateMockedGatewayAPI()
     {
         var gatewayAPIMock = new Mock<IDiscordRestGatewayAPI>();