From 4df627029ca37675d1a6e57376bd4b663fd9947b Mon Sep 17 00:00:00 2001
From: Jonas Herzig <me@johni0702.de>
Date: Sun, 19 May 2024 14:27:15 +0200
Subject: [PATCH] Fix race condition in PacketListener

The ChannelHandlerContext was only set on the next received packet,
however `save` could be called before that happened.
---
 .../handler/ConnectionEventHandler.java       |  3 ++-
 .../recording/packet/PacketListener.java      | 23 +++++--------------
 2 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/src/main/java/com/replaymod/recording/handler/ConnectionEventHandler.java b/src/main/java/com/replaymod/recording/handler/ConnectionEventHandler.java
index ef71a08b..5318013f 100644
--- a/src/main/java/com/replaymod/recording/handler/ConnectionEventHandler.java
+++ b/src/main/java/com/replaymod/recording/handler/ConnectionEventHandler.java
@@ -153,8 +153,9 @@ public class ConnectionEventHandler {
             metaData.setGenerator("ReplayMod v" + ReplayMod.instance.getVersion());
             metaData.setDate(System.currentTimeMillis());
             metaData.setMcVersion(ReplayMod.instance.getMinecraftVersion());
-            packetListener = new PacketListener(core, outputPath, replayFile, metaData);
+
             Channel channel = ((NetworkManagerAccessor) networkManager).getChannel();
+            packetListener = new PacketListener(core, channel, outputPath, replayFile, metaData);
             if (channel.pipeline().get(PacketListener.DECODER_KEY) != null) {
                 // Regular channel, we'll inject our recorder directly before the decoder
                 channel.pipeline().addBefore(PacketListener.DECODER_KEY, PacketListener.RAW_RECORDER_KEY, packetListener);
diff --git a/src/main/java/com/replaymod/recording/packet/PacketListener.java b/src/main/java/com/replaymod/recording/packet/PacketListener.java
index 3faaa5c8..c04f44bc 100644
--- a/src/main/java/com/replaymod/recording/packet/PacketListener.java
+++ b/src/main/java/com/replaymod/recording/packet/PacketListener.java
@@ -22,6 +22,7 @@ import com.replaymod.replaystudio.replay.ReplayMetaData;
 import de.johni0702.minecraft.gui.container.VanillaGuiScreen;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
@@ -102,7 +103,7 @@ public class PacketListener extends ChannelInboundHandlerAdapter {
 
     private ReplayMetaData metaData;
 
-    private ChannelHandlerContext context = null;
+    private final Channel channel;
     private Packet currentRawPacket;
 
     private final long startTime;
@@ -116,8 +117,9 @@ public class PacketListener extends ChannelInboundHandlerAdapter {
      */
     private final AtomicInteger lastSaveMetaDataId = new AtomicInteger();
 
-    public PacketListener(ReplayMod core, Path outputPath, ReplayFile replayFile, ReplayMetaData metaData) throws IOException {
+    public PacketListener(ReplayMod core, Channel channel, Path outputPath, ReplayFile replayFile, ReplayMetaData metaData) throws IOException {
         this.core = core;
+        this.channel = channel;
         this.outputPath = outputPath;
         this.replayFile = replayFile;
         this.metaData = metaData;
@@ -309,15 +311,6 @@ public class PacketListener extends ChannelInboundHandlerAdapter {
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-        if(ctx == null) {
-            if(context == null) {
-                return;
-            } else {
-                ctx = context;
-            }
-        }
-        this.context = ctx;
-
         NetworkState connectionState = getConnectionState();
 
         Packet packet = null;
@@ -367,16 +360,12 @@ public class PacketListener extends ChannelInboundHandlerAdapter {
     }
 
     private NetworkState getConnectionState() {
-        ChannelHandlerContext ctx = context;
-        if (ctx == null) {
-            return NetworkState.LOGIN;
-        }
         //#if MC>=12002
         //$$ AttributeKey<NetworkState.PacketHandler<?>> key = ClientConnection.CLIENTBOUND_PROTOCOL_KEY;
-        //$$ return ctx.channel().attr(key).get().getState();
+        //$$ return channel.attr(key).get().getState();
         //#else
         AttributeKey<NetworkState> key = ClientConnection.ATTR_KEY_PROTOCOL;
-        return ctx.channel().attr(key).get();
+        return channel.attr(key).get();
         //#endif
     }