Fix race conditions with login plugin messages

SpongePowered · Nov 22, 2023 · dfe4de5 · dfe4de5
1 parent df736b8
commit dfe4de5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/main/java/org/spongepowered/common/network/channel/raw/SpongeRawLoginDataChannel.java b/src/main/java/org/spongepowered/common/network/channel/raw/SpongeRawLoginDataChannel.java
@@ -167,10 +167,10 @@ public CompletableFuture<ChannelBuf> sendTo(final EngineConnection connection, f
 
         final Packet<?> mcPacket = PacketUtil.createLoginPayloadRequest(this.parent.key(), buf, transactionId);
 
+        transactionStore.put(transactionId, this.parent, resultConsumer);
+
         PacketSender.sendTo(connection, mcPacket, throwable -> {
-            if (throwable == null) {
-                transactionStore.put(transactionId, this.parent, resultConsumer);
-            } else {
+            if (throwable != null && transactionStore.remove(transactionId) != null) {
                 // The packet already failed before it could reach the client
                 future.completeExceptionally(throwable);
             }

diff --git a/...powered/vanilla/mixin/core/server/network/ServerLoginPacketListenerImplMixin_Vanilla.java b/...powered/vanilla/mixin/core/server/network/ServerLoginPacketListenerImplMixin_Vanilla.java
@@ -61,7 +61,8 @@ public abstract class ServerLoginPacketListenerImplMixin_Vanilla implements Serv
     private static final int HANDSHAKE_NOT_STARTED = 0;
     private static final int HANDSHAKE_CLIENT_TYPE = 1;
     private static final int HANDSHAKE_SYNC_CHANNEL_REGISTRATIONS = 2;
-    private static final int HANDSHAKE_SYNC_PLUGIN_DATA = 3;
+    private static final int HANDSHAKE_CHANNEL_REGISTRATION = 3;
+    private static final int HANDSHAKE_SYNC_PLUGIN_DATA = 4;
 
     private int impl$handshakeState = ServerLoginPacketListenerImplMixin_Vanilla.HANDSHAKE_NOT_STARTED;
 
@@ -85,6 +86,8 @@ private void onResponsePayload(final ServerboundCustomQueryPacket packet, final
                 });
 
             } else if (this.impl$handshakeState == ServerLoginPacketListenerImplMixin_Vanilla.HANDSHAKE_SYNC_CHANNEL_REGISTRATIONS) {
+                this.impl$handshakeState = ServerLoginPacketListenerImplMixin_Vanilla.HANDSHAKE_CHANNEL_REGISTRATION;
+
                 ((SpongeChannelManager) Sponge.channelManager()).sendLoginChannelRegistry(connection).thenAccept(result -> {
                     final Cause cause = Cause.of(EventContext.empty(), this);
                     final ServerSideConnectionEvent.Handshake event =