Some fixes to c2c packets

Earthcomputer · May 29, 2024 · d6a4759 · d6a4759
1 parent eb7e3a7
commit d6a4759
Show file tree

Hide file tree

Showing 8 changed files with 155 additions and 75 deletions.
diff --git a/src/main/java/net/earthcomputer/clientcommands/c2c/C2CPacketHandler.java b/src/main/java/net/earthcomputer/clientcommands/c2c/C2CPacketHandler.java
@@ -3,11 +3,16 @@
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
 import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
 import com.mojang.brigadier.exceptions.SimpleCommandExceptionType;
+import com.mojang.logging.LogUtils;
+import io.netty.buffer.Unpooled;
 import net.earthcomputer.clientcommands.c2c.packets.MessageC2CPacket;
 import net.earthcomputer.clientcommands.command.ListenCommand;
+import net.earthcomputer.clientcommands.interfaces.IClientPacketListener_C2C;
 import net.fabricmc.fabric.api.networking.v1.PacketByteBufs;
 import net.minecraft.ChatFormatting;
 import net.minecraft.client.Minecraft;
+import net.minecraft.client.multiplayer.AccountProfileKeyPairManager;
+import net.minecraft.client.multiplayer.ClientPacketListener;
 import net.minecraft.client.multiplayer.PlayerInfo;
 import net.minecraft.network.ConnectionProtocol;
 import net.minecraft.network.FriendlyByteBuf;
@@ -19,19 +24,25 @@
 import net.minecraft.network.protocol.Packet;
 import net.minecraft.network.protocol.PacketFlow;
 import net.minecraft.network.protocol.ProtocolInfoBuilder;
+import net.minecraft.world.entity.player.ProfileKeyPair;
 import net.minecraft.world.entity.player.ProfilePublicKey;
+import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
 
 import java.security.PublicKey;
+import java.util.Arrays;
+import java.util.Optional;
 
 public class C2CPacketHandler implements C2CPacketListener {
+    private static final Logger LOGGER = LogUtils.getLogger();
 
     private static final DynamicCommandExceptionType MESSAGE_TOO_LONG_EXCEPTION = new DynamicCommandExceptionType(d -> Component.translatable("c2cpacket.messageTooLong", d));
     private static final SimpleCommandExceptionType PUBLIC_KEY_NOT_FOUND_EXCEPTION = new SimpleCommandExceptionType(Component.translatable("c2cpacket.publicKeyNotFound"));
     private static final SimpleCommandExceptionType ENCRYPTION_FAILED_EXCEPTION = new SimpleCommandExceptionType(Component.translatable("c2cpacket.encryptionFailed"));
 
-    public static final ProtocolInfo<C2CPacketListener> C2C = ProtocolInfoBuilder.<C2CPacketListener, RegistryFriendlyByteBuf>protocolUnbound(ConnectionProtocol.PLAY, PacketFlow.CLIENTBOUND, builder -> builder
+    public static final ProtocolInfo.Unbound<C2CPacketListener, RegistryFriendlyByteBuf> PROTOCOL_UNBOUND = ProtocolInfoBuilder.protocolUnbound(ConnectionProtocol.PLAY, PacketFlow.CLIENTBOUND, builder -> builder
         .addPacket(MessageC2CPacket.ID, MessageC2CPacket.CODEC)
-    ).bind(RegistryFriendlyByteBuf.decorator(Minecraft.getInstance().getConnection().registryAccess()));
+    );
 
     private static final C2CPacketHandler instance = new C2CPacketHandler();
 
@@ -53,7 +64,11 @@ public void sendPacket(Packet<C2CPacketListener> packet, PlayerInfo recipient) t
         }
         PublicKey key = ppk.data().key();
         FriendlyByteBuf buf = PacketByteBufs.create();
-        C2C.codec().encode(buf, packet);
+        ProtocolInfo<C2CPacketListener> protocolInfo = getCurrentProtocolInfo();
+        if (protocolInfo == null) {
+            return;
+        }
+        protocolInfo.codec().encode(buf, packet);
         byte[] uncompressed = new byte[buf.readableBytes()];
         buf.getBytes(0, uncompressed);
         byte[] compressed = ConversionHelper.Gzip.compress(uncompressed);
@@ -102,14 +117,85 @@ public void onMessageC2CPacket(MessageC2CPacket packet) {
         Minecraft.getInstance().gui.getChat().addMessage(component);
     }
 
+    public static boolean handleC2CPacket(String content) {
+        byte[] encrypted = ConversionHelper.BaseUTF8.fromUnicode(content);
+        // round down to multiple of 256 bytes
+        int length = encrypted.length & ~0xFF;
+        // copy to new array of arrays
+        byte[][] encryptedArrays = new byte[length / 256][];
+        for (int i = 0; i < length; i += 256) {
+            encryptedArrays[i / 256] = Arrays.copyOfRange(encrypted, i, i + 256);
+        }
+        if (!(Minecraft.getInstance().getProfileKeyPairManager() instanceof AccountProfileKeyPairManager profileKeyPairManager)) {
+            return false;
+        }
+        Optional<ProfileKeyPair> keyPair = profileKeyPairManager.keyPair.join();
+        if (keyPair.isEmpty()) {
+            return false;
+        }
+        // decrypt
+        int len = 0;
+        byte[][] decryptedArrays = new byte[encryptedArrays.length][];
+        for (int i = 0; i < encryptedArrays.length; i++) {
+            decryptedArrays[i] = ConversionHelper.RsaEcb.decrypt(encryptedArrays[i], keyPair.get().privateKey());
+            if (decryptedArrays[i] == null) {
+                return false;
+            }
+            len += decryptedArrays[i].length;
+        }
+        // copy to new array
+        byte[] decrypted = new byte[len];
+        int pos = 0;
+        for (byte[] decryptedArray : decryptedArrays) {
+            System.arraycopy(decryptedArray, 0, decrypted, pos, decryptedArray.length);
+            pos += decryptedArray.length;
+        }
+        byte[] uncompressed = ConversionHelper.Gzip.decompress(decrypted);
+        if (uncompressed == null) {
+            return false;
+        }
+        ProtocolInfo<C2CPacketListener> protocolInfo = getCurrentProtocolInfo();
+        if (protocolInfo == null) {
+            return false;
+        }
+        FriendlyByteBuf buf = new FriendlyByteBuf(Unpooled.wrappedBuffer(uncompressed));
+        Packet<? super C2CPacketListener> packet;
+        try {
+            packet = protocolInfo.codec().decode(buf);
+        } catch (Throwable e) {
+            LOGGER.error("Error decoding C2C packet", e);
+            return false;
+        }
+        if (buf.readableBytes() > 0) {
+            return false;
+        }
+        ListenCommand.onPacket(packet, ListenCommand.PacketFlow.C2C_INBOUND);
+        try {
+            packet.handle(C2CPacketHandler.getInstance());
+        } catch (Throwable e) {
+            Minecraft.getInstance().gui.getChat().addMessage(Component.nullToEmpty(e.getMessage()));
+            LogUtils.getLogger().error("Error handling C2C packet", e);
+        }
+        return true;
+    }
+
+    @Nullable
+    public static ProtocolInfo<C2CPacketListener> getCurrentProtocolInfo() {
+        ClientPacketListener connection = Minecraft.getInstance().getConnection();
+        if (connection == null) {
+            return null;
+        }
+        return ((IClientPacketListener_C2C) connection).clientcommands_getC2CProtocolInfo();
+    }
+
     @Override
     public PacketFlow flow() {
-        return C2C.flow();
+        return PacketFlow.CLIENTBOUND;
     }
 
     @Override
     public ConnectionProtocol protocol() {
-        return C2C.id();
+        return ConnectionProtocol.PLAY;
     }
 
     @Override

diff --git a/src/main/java/net/earthcomputer/clientcommands/c2c/ConversionHelper.java b/src/main/java/net/earthcomputer/clientcommands/c2c/ConversionHelper.java
@@ -1,5 +1,9 @@
 package net.earthcomputer.clientcommands.c2c;
 
+import com.mojang.logging.LogUtils;
+import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+
 import javax.crypto.Cipher;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -12,6 +16,7 @@
 import java.util.zip.GZIPOutputStream;
 
 public class ConversionHelper {
+    private static final Logger LOGGER = LogUtils.getLogger();
 
     /**
      * @author Wagyourtail
@@ -91,20 +96,21 @@ private static int unmapCodepoint(int codepoint) {
      * @author Wagyourtail
      */
     public static class Gzip {
-        public static byte[] compress(byte[] bytes) {
+        public static byte @Nullable [] compress(byte[] bytes) {
             if (bytes == null || bytes.length == 0) {
                 return null;
             }
             ByteArrayOutputStream out = new ByteArrayOutputStream();
             try (GZIPOutputStream gzip = new GZIPOutputStream(out)) {
                 gzip.write(bytes);
             } catch (IOException e) {
-                e.printStackTrace();
+                LOGGER.error("Error compressing", e);
+                return null;
             }
             return out.toByteArray();
         }
 
-        public static byte[] uncompress(byte[] bytes) {
+        public static byte @Nullable [] decompress(byte[] bytes) {
             if (bytes == null || bytes.length == 0) {
                 return null;
             }
@@ -117,7 +123,8 @@ public static byte[] uncompress(byte[] bytes) {
                     out.write(buffer, 0, n);
                 }
             } catch (IOException e) {
-                e.printStackTrace();
+                LOGGER.error("Error decompressing", e);
+                return null;
             }
             return out.toByteArray();
         }
@@ -130,7 +137,7 @@ public static byte[] encrypt(byte[] bytes, PublicKey key) {
                 cipher.init(Cipher.ENCRYPT_MODE, key);
                 return cipher.doFinal(bytes);
             } catch (GeneralSecurityException e) {
-                e.printStackTrace();
+                LOGGER.error("Error encrypting", e);
                 return null;
             }
         }
@@ -141,7 +148,7 @@ public static byte[] decrypt(byte[] bytes, PrivateKey key) {
                 cipher.init(Cipher.DECRYPT_MODE, key);
                 return cipher.doFinal(bytes);
             } catch (GeneralSecurityException e) {
-                e.printStackTrace();
+                LOGGER.error("Error decrypting", e);
                 return null;
             }
         }

diff --git a/src/main/java/net/earthcomputer/clientcommands/command/arguments/PacketTypeArgument.java b/src/main/java/net/earthcomputer/clientcommands/command/arguments/PacketTypeArgument.java
@@ -10,6 +10,7 @@
 import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import io.netty.channel.ChannelPipeline;
 import net.earthcomputer.clientcommands.c2c.C2CPacketHandler;
+import net.earthcomputer.clientcommands.c2c.C2CPacketListener;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.multiplayer.ClientPacketListener;
@@ -86,12 +87,17 @@ private static PacketTypes get() {
             if (connection == null) {
                 return null;
             }
+            ProtocolInfo<C2CPacketListener> protocolInfo = C2CPacketHandler.getCurrentProtocolInfo();
+            if (protocolInfo == null) {
+                return null;
+            }
+
             ChannelPipeline pipeline = connection.getConnection().channel.pipeline();
             var decoder = (PacketDecoder<?>) pipeline.get("decoder");
             var clientbound = packetTypesCache.computeIfAbsent(decoder, k -> getPacketTypes(decoder.protocolInfo));
             var encoder = (PacketEncoder<?>) pipeline.get("encoder");
             var serverbound = packetTypesCache.computeIfAbsent(encoder, k -> getPacketTypes(encoder.protocolInfo));
-            var c2cbound = packetTypesCache.computeIfAbsent("c2c", k -> getPacketTypes(C2CPacketHandler.C2C));
+            var c2cbound = packetTypesCache.computeIfAbsent("c2c", k -> getPacketTypes(protocolInfo));
             return new PacketTypes(clientbound, serverbound, c2cbound);
         }
     }

diff --git a/src/main/java/net/earthcomputer/clientcommands/features/PacketDumper.java b/src/main/java/net/earthcomputer/clientcommands/features/PacketDumper.java
@@ -28,6 +28,7 @@
 import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.network.PacketDecoder;
 import net.minecraft.network.PacketEncoder;
+import net.minecraft.network.ProtocolInfo;
 import net.minecraft.network.codec.StreamCodec;
 import net.minecraft.network.codec.StreamEncoder;
 import net.minecraft.network.protocol.Packet;
@@ -76,8 +77,12 @@ public static void dumpPacket(Packet<?> packet, JsonWriter writer) throws IOExce
         writer.beginArray();
         try {
             if (packet.type().id().getNamespace().equals("clientcommands")) {
+                ProtocolInfo<C2CPacketListener> protocolInfo = C2CPacketHandler.getCurrentProtocolInfo();
+                if (protocolInfo == null) {
+                    throw new IOException("Not currently logged in");
+                }
                 //noinspection unchecked
-                C2CPacketHandler.C2C.codec().encode(new PacketDumpByteBuf(writer), (Packet<? super C2CPacketListener>) packet);
+                protocolInfo.codec().encode(new PacketDumpByteBuf(writer), (Packet<? super C2CPacketListener>) packet);
             } else {
                 ChannelPipeline pipeline = Minecraft.getInstance().getConnection().getConnection().channel.pipeline();
                 @SuppressWarnings("unchecked")

diff --git a/src/main/java/net/earthcomputer/clientcommands/interfaces/IClientPacketListener_C2C.java b/src/main/java/net/earthcomputer/clientcommands/interfaces/IClientPacketListener_C2C.java
@@ -0,0 +1,8 @@
+package net.earthcomputer.clientcommands.interfaces;
+
+import net.earthcomputer.clientcommands.c2c.C2CPacketListener;
+import net.minecraft.network.ProtocolInfo;
+
+public interface IClientPacketListener_C2C {
+    ProtocolInfo<C2CPacketListener> clientcommands_getC2CProtocolInfo();
+}
diff --git a/src/main/java/net/earthcomputer/clientcommands/mixin/c2c/ChatComponentMixin.java b/src/main/java/net/earthcomputer/clientcommands/mixin/c2c/ChatComponentMixin.java
@@ -1,23 +1,15 @@
 package net.earthcomputer.clientcommands.mixin.c2c;
 
-import io.netty.buffer.Unpooled;
 import net.earthcomputer.clientcommands.Configs;
 import net.earthcomputer.clientcommands.c2c.C2CPacketHandler;
-import net.earthcomputer.clientcommands.c2c.C2CPacketListener;
-import net.earthcomputer.clientcommands.c2c.ConversionHelper;
 import net.earthcomputer.clientcommands.c2c.OutgoingPacketFilter;
-import net.earthcomputer.clientcommands.command.ListenCommand;
 import net.minecraft.ChatFormatting;
 import net.minecraft.client.GuiMessageTag;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.gui.components.ChatComponent;
-import net.minecraft.client.multiplayer.AccountProfileKeyPairManager;
-import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.network.chat.Component;
 import net.minecraft.network.chat.HoverEvent;
 import net.minecraft.network.chat.MessageSignature;
-import net.minecraft.network.protocol.Packet;
-import net.minecraft.world.entity.player.ProfileKeyPair;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -26,9 +18,6 @@
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
-import java.util.Arrays;
-import java.util.Optional;
-
 @Mixin(ChatComponent.class)
 public class ChatComponentMixin {
 
@@ -60,60 +49,10 @@ private void handleIfPacket(Component content, CallbackInfo ci) {
             ci.cancel();
             return;
         }
-        if (handleC2CPacket(packetString)) {
+        if (C2CPacketHandler.handleC2CPacket(packetString)) {
             ci.cancel();
         } else {
             this.minecraft.gui.getChat().addMessage(Component.translatable("c2cpacket.malformedPacket").withStyle(ChatFormatting.RED));
         }
     }
-
-    @Unique
-    private static boolean handleC2CPacket(String content) {
-        byte[] encrypted = ConversionHelper.BaseUTF8.fromUnicode(content);
-        // round down to multiple of 256 bytes
-        int length = encrypted.length & ~0xFF;
-        // copy to new array of arrays
-        byte[][] encryptedArrays = new byte[length / 256][];
-        for (int i = 0; i < length; i += 256) {
-            encryptedArrays[i / 256] = Arrays.copyOfRange(encrypted, i, i + 256);
-        }
-        if (!(Minecraft.getInstance().getProfileKeyPairManager() instanceof AccountProfileKeyPairManager profileKeyPairManager)) {
-            return false;
-        }
-        Optional<ProfileKeyPair> keyPair = profileKeyPairManager.keyPair.join();
-        if (keyPair.isEmpty()) {
-            return false;
-        }
-        // decrypt
-        int len = 0;
-        byte[][] decryptedArrays = new byte[encryptedArrays.length][];
-        for (int i = 0; i < encryptedArrays.length; i++) {
-            decryptedArrays[i] = ConversionHelper.RsaEcb.decrypt(encryptedArrays[i], keyPair.get().privateKey());
-            if (decryptedArrays[i] == null) {
-                return false;
-            }
-            len += decryptedArrays[i].length;
-        }
-        // copy to new array
-        byte[] decrypted = new byte[len];
-        int pos = 0;
-        for (byte[] decryptedArray : decryptedArrays) {
-            System.arraycopy(decryptedArray, 0, decrypted, pos, decryptedArray.length);
-            pos += decryptedArray.length;
-        }
-        byte[] uncompressed = ConversionHelper.Gzip.uncompress(decrypted);
-        FriendlyByteBuf buf = new FriendlyByteBuf(Unpooled.wrappedBuffer(uncompressed));
-        Packet<? super C2CPacketListener> packet = C2CPacketHandler.C2C.codec().decode(buf);
-        if (buf.readableBytes() > 0) {
-            return false;
-        }
-        ListenCommand.onPacket(packet, ListenCommand.PacketFlow.C2C_INBOUND);
-        try {
-            packet.handle(C2CPacketHandler.getInstance());
-        } catch (Exception e) {
-            Minecraft.getInstance().gui.getChat().addMessage(Component.nullToEmpty(e.getMessage()));
-            e.printStackTrace();
-        }
-        return true;
-    }
 }
diff --git a/src/main/java/net/earthcomputer/clientcommands/mixin/c2c/ClientPacketListenerMixin.java b/src/main/java/net/earthcomputer/clientcommands/mixin/c2c/ClientPacketListenerMixin.java
@@ -0,0 +1,28 @@
+package net.earthcomputer.clientcommands.mixin.c2c;
+
+import net.earthcomputer.clientcommands.c2c.C2CPacketHandler;
+import net.earthcomputer.clientcommands.c2c.C2CPacketListener;
+import net.earthcomputer.clientcommands.interfaces.IClientPacketListener_C2C;
+import net.minecraft.client.multiplayer.ClientPacketListener;
+import net.minecraft.core.RegistryAccess;
+import net.minecraft.network.ProtocolInfo;
+import net.minecraft.network.RegistryFriendlyByteBuf;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
+
+@Mixin(ClientPacketListener.class)
+public class ClientPacketListenerMixin implements IClientPacketListener_C2C {
+    @Shadow
+    @Final
+    private RegistryAccess.Frozen registryAccess;
+
+    @Unique
+    private final ProtocolInfo<C2CPacketListener> c2cProtocolInfo = C2CPacketHandler.PROTOCOL_UNBOUND.bind(RegistryFriendlyByteBuf.decorator(registryAccess));
+
+    @Override
+    public ProtocolInfo<C2CPacketListener> clientcommands_getC2CProtocolInfo() {
+        return c2cProtocolInfo;
+    }
+}