From 3e2118b20e74317ea3653992082807999e98f1be Mon Sep 17 00:00:00 2001
From: Yuji Ueki <unhappychoice@gmail.com>
Date: Wed, 9 Jun 2021 08:26:10 +0900
Subject: [PATCH] :bug: Fix cannot record vanilla server

---
 .../recording/packet/PacketListener.java      | 51 +++++++++++--------
 .../mixin/MixinNetHandlerLoginClient.java     |  7 ++-
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/src/main/java/com/replaymod/recording/packet/PacketListener.java b/src/main/java/com/replaymod/recording/packet/PacketListener.java
index b8a8a576..9a2c5d95 100644
--- a/src/main/java/com/replaymod/recording/packet/PacketListener.java
+++ b/src/main/java/com/replaymod/recording/packet/PacketListener.java
@@ -103,18 +103,8 @@ public class PacketListener extends ChannelInboundHandlerAdapter {
     private long lastSentPacket;
     private long timePassedWhilePaused;
     private volatile boolean serverWasPaused;
-    //#if FABRIC>1
     private NetworkState connectionState = NetworkState.LOGIN;
     private boolean loginPhase = true;
-    //#else
-    //#if MC>=11600
-    //$$ private ProtocolType connectionState = ProtocolType.LOGIN;
-    //$$ private boolean loginPhase = true;
-    //#else
-    //$$ private EnumConnectionState connectionState = EnumConnectionState.PLAY;
-    //$$ private boolean loginPhase = false;
-    //#endif
-    //#endif
 
     /**
      * Used to keep track of the last metadata save job submitted to the save service and
@@ -232,8 +222,7 @@ public void save(Packet packet) {
 
             //#if MC>=11400
             if (packet instanceof LoginSuccessS2CPacket) {
-                connectionState = NetworkState.PLAY;
-                loginPhase = false;
+                moveToPlayState();
             }
             //#endif
         } catch(Exception e) {
@@ -465,14 +454,7 @@ private PacketData getPacketData(int timestamp, Packet packet) throws Exception
         //$$ }
         //#endif
 
-        //#if MC>=10800
-        Integer packetId = connectionState.getPacketId(NetworkSide.CLIENTBOUND, packet);
-        //#else
-        //$$ Integer packetId = (Integer) connectionState.func_150755_b().inverse().get(packet.getClass());
-        //#endif
-        if (packetId == null) {
-            throw new IOException("Unknown packet type:" + packet.getClass());
-        }
+        Integer packetId = getPacketId(packet);
         ByteBuf byteBuf = Unpooled.buffer();
         try {
             packet.write(new PacketByteBuf(byteBuf));
@@ -534,4 +516,33 @@ public long getCurrentDuration() {
     public void setServerWasPaused() {
         this.serverWasPaused = true;
     }
+
+    private Integer getPacketId(Packet packet) throws IOException {
+        //#if MC>=10800
+        Integer packetId = connectionState.getPacketId(NetworkSide.CLIENTBOUND, packet);
+        //#else
+        //$$ Integer packetId = (Integer) connectionState.func_150755_b().inverse().get(packet.getClass());
+        //#endif
+
+        if (packetId == null) {
+            /*
+               Retrying as the state is PLAY
+               TODO: Investigate packet order and refactor to be more consistent
+            */
+            packetId = NetworkState.PLAY.getPacketId(NetworkSide.CLIENTBOUND, packet);
+
+            if (packetId == null) {
+                throw new IOException("Unknown packet type:" + packet.getClass());
+            }
+
+            moveToPlayState();
+        }
+
+        return packetId;
+    }
+
+    private void moveToPlayState() {
+        connectionState = NetworkState.PLAY;
+        loginPhase = false;
+    }
 }
diff --git a/versions/1.16.4-forge/src/main/java/com/replaymod/mixin/MixinNetHandlerLoginClient.java b/versions/1.16.4-forge/src/main/java/com/replaymod/mixin/MixinNetHandlerLoginClient.java
index 91234a81..4e1a5274 100644
--- a/versions/1.16.4-forge/src/main/java/com/replaymod/mixin/MixinNetHandlerLoginClient.java
+++ b/versions/1.16.4-forge/src/main/java/com/replaymod/mixin/MixinNetHandlerLoginClient.java
@@ -19,6 +19,11 @@ public abstract class MixinNetHandlerLoginClient {
     @Final @Shadow
     private NetworkManager networkManager;
 
+    @Inject(method = "handleLoginSuccess", at=@At("HEAD"))
+    public void replayModRecording_initiateRecording(CallbackInfo cb) {
+        initiateRecording(null);
+    }
+
     /**
      * Starts the recording right before switching into PLAY state.
      * We cannot use the {@link FMLNetworkEvent.ClientConnectedToServerEvent}
@@ -35,7 +40,7 @@ private void initiateRecording(SCustomPayloadLoginPacket packet) {
             return; // already recording
         }
         ReplayModRecording.instance.initiateRecording(this.networkManager);
-        if (eventSender.getRecordingEventHandler() != null) {
+        if (eventSender.getRecordingEventHandler() != null && packet != null) {
             eventSender.getRecordingEventHandler().onPacket(packet);
         }
     }