From 53aca791ddd33449f6b49037698b7e97d363545e Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Tue, 19 Sep 2023 00:02:06 +1000
Subject: [PATCH] Fix reentering of the configuration state

---
 .../Protocol1_20_2To1_20.java                 | 41 +++++++++++++++++--
 .../rewriter/EntityPacketRewriter1_20_2.java  | 25 +++++------
 .../storage/ConfigurationState.java           | 25 ++++++++++-
 settings.gradle.kts                           |  2 +-
 4 files changed, 76 insertions(+), 17 deletions(-)

diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/Protocol1_20_2To1_20.java b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/Protocol1_20_2To1_20.java
index b1731042e..df94577fe 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/Protocol1_20_2To1_20.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/Protocol1_20_2To1_20.java
@@ -17,6 +17,7 @@
  */
 package com.viaversion.viaversion.protocols.protocol1_20_2to1_20;
 
+import com.github.steveice10.opennbt.tag.builtin.CompoundTag;
 import com.viaversion.viaversion.api.Via;
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.data.MappingData;
@@ -117,7 +118,19 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
         registerServerbound(State.CONFIGURATION, ServerboundConfigurationPackets1_20_2.RESOURCE_PACK.getId(), -1, queueServerboundPacket(ServerboundPackets1_20_2.RESOURCE_PACK_STATUS));
 
         cancelClientbound(ClientboundPackets1_19_4.UPDATE_ENABLED_FEATURES); // TODO Sad emoji
-        cancelServerbound(ServerboundPackets1_20_2.CONFIGURATION_ACKNOWLEDGED);
+        registerServerbound(ServerboundPackets1_20_2.CONFIGURATION_ACKNOWLEDGED, null, wrapper -> {
+            wrapper.cancel();
+
+            final ConfigurationState configurationState = wrapper.user().get(ConfigurationState.class);
+            if (configurationState.getReenterInfo() == null) {
+                return;
+            }
+
+            // Reenter the configuration state
+            configurationState.setBridgePhase(BridgePhase.CONFIGURATION);
+            sendConfigurationPackets(wrapper.user(), configurationState.getReenterInfo().dimensionRegistry());
+            configurationState.setReenterInfo(null);
+        });
         cancelServerbound(ServerboundPackets1_20_2.CHUNK_BATCH_RECEIVED);
 
         registerServerbound(ServerboundPackets1_20_2.PING_REQUEST, null, wrapper -> {
@@ -141,19 +154,25 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
     @Override
     public void transform(final Direction direction, final State state, final PacketWrapper packetWrapper) throws Exception {
         final ConfigurationState configurationBridge = packetWrapper.user().get(ConfigurationState.class);
+        if (configurationBridge == null) {
+            // Bad state during an unexpected disconnect
+            return;
+        }
+
         final BridgePhase phase = configurationBridge.bridgePhase();
         if (phase == BridgePhase.NONE) {
             super.transform(direction, state, packetWrapper);
             return;
         }
 
+        final boolean reentering = configurationBridge.getReenterInfo() != null;
         if (direction == Direction.SERVERBOUND) {
             // Client and server will be on different protocol states, pick the right client protocol state
-            super.transform(direction, phase == BridgePhase.PROFILE_SENT ? State.LOGIN : State.CONFIGURATION, packetWrapper);
+            super.transform(direction, phase == BridgePhase.PROFILE_SENT ? State.LOGIN : reentering ? State.PLAY : State.CONFIGURATION, packetWrapper);
             return;
         }
 
-        if (phase == BridgePhase.PROFILE_SENT) {
+        if (phase == BridgePhase.PROFILE_SENT || reentering) {
             // Queue packets sent by the server while we wait for the client to transition to the configuration state
             configurationBridge.addPacketToQueue(packetWrapper, true);
             throw CancelException.generate();
@@ -207,6 +226,22 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
         super.transform(direction, State.CONFIGURATION, packetWrapper);
     }
 
+    public static void sendConfigurationPackets(final UserConnection connection, final CompoundTag dimensionRegistry) throws Exception {
+        final PacketWrapper registryDataPacket = PacketWrapper.create(ClientboundConfigurationPackets1_20_2.REGISTRY_DATA, connection);
+        registryDataPacket.write(Type.NAMELESS_NBT, dimensionRegistry);
+        registryDataPacket.send(Protocol1_20_2To1_20.class);
+
+        // Enabling features is only possible during the configuration phase
+        // TODO Sad emoji
+        final PacketWrapper enableFeaturesPacket = PacketWrapper.create(ClientboundConfigurationPackets1_20_2.UPDATE_ENABLED_FEATURES, connection);
+        enableFeaturesPacket.write(Type.VAR_INT, 1);
+        enableFeaturesPacket.write(Type.STRING, "minecraft:vanilla");
+        enableFeaturesPacket.send(Protocol1_20_2To1_20.class);
+
+        final PacketWrapper finishConfigurationPacket = PacketWrapper.create(ClientboundConfigurationPackets1_20_2.FINISH_CONFIGURATION, connection);
+        finishConfigurationPacket.send(Protocol1_20_2To1_20.class);
+    }
+
     @Override
     public MappingData getMappingData() {
         return MAPPINGS;
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/rewriter/EntityPacketRewriter1_20_2.java b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/rewriter/EntityPacketRewriter1_20_2.java
index 47a0cc1a2..9ff1a7cd5 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/rewriter/EntityPacketRewriter1_20_2.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/rewriter/EntityPacketRewriter1_20_2.java
@@ -100,22 +100,23 @@ public final class EntityPacketRewriter1_20_2 extends EntityRewriter<Clientbound
                     // Debug, flat, last death pos, and portal cooldown at the end unchanged
 
                     // Send configuration packets first before going into the play protocol state
-                    final PacketWrapper registryDataPacket = wrapper.create(ClientboundConfigurationPackets1_20_2.REGISTRY_DATA);
-                    registryDataPacket.write(Type.NAMELESS_NBT, dimensionRegistry);
-                    registryDataPacket.send(Protocol1_20_2To1_20.class);
+                    ConfigurationState configurationBridge = wrapper.user().get(ConfigurationState.class);
+                    if (configurationBridge.bridgePhase() == ConfigurationState.BridgePhase.NONE) {
+                        // Reenter the configuration state
+                        final PacketWrapper configurationPacket = wrapper.create(ClientboundPackets1_20_2.START_CONFIGURATION);
+                        configurationPacket.send(Protocol1_20_2To1_20.class);
 
-                    // Enabling features is only possible during the configuraton phase
-                    // TODO Sad emoji
-                    final PacketWrapper enableFeaturesPacket = wrapper.create(ClientboundConfigurationPackets1_20_2.UPDATE_ENABLED_FEATURES);
-                    enableFeaturesPacket.write(Type.VAR_INT, 1);
-                    enableFeaturesPacket.write(Type.STRING, "minecraft:vanilla");
-                    enableFeaturesPacket.send(Protocol1_20_2To1_20.class);
+                        // TODO The client clears the resource pack when reentering (?)
+                        configurationBridge.setBridgePhase(ConfigurationState.BridgePhase.REENTERING_CONFIGURATION);
+                        configurationBridge.setJoinGamePacket(wrapper);
+                        configurationBridge.setReenterInfo(new ConfigurationState.ReenterInfo(dimensionRegistry));
+                        wrapper.cancel();
+                        return;
+                    }
 
-                    final PacketWrapper finishConfigurationPacket = wrapper.create(ClientboundConfigurationPackets1_20_2.FINISH_CONFIGURATION);
-                    finishConfigurationPacket.send(Protocol1_20_2To1_20.class);
+                    Protocol1_20_2To1_20.sendConfigurationPackets(wrapper.user(), dimensionRegistry);
 
                     // Manually send it at the end and hope nothing breaks
-                    final ConfigurationState configurationBridge = wrapper.user().get(ConfigurationState.class);
                     configurationBridge.setJoinGamePacket(wrapper);
                     wrapper.cancel();
                 });
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
index 79e669499..81b7f1295 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
@@ -17,6 +17,7 @@
  */
 package com.viaversion.viaversion.protocols.protocol1_20_2to1_20.storage;
 
+import com.github.steveice10.opennbt.tag.builtin.CompoundTag;
 import com.viaversion.viaversion.api.connection.StorableObject;
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.protocol.packet.PacketType;
@@ -34,6 +35,7 @@ public class ConfigurationState implements StorableObject {
     private BridgePhase bridgePhase = BridgePhase.NONE;
     private QueuedPacket joinGamePacket;
     private boolean queuedJoinGame;
+    private ReenterInfo reenterInfo;
 
     public BridgePhase bridgePhase() {
         return bridgePhase;
@@ -43,6 +45,14 @@ public class ConfigurationState implements StorableObject {
         this.bridgePhase = bridgePhase;
     }
 
+    public @Nullable ReenterInfo getReenterInfo() {
+        return reenterInfo;
+    }
+
+    public void setReenterInfo(@Nullable final ReenterInfo reenterInfo) {
+        this.reenterInfo = reenterInfo;
+    }
+
     public void addPacketToQueue(final PacketWrapper wrapper, final boolean clientbound) throws Exception {
         packetQueue.add(toQueuedPacket(wrapper, clientbound, false));
     }
@@ -113,6 +123,7 @@ public class ConfigurationState implements StorableObject {
         packetQueue.clear();
         bridgePhase = BridgePhase.NONE;
         queuedJoinGame = false;
+        reenterInfo = null;
     }
 
     public boolean queuedOrSentJoinGame() {
@@ -120,7 +131,7 @@ public class ConfigurationState implements StorableObject {
     }
 
     public enum BridgePhase {
-        NONE, PROFILE_SENT, CONFIGURATION
+        NONE, PROFILE_SENT, CONFIGURATION, REENTERING_CONFIGURATION
     }
 
     public static final class QueuedPacket {
@@ -170,4 +181,16 @@ public class ConfigurationState implements StorableObject {
                     '}';
         }
     }
+
+    public static final class ReenterInfo {
+        private final CompoundTag dimensionRegistry;
+
+        public ReenterInfo(final CompoundTag dimensionRegistry) {
+            this.dimensionRegistry = dimensionRegistry;
+        }
+
+        public CompoundTag dimensionRegistry() {
+            return dimensionRegistry;
+        }
+    }
 }
diff --git a/settings.gradle.kts b/settings.gradle.kts
index e29d41b00..dc9697336 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -17,7 +17,7 @@ dependencyResolutionManagement {
 pluginManagement {
     // default plugin versions
     plugins {
-        id("net.kyori.blossom") version "2.0.1"
+        id("net.kyori.blossom") version "2.1.0"
         id("org.jetbrains.gradle.plugin.idea-ext") version "1.1.7"
         id("com.github.johnrengelman.shadow") version "8.1.1"
     }