From 1bc3d407d19c8c2e6cfe5f554de25831b5b40c00 Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Thu, 3 Aug 2023 23:03:42 +1000
Subject: [PATCH] More juicy, untested work

---
 .../Protocol1_20_2To1_20.java                 | 112 ++++++++++--------
 .../rewriter/EntityPacketRewriter1_20_2.java  |   2 -
 ...ocolState.java => ConfigurationState.java} |  53 ++++++---
 3 files changed, 95 insertions(+), 72 deletions(-)
 rename common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/{FakeProtocolState.java => ConfigurationState.java} (56%)

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 705c1f0e9..0cab27e6f 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,7 +17,6 @@
  */
 package com.viaversion.viaversion.protocols.protocol1_20_2to1_20;
 
-import com.viaversion.viaversion.api.Via;
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.minecraft.entities.Entity1_19_4Types;
 import com.viaversion.viaversion.api.protocol.AbstractProtocol;
@@ -29,7 +28,6 @@ import com.viaversion.viaversion.api.rewriter.ItemRewriter;
 import com.viaversion.viaversion.api.type.Type;
 import com.viaversion.viaversion.data.entity.EntityTrackerBase;
 import com.viaversion.viaversion.exception.CancelException;
-import com.viaversion.viaversion.protocol.packet.PacketWrapperImpl;
 import com.viaversion.viaversion.protocols.base.ClientboundLoginPackets;
 import com.viaversion.viaversion.protocols.base.ServerboundLoginPackets;
 import com.viaversion.viaversion.protocols.protocol1_19_4to1_19_3.ClientboundPackets1_19_4;
@@ -40,12 +38,10 @@ import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.packet.Serverbou
 import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.packet.ServerboundPackets1_20_2;
 import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.rewriter.BlockItemPacketRewriter1_20_2;
 import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.rewriter.EntityPacketRewriter1_20_2;
-import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.storage.FakeProtocolState;
+import com.viaversion.viaversion.protocols.protocol1_20_2to1_20.storage.ConfigurationState;
 import com.viaversion.viaversion.rewriter.SoundRewriter;
 import com.viaversion.viaversion.rewriter.StatisticsRewriter;
 import com.viaversion.viaversion.rewriter.TagRewriter;
-import io.netty.buffer.ByteBuf;
-import io.netty.buffer.ByteBufAllocator;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.util.UUID;
@@ -62,10 +58,8 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
 
     @Override
     protected void registerPackets() {
-        // TODO Handle Enabled features and tags before configuration phase end
-        // TODO Make sure Paper/Velocity handle a 0,0 uuid fine during login
-
-        // Lesser:
+        // Close your eyes and turn around while you still can
+        // TODO Handle Enabled features and tags before configuration phase end?
         // TODO Player info, replace profile with missing name with null?
         // TODO Scoreboard objective probably okay, but there are refactors to the id
         super.registerPackets();
@@ -92,31 +86,48 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
         });
 
         // Deal with the new CONFIGURATION protocol state the client expects
+        // After the game profile is received by the client, it will send its login ack,
+        // switch to the configration protocol state and send its brand.
+        // We need to wait for it send the login ack before actually sending the play login,
+        // hence packets are added to a queue. With the data from the login packet, we sent what is needed
+        // during the configuration phase before finally transitioning to the play state with the client as well.
         registerClientbound(State.LOGIN, ClientboundLoginPackets.GAME_PROFILE.getId(), ClientboundLoginPackets.GAME_PROFILE.getId(), wrapper -> {
-            wrapper.user().get(FakeProtocolState.class).setGameProfileSent(true);
+            wrapper.user().get(ConfigurationState.class).setBridgePhase(ConfigurationState.BridgePhase.PROFILE_SENT);
         });
 
         registerServerbound(State.LOGIN, ServerboundLoginPackets.LOGIN_ACKNOWLEDGED.getId(), -1, wrapper -> {
             System.out.println("Login acknowleged!");
-            wrapper.user().getProtocolInfo().setState(State.PLAY);
-            wrapper.user().get(FakeProtocolState.class).setConfigurationState(true);
             wrapper.cancel();
+            final ConfigurationState configurationState = wrapper.user().get(ConfigurationState.class);
+            configurationState.setBridgePhase(ConfigurationState.BridgePhase.CONFIGURATION);
+            wrapper.user().getProtocolInfo().setState(State.PLAY);
+
+            for (final ConfigurationState.QueuedPacket packet : configurationState.packetQueue()) {
+                final PacketWrapper queuedWrapper;
+                if (packet.packetType() != null) {
+                    queuedWrapper = PacketWrapper.create(packet.packetType(), packet.buf(), wrapper.user());
+                } else {
+                    //noinspection deprecation
+                    queuedWrapper = PacketWrapper.create(packet.packetId(), packet.buf(), wrapper.user());
+                }
+
+                queuedWrapper.send(Protocol1_20_2To1_20.class, false);
+            }
+
+            configurationState.packetQueue().clear();
         });
         cancelServerbound(State.LOGIN, ServerboundLoginPackets.CUSTOM_QUERY_ANSWER.getId()); // TODO ?
 
         // TODO Make sure this is called in other protocols as well/the base protocol
         registerServerbound(State.CONFIGURATION, ServerboundConfigurationPackets1_20_2.FINISH_CONFIGURATION.getId(), -1, wrapper -> {
-            wrapper.user().get(FakeProtocolState.class).setConfigurationState(false);
-            System.out.println("NOW PLAYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYY");
+            wrapper.user().get(ConfigurationState.class).reset();
             wrapper.cancel();
+            System.out.println("CLIENT NOW ALSO ENTERING PLAY STATE");
         });
+
         registerServerbound(State.CONFIGURATION, ServerboundConfigurationPackets1_20_2.CUSTOM_PAYLOAD.getId(), -1, wrapper -> {
-            if (wrapper.user().getProtocolInfo().getState() == State.PLAY) {
-                wrapper.setPacketType(ServerboundPackets1_19_4.PLUGIN_MESSAGE);
-            } else {
-                wrapper.user().get(FakeProtocolState.class).addPacketToQueue(wrapper, false);
-                wrapper.cancel();
-            }
+            wrapper.setPacketType(ServerboundPackets1_20_2.PLUGIN_MESSAGE);
+            wrapper.user().get(ConfigurationState.class).addPacketToQueue(wrapper, false);
         });
         registerServerbound(State.CONFIGURATION, ServerboundConfigurationPackets1_20_2.KEEP_ALIVE.getId(), ServerboundPackets1_19_4.KEEP_ALIVE.getId(), wrapper -> {
         });
@@ -125,13 +136,8 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
         registerServerbound(State.CONFIGURATION, ServerboundConfigurationPackets1_20_2.RESOURCE_PACK.getId(), ServerboundPackets1_19_4.RESOURCE_PACK_STATUS.getId(), wrapper -> {
         });
 
-        // Sad emoji
-        cancelClientbound(ClientboundPackets1_19_4.UPDATE_ENABLED_FEATURES);
-
-        registerServerbound(ServerboundPackets1_20_2.CONFIGURATION_ACKNOWLEDGED, null, wrapper -> {
-            wrapper.user().get(FakeProtocolState.class).setConfigurationState(false);
-            wrapper.cancel();
-        });
+        cancelClientbound(ClientboundPackets1_19_4.UPDATE_ENABLED_FEATURES); // Sad emoji
+        cancelServerbound(ServerboundPackets1_20_2.CONFIGURATION_ACKNOWLEDGED);
 
         // TODO Check if we can just not send batches (probably fine like this)
         cancelServerbound(ServerboundPackets1_20_2.CHUNK_BATCH_RECEIVED);
@@ -139,26 +145,31 @@ 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 FakeProtocolState fakeProtocolState = packetWrapper.user().get(FakeProtocolState.class);
-        if (fakeProtocolState.hasGameProfileBeenSent() && !fakeProtocolState.isConfigurationState()) {
-            fakeProtocolState.addPacketToQueue(packetWrapper, direction == Direction.CLIENTBOUND);
-            System.out.println("added to queue " + packetWrapper.getId() + " " + direction + " " + state);
-            super.transform(direction, State.PLAY, packetWrapper);
-            throw CancelException.generate();
+        final ConfigurationState configurationBridge = packetWrapper.user().get(ConfigurationState.class);
+        if (configurationBridge.bridgePhase() == ConfigurationState.BridgePhase.NONE) {
+            super.transform(direction, state, packetWrapper);
         }
 
-        System.out.println("Transforming " + packetWrapper.getId() + " " + direction + " " + state);
-        System.out.println(fakeProtocolState.isConfigurationState());
-        if (!fakeProtocolState.isConfigurationState()) {
-            super.transform(direction, state, packetWrapper);
+        if (direction == Direction.SERVERBOUND) {
+            // Client and server might be on two different protocol states - always let the client packets go through
+            super.transform(direction, configurationBridge.bridgePhase() == ConfigurationState.BridgePhase.CONFIGURATION
+                    ? State.CONFIGURATION : state, packetWrapper);
             return;
         }
 
-        if (direction == Direction.CLIENTBOUND
-                && (packetWrapper.getPacketType() == null || packetWrapper.getPacketType().state() != State.CONFIGURATION)) {
-            // Map some of them to their configuration state counterparts
+        // Queue packets sent by the serverwhile we wait for the client to transition to the configuration state
+        if (configurationBridge.bridgePhase() == ConfigurationState.BridgePhase.PROFILE_SENT) {
+            System.out.println("added to queue " + packetWrapper.getId() + " " + direction + " " + state);
+            configurationBridge.addPacketToQueue(packetWrapper, direction == Direction.CLIENTBOUND);
+            throw CancelException.generate();
+        }
+
+        // Map some of them to their configuration state counterparts
+        System.out.println("Transforming " + packetWrapper.getId() + " " + direction + " " + state);
+        System.out.println(configurationBridge.bridgePhase());
+        if (packetWrapper.getPacketType() == null || packetWrapper.getPacketType().state() != State.CONFIGURATION) {
             final int unmappedId = packetWrapper.getId();
-            if (state == State.PLAY && unmappedId != ClientboundPackets1_19_4.JOIN_GAME.getId()) {
+            if (state == State.PLAY) {
                 if (unmappedId == ClientboundPackets1_19_4.PLUGIN_MESSAGE.getId()) {
                     packetWrapper.setPacketType(ClientboundConfigurationPackets1_20_2.CUSTOM_PAYLOAD);
                 } else if (unmappedId == ClientboundPackets1_19_4.DISCONNECT.getId()) {
@@ -174,19 +185,16 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
                 } else if (unmappedId == ClientboundPackets1_19_4.TAGS.getId()) {
                     packetWrapper.setPacketType(ClientboundConfigurationPackets1_20_2.UPDATE_TAGS);
                 } else {
-                    // TODO Queue these packets to send them later
-                    Via.getPlatform().getLogger().warning("Cancelled packet " + packetWrapper.getId() + " sent during configuration state");
-                    packetWrapper.cancel();
-                    return;
+                    configurationBridge.addPacketToQueue(packetWrapper, direction == Direction.CLIENTBOUND);
+                    throw CancelException.generate();
                 }
-
-                super.transform(direction, State.CONFIGURATION, packetWrapper);
             }
-        } else {
-            // Redirect packets during the fake configuration phase
-            // This might mess up people using Via API/other protocols down the line, but such is life. We can't have different states for server and client
-            super.transform(direction, State.CONFIGURATION, packetWrapper);
         }
+
+
+        // Redirect packets during the fake configuration phase
+        // This might mess up people using Via API/other protocols down the line, but such is life. We can't have different states for server and client
+        super.transform(direction, State.CONFIGURATION, packetWrapper);
     }
 
     @Override
@@ -196,7 +204,7 @@ public final class Protocol1_20_2To1_20 extends AbstractProtocol<ClientboundPack
 
     @Override
     public void init(final UserConnection user) {
-        user.put(new FakeProtocolState());
+        user.put(new ConfigurationState());
         addEntityTracker(user, new EntityTrackerBase(user, Entity1_19_4Types.PLAYER));
     }
 
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 3444edc39..971bbe952 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
@@ -95,8 +95,6 @@ public final class EntityPacketRewriter1_20_2 extends EntityRewriter<Clientbound
                     final PacketWrapper finishConfigurationPacket = wrapper.create(ClientboundConfigurationPackets1_20_2.FINISH_CONFIGURATION);
                     finishConfigurationPacket.send(Protocol1_20_2To1_20.class);
 
-                    System.out.println("PLEASEEEEEEEE");
-
                     // Manually send it at the end and hope nothing breaks
                     wrapper.send(Protocol1_20_2To1_20.class);
                     wrapper.cancel();
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/FakeProtocolState.java b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
similarity index 56%
rename from common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/FakeProtocolState.java
rename to common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
index 6f433306d..cb579bb5a 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/FakeProtocolState.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/protocol1_20_2to1_20/storage/ConfigurationState.java
@@ -18,55 +18,64 @@
 package com.viaversion.viaversion.protocols.protocol1_20_2to1_20.storage;
 
 import com.viaversion.viaversion.api.connection.StorableObject;
+import com.viaversion.viaversion.api.protocol.packet.PacketType;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
-import com.viaversion.viaversion.protocol.packet.PacketWrapperImpl;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class FakeProtocolState implements StorableObject {
+public class ConfigurationState implements StorableObject {
 
     private final List<QueuedPacket> packetQueue = new ArrayList<>();
-    private boolean configurationState;
-    private boolean gameProfileSent;
+    private BridgePhase bridgePhase = BridgePhase.NONE;
 
-    public boolean isConfigurationState() {
-        return configurationState;
+    public BridgePhase bridgePhase() {
+        return bridgePhase;
     }
 
-    public void setConfigurationState(final boolean configurationState) {
-        this.configurationState = configurationState;
+    public void setBridgePhase(final BridgePhase bridgePhase) {
+        this.bridgePhase = bridgePhase;
     }
 
-    public void addPacketToQueue(final PacketWrapper wrapper, final boolean clientbound) {
-        packetQueue.add(new QueuedPacket(((PacketWrapperImpl) wrapper).getInputBuffer().copy(), clientbound)); // TODO
+    public void addPacketToQueue(final PacketWrapper wrapper, final boolean clientbound) throws Exception {
+        // Caching packet buffers is cursed, copy to heap buffers to make sure we don't start leaking in dumb cases
+        final ByteBuf copy = Unpooled.buffer();
+        wrapper.writeToBuffer(copy);
+        packetQueue.add(new QueuedPacket(copy, clientbound, wrapper.getPacketType(), wrapper.getId()));
     }
 
     public List<QueuedPacket> packetQueue() {
         return packetQueue;
     }
 
-    public boolean hasGameProfileBeenSent() {
-        return gameProfileSent;
-    }
-
-    public void setGameProfileSent(final boolean gameProfileSent) {
-        this.gameProfileSent = gameProfileSent;
+    public void reset() {
+        packetQueue.clear();
+        bridgePhase = BridgePhase.NONE;
     }
 
     @Override
     public boolean clearOnServerSwitch() {
-        return false;
+        return false; // This might be bad
+    }
+
+    public enum BridgePhase {
+        NONE, PROFILE_SENT, CONFIGURATION
     }
 
     public static final class QueuedPacket {
         private final ByteBuf buf;
         private final boolean clientbound;
+        private final PacketType packetType;
+        private final int packetId;
 
-        private QueuedPacket(final ByteBuf buf, final boolean clientbound) {
+        private QueuedPacket(final ByteBuf buf, final boolean clientbound, final PacketType packetType, final int packetId) {
             this.buf = buf;
             this.clientbound = clientbound;
+            this.packetType = packetType;
+            this.packetId = packetId;
         }
 
         public ByteBuf buf() {
@@ -76,5 +85,13 @@ public class FakeProtocolState implements StorableObject {
         public boolean clientbound() {
             return clientbound;
         }
+
+        public int packetId() {
+            return packetId;
+        }
+
+        public @Nullable PacketType packetType() {
+            return packetType;
+        }
     }
 }