simplify handler mixins, fixed memory leak

closes #19
badasintended · Jan 8, 2025 · d83fc41 · d83fc41
1 parent f86019c
commit d83fc41
Show file tree

Hide file tree

Showing 11 changed files with 70 additions and 83 deletions.
diff --git a/src/main/java/lol/bai/badpackets/impl/handler/ClientPlayPacketHandler.java b/src/main/java/lol/bai/badpackets/impl/handler/ClientPlayPacketHandler.java
@@ -30,7 +30,7 @@ public static ClientPlayPacketHandler get() {
             throw new IllegalStateException("Cannot get c2s sender when not in play stage!");
         }
 
-        return ((ClientPlayPacketHandler.Holder) listener).badpackets_getHandler();
+        return ((PacketHandlerHolder<ClientPlayPacketHandler>) listener).badpackets_handler();
     }
 
     @Override
@@ -51,10 +51,4 @@ protected void receiveUnsafe(ClientPlayPacketReceiver<CustomPacketPayload> recei
         receiver.receive(client, listener, payload, this);
     }
 
-    public interface Holder {
-
-        ClientPlayPacketHandler badpackets_getHandler();
-
-    }
-
 }
diff --git a/src/main/java/lol/bai/badpackets/impl/handler/PacketHandlerHolder.java b/src/main/java/lol/bai/badpackets/impl/handler/PacketHandlerHolder.java
@@ -0,0 +1,7 @@
+package lol.bai.badpackets.impl.handler;
+
+public interface PacketHandlerHolder<T extends AbstractPacketHandler<?>> {
+
+    T badpackets_handler();
+
+}
diff --git a/src/main/java/lol/bai/badpackets/impl/handler/ServerPlayPacketHandler.java b/src/main/java/lol/bai/badpackets/impl/handler/ServerPlayPacketHandler.java
@@ -26,7 +26,7 @@ public ServerPlayPacketHandler(MinecraftServer server, ServerGamePacketListenerI
     }
 
     public static ServerPlayPacketHandler get(ServerPlayer player) {
-        return ((ServerPlayPacketHandler.Holder) player.connection).badpackets_getHandler();
+        return ((PacketHandlerHolder<ServerPlayPacketHandler>) player.connection).badpackets_handler();
     }
 
     @Override
@@ -46,10 +46,4 @@ protected void receiveUnsafe(ServerPlayPacketReceiver<CustomPacketPayload> recei
         receiver.receive(server, handler.getPlayer(), handler, payload, this);
     }
 
-    public interface Holder {
-
-        ServerPlayPacketHandler badpackets_getHandler();
-
-    }
-
 }
diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/MixinConnection.java b/src/main/java/lol/bai/badpackets/impl/mixin/MixinConnection.java
@@ -0,0 +1,44 @@
+package lol.bai.badpackets.impl.mixin;
+
+import io.netty.channel.ChannelHandlerContext;
+import lol.bai.badpackets.impl.handler.PacketHandlerHolder;
+import net.minecraft.network.Connection;
+import net.minecraft.network.PacketListener;
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(Connection.class)
+public class MixinConnection {
+
+    @Shadow
+    @Nullable
+    private volatile PacketListener packetListener;
+
+    @Inject(method = "setListener", at = @At("HEAD"))
+    private void badpackets_cleanOldListener(PacketListener $$0, CallbackInfo ci) {
+        badpackets_cleanListener();
+    }
+
+    @Inject(method = "channelInactive", at = @At("HEAD"))
+    private void badpackets_cleanListener(ChannelHandlerContext $$0, CallbackInfo ci) {
+        badpackets_cleanListener();
+    }
+
+    @Inject(method = "handleDisconnection", at = @At(value = "INVOKE", target = "Lnet/minecraft/network/PacketListener;onDisconnect(Lnet/minecraft/network/chat/Component;)V"))
+    private void badpackets_cleanListener(CallbackInfo ci) {
+        badpackets_cleanListener();
+    }
+
+    @Unique
+    private void badpackets_cleanListener() {
+        if (this.packetListener instanceof PacketHandlerHolder<?> holder) {
+            holder.badpackets_handler().remove();
+        }
+    }
+
+}
diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerCommonPacketListenerImpl.java b/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerCommonPacketListenerImpl.java
@@ -1,6 +1,5 @@
 package lol.bai.badpackets.impl.mixin;
 
-import net.minecraft.network.chat.Component;
 import net.minecraft.network.protocol.common.ServerboundCustomPayloadPacket;
 import net.minecraft.network.protocol.common.ServerboundPongPacket;
 import net.minecraft.server.network.ServerCommonPacketListenerImpl;
@@ -23,20 +22,11 @@ private void badpackets_onPong(ServerboundPongPacket packet, CallbackInfo ci) {
         badpackets_onPong(packet.getId());
     }
 
-    @Inject(method = "onDisconnect", at = @At("HEAD"))
-    private void badpackets_removePacketHandler(Component reason, CallbackInfo ci) {
-        badpackets_removePacketHandler();
-    }
-
     @Unique
     protected boolean badpackets_handleCustomPayload(ServerboundCustomPayloadPacket packet) {
         return false;
     }
 
-    @Unique
-    protected void badpackets_removePacketHandler() {
-    }
-
     @Unique
     protected void badpackets_onPong(int id) {
     }

diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerConfigurationPacketListenerImpl.java b/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerConfigurationPacketListenerImpl.java
@@ -1,12 +1,10 @@
 package lol.bai.badpackets.impl.mixin;
 
-import java.util.Queue;
-
 import lol.bai.badpackets.impl.Constants;
+import lol.bai.badpackets.impl.handler.PacketHandlerHolder;
 import lol.bai.badpackets.impl.handler.ServerConfigPacketHandler;
 import net.minecraft.network.Connection;
 import net.minecraft.network.protocol.common.ServerboundCustomPayloadPacket;
-import net.minecraft.network.protocol.configuration.ServerboundFinishConfigurationPacket;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.network.CommonListenerCookie;
 import net.minecraft.server.network.ConfigurationTask;
@@ -22,8 +20,10 @@
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
 
+import java.util.Queue;
+
 @Mixin(ServerConfigurationPacketListenerImpl.class)
-public abstract class MixinServerConfigurationPacketListenerImpl extends MixinServerCommonPacketListenerImpl implements ServerConfigPacketHandler.TaskFinisher {
+public abstract class MixinServerConfigurationPacketListenerImpl extends MixinServerCommonPacketListenerImpl implements ServerConfigPacketHandler.TaskFinisher, PacketHandlerHolder<ServerConfigPacketHandler> {
 
     @Shadow
     @Final
@@ -50,11 +50,6 @@ private void badpackets_initPacketHandler(CallbackInfo ci) {
         configurationTasks.addAll(ServerConfigPacketHandler.CUSTOM_TASKS.values());
     }
 
-    @Inject(method = "handleConfigurationFinished", at = @At("RETURN"))
-    private void badpackets_removePacketHandler(ServerboundFinishConfigurationPacket packet, CallbackInfo ci) {
-        badpackets_packetHandler.remove();
-    }
-
     @Inject(method = "startNextTask", locals = LocalCapture.CAPTURE_FAILHARD, at = @At(value = "FIELD", opcode = Opcodes.PUTFIELD, target = "Lnet/minecraft/server/network/ServerConfigurationPacketListenerImpl;currentTask:Lnet/minecraft/server/network/ConfigurationTask;"))
     private void badpackets_attachCustomTaskContext(CallbackInfo ci, ConfigurationTask task) {
         if (task instanceof ServerConfigPacketHandler.CustomTask custom) {
@@ -70,8 +65,8 @@ private void badpackets_detachCustomTaskContext(ConfigurationTask.Type type, Cal
     }
 
     @Override
-    protected void badpackets_removePacketHandler() {
-        badpackets_packetHandler.remove();
+    public ServerConfigPacketHandler badpackets_handler() {
+        return badpackets_packetHandler;
     }
 
     @Override

diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerGamePacketListenerImpl.java b/src/main/java/lol/bai/badpackets/impl/mixin/MixinServerGamePacketListenerImpl.java
@@ -1,10 +1,9 @@
 package lol.bai.badpackets.impl.mixin;
 
+import lol.bai.badpackets.impl.handler.PacketHandlerHolder;
 import lol.bai.badpackets.impl.handler.ServerPlayPacketHandler;
 import net.minecraft.network.Connection;
-import net.minecraft.network.chat.Component;
 import net.minecraft.network.protocol.common.ServerboundCustomPayloadPacket;
-import net.minecraft.network.protocol.game.ServerboundConfigurationAcknowledgedPacket;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.server.network.CommonListenerCookie;
@@ -16,7 +15,7 @@
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 @Mixin(ServerGamePacketListenerImpl.class)
-public class MixinServerGamePacketListenerImpl extends MixinServerCommonPacketListenerImpl implements ServerPlayPacketHandler.Holder {
+public class MixinServerGamePacketListenerImpl extends MixinServerCommonPacketListenerImpl implements PacketHandlerHolder<ServerPlayPacketHandler> {
 
     @Unique
     private ServerPlayPacketHandler badpacket_packetHandler;
@@ -26,23 +25,13 @@ private void badpackets_createServerPacketHandler(MinecraftServer minecraftServe
         badpacket_packetHandler = new ServerPlayPacketHandler(minecraftServer, (ServerGamePacketListenerImpl) (Object) this, connection);
     }
 
-    @Inject(method = "onDisconnect", at = @At("HEAD"))
-    private void badpackets_removeServerPacketHandler(Component reason, CallbackInfo ci) {
-        badpacket_packetHandler.remove();
-    }
-
-    @Inject(method = "handleConfigurationAcknowledged", at = @At("RETURN"))
-    private void badpacekts_removePacketHandler(ServerboundConfigurationAcknowledgedPacket packet, CallbackInfo ci) {
-        badpacket_packetHandler.remove();
-    }
-
     @Override
     protected boolean badpackets_handleCustomPayload(ServerboundCustomPayloadPacket packet) {
         return badpacket_packetHandler.receive(packet.payload());
     }
 
     @Override
-    public ServerPlayPacketHandler badpackets_getHandler() {
+    public ServerPlayPacketHandler badpackets_handler() {
         return badpacket_packetHandler;
     }
 

diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/client/MixinClientCommonPacketListenerImpl.java b/src/main/java/lol/bai/badpackets/impl/mixin/client/MixinClientCommonPacketListenerImpl.java
@@ -1,7 +1,6 @@
 package lol.bai.badpackets.impl.mixin.client;
 
 import net.minecraft.client.multiplayer.ClientCommonPacketListenerImpl;
-import net.minecraft.network.chat.Component;
 import net.minecraft.network.protocol.common.ClientboundCustomPayloadPacket;
 import net.minecraft.network.protocol.common.ClientboundPingPacket;
 import org.spongepowered.asm.mixin.Mixin;
@@ -13,11 +12,6 @@
 @Mixin(ClientCommonPacketListenerImpl.class)
 public abstract class MixinClientCommonPacketListenerImpl {
 
-    @Inject(method = "onDisconnect", at = @At("HEAD"))
-    private void badpackets_removeClientPacketHandler(Component reason, CallbackInfo ci) {
-        badpackets_removeClientPacketHandler(reason);
-    }
-
     @Inject(method = "handlePing", at = @At("HEAD"), cancellable = true)
     private void badpackets_handlePing(ClientboundPingPacket packet, CallbackInfo ci) {
         if (badpackets_handlePing(packet.getId())) ci.cancel();
@@ -28,10 +22,6 @@ private void badpackets_handleCustomPayload(ClientboundCustomPayloadPacket packe
         if (badpackets_handleCustomPayload(packet)) ci.cancel();
     }
 
-    @Unique
-    protected void badpackets_removeClientPacketHandler(Component reason) {
-    }
-
     @Unique
     protected boolean badpackets_handlePing(int id) {
         return false;

diff --git a/...java/lol/bai/badpackets/impl/mixin/client/MixinClientConfigurationPacketListenerImpl.java b/...java/lol/bai/badpackets/impl/mixin/client/MixinClientConfigurationPacketListenerImpl.java
@@ -2,21 +2,20 @@
 
 import lol.bai.badpackets.impl.Constants;
 import lol.bai.badpackets.impl.handler.ClientConfigPacketHandler;
+import lol.bai.badpackets.impl.handler.PacketHandlerHolder;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.multiplayer.ClientConfigurationPacketListenerImpl;
 import net.minecraft.client.multiplayer.CommonListenerCookie;
 import net.minecraft.network.Connection;
-import net.minecraft.network.chat.Component;
 import net.minecraft.network.protocol.common.ClientboundCustomPayloadPacket;
-import net.minecraft.network.protocol.configuration.ClientboundFinishConfigurationPacket;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 @Mixin(ClientConfigurationPacketListenerImpl.class)
-public class MixinClientConfigurationPacketListenerImpl extends MixinClientCommonPacketListenerImpl {
+public class MixinClientConfigurationPacketListenerImpl extends MixinClientCommonPacketListenerImpl implements PacketHandlerHolder<ClientConfigPacketHandler> {
 
     @Unique
     private ClientConfigPacketHandler badpackets_packetHandler;
@@ -26,14 +25,9 @@ private void badpackets_createPacketHandler(Minecraft minecraft, Connection conn
         badpackets_packetHandler = new ClientConfigPacketHandler(minecraft, (ClientConfigurationPacketListenerImpl) (Object) this, connection);
     }
 
-    @Inject(method = "handleConfigurationFinished", at = @At("RETURN"))
-    private void badpacekts_removePacketHandler(ClientboundFinishConfigurationPacket $$0, CallbackInfo ci) {
-        badpackets_packetHandler.remove();
-    }
-
     @Override
-    protected void badpackets_removeClientPacketHandler(Component reason) {
-        badpackets_packetHandler.remove();
+    public ClientConfigPacketHandler badpackets_handler() {
+        return badpackets_packetHandler;
     }
 
     @Override

diff --git a/src/main/java/lol/bai/badpackets/impl/mixin/client/MixinClientPacketListener.java b/src/main/java/lol/bai/badpackets/impl/mixin/client/MixinClientPacketListener.java
@@ -1,21 +1,20 @@
 package lol.bai.badpackets.impl.mixin.client;
 
 import lol.bai.badpackets.impl.handler.ClientPlayPacketHandler;
+import lol.bai.badpackets.impl.handler.PacketHandlerHolder;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.multiplayer.ClientPacketListener;
 import net.minecraft.client.multiplayer.CommonListenerCookie;
 import net.minecraft.network.Connection;
-import net.minecraft.network.chat.Component;
 import net.minecraft.network.protocol.common.ClientboundCustomPayloadPacket;
-import net.minecraft.network.protocol.game.ClientboundStartConfigurationPacket;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 @Mixin(ClientPacketListener.class)
-public abstract class MixinClientPacketListener extends MixinClientCommonPacketListenerImpl implements ClientPlayPacketHandler.Holder {
+public abstract class MixinClientPacketListener extends MixinClientCommonPacketListenerImpl implements PacketHandlerHolder<ClientPlayPacketHandler> {
 
     @Unique
     private ClientPlayPacketHandler badpacket_packetHandler;
@@ -26,17 +25,7 @@ private void badpackets_createClientPacketHandler(Minecraft client, Connection c
     }
 
     @Override
-    protected void badpackets_removeClientPacketHandler(Component reason) {
-        badpacket_packetHandler.remove();
-    }
-
-    @Inject(method = "handleConfigurationStart", at = @At("RETURN"))
-    private void badpackets_removeHandler(ClientboundStartConfigurationPacket $$0, CallbackInfo ci) {
-        badpacket_packetHandler.remove();
-    }
-
-    @Override
-    public ClientPlayPacketHandler badpackets_getHandler() {
+    public ClientPlayPacketHandler badpackets_handler() {
         return badpacket_packetHandler;
     }
 

diff --git a/src/main/resources/badpackets.mixins.json b/src/main/resources/badpackets.mixins.json
@@ -6,6 +6,7 @@
   "refmap"            : "badpackets.refmap.json",
   "mixins"            : [
     "AccessServerboundCustomPayloadPacket",
+    "MixinConnection",
     "MixinPlayerList",
     "MixinServerboundCustomPayloadPacket",
     "MixinServerCommonPacketListenerImpl",