fix packet type collection in new versions

Author: derklaro

src/main/java/dev/derklaro/protocolgenerator/protocol/McClassNames.java

@@ -29,8 +29,7 @@ final class McClassNames {
   public static final String BOOTSTRAP = "net.minecraft.server.Bootstrap";
   public static final String SHARED_CONSTANTS = "net.minecraft.SharedConstants";
 
-  public static final String REGISTRY_ACCESS = "net.minecraft.core.RegistryAccess";
-  public static final String REGISTRY_FRIENDLY_BB = "net.minecraft.network.RegistryFriendlyByteBuf";
+  public static final String PACKET_VISITOR = "net.minecraft.network.ProtocolInfo$Details$PacketVisitor";
 
   private McClassNames() {
     throw new UnsupportedOperationException();

src/main/java/dev/derklaro/protocolgenerator/protocol/McProtocolsDecoder.java

@@ -26,6 +26,7 @@
 
 import dev.derklaro.reflexion.Reflexion;
 import dev.derklaro.reflexion.matcher.FieldMatcher;
+import dev.derklaro.reflexion.matcher.MethodMatcher;
 import java.io.IOException;
 import java.lang.reflect.Modifier;
 import java.util.EnumMap;
@@ -34,10 +35,12 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import lombok.NonNull;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.Handle;
+import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.ClassNode;
 import org.objectweb.asm.tree.FieldInsnNode;
@@ -47,13 +50,30 @@
 
 final class McProtocolsDecoder {
 
+  // matcher for the UnboundProtocol constants in XXXProtocols classes
   private static final FieldMatcher PROTOCOL_INFO_MATCHER = FieldMatcher.newMatcher()
     .hasModifier(Modifier.STATIC)
     .and(field -> {
-      var name = field.getType().getCanonicalName();
-      return name.endsWith(".ProtocolInfo") || name.endsWith(".ProtocolInfo.Unbound");
+      var typeName = field.getType().getCanonicalName();
+      return typeName.endsWith(".UnboundProtocol") || typeName.endsWith(".SimpleUnboundProtocol");
     });
 
+  // matcher to find the listPackets(PacketVisitor) method in ProtocolInfo$Details
+  private static final MethodMatcher LIST_PACKETS_MATCHER = MethodMatcher.newMatcher()
+    .parameterCount(1)
+    .hasName("listPackets");
+
+  // matcher for InvokeDynamicInsnNode that constructs the info for a protocol
+  private static final Predicate<InvokeDynamicInsnNode> PROTOCOL_LAMBDA_BUILDER_FILTER = node -> {
+    var bsmArgs = node.bsmArgs;
+    return bsmArgs != null
+      && bsmArgs.length == 3
+      && bsmArgs[2] instanceof Type type
+      && type.getArgumentCount() == 1
+      && type.getReturnType() == Type.VOID_TYPE
+      && type.getArgumentTypes()[0].getClassName().endsWith(".ProtocolInfoBuilder");
+  };
+
   private static final String STREAM_CODEC_DESC = "Lnet/minecraft/network/codec/StreamCodec;";
   private static final String PACKET_TYPE_DESC = "Lnet/minecraft/network/protocol/PacketType;";
 
@@ -75,8 +95,16 @@ public McProtocolsDecoder(@NonNull Class<?> targetClass, @NonNull ClassLoader mc
     }
 
     // there can be two calls: one for serverbound, one for clientbound
-    var firstInvokeDynamic = this.findNodeOfType(InvokeDynamicInsnNode.class, clinit, null);
-    var secondInvokeDynamic = this.findNodeOfType(InvokeDynamicInsnNode.class, clinit, firstInvokeDynamic);
+    var firstInvokeDynamic = this.findNodeOfType(
+      InvokeDynamicInsnNode.class,
+      clinit,
+      null,
+      PROTOCOL_LAMBDA_BUILDER_FILTER);
+    var secondInvokeDynamic = this.findNodeOfType(
+      InvokeDynamicInsnNode.class,
+      clinit,
+      firstInvokeDynamic,
+      PROTOCOL_LAMBDA_BUILDER_FILTER);
 
     // resolve the protocol info (ProtocolInfoBuilder type -> Invocation Target for configuration)
     var firstProtocolInfo = this.resolveProtocolInvocationInfo(clinit, firstInvokeDynamic);
@@ -90,36 +118,34 @@ public McProtocolsDecoder(@NonNull Class<?> targetClass, @NonNull ClassLoader mc
     return result;
   }
 
-  public @NonNull Map<McPacketFlow, Map<Object, Integer>> resolvePacketIds() {
+  public @NonNull Map<McPacketFlow, Map<Object, Integer>> resolvePacketIds(@NonNull Class<?> packetVisitorClass) {
     Map<McPacketFlow, Map<Object, Integer>> result = new EnumMap<>(McPacketFlow.class);
 
     // resolve the packet ids from the constants in the given class
     var protocolInfoFields = Reflexion.on(this.targetClass).findFields(PROTOCOL_INFO_MATCHER);
     for (var protocolInfoField : protocolInfoFields) {
+      // get the details from the protocol info
       var protocolInfo = protocolInfoField.getValue().get();
-      if (McUnboundProtocolBinder.isUnbound(protocolInfo.getClass())) {
-        // need to bind protocol info first
-        var binder = new McUnboundProtocolBinder(protocolInfo, this.mcClassLoader);
-        protocolInfo = binder.bind();
-      }
-
-      // get codec and mc packet flow
-      var codec = Reflexion.onBound(protocolInfo)
-        .findMethod("codec")
+      var details = Reflexion.onBound(protocolInfo)
+        .findMethod("details")
         .flatMap(accessor -> accessor.invoke().asOptional())
         .orElseThrow();
-      var flow = Reflexion.onBound(protocolInfo)
+
+      // resolve the packet list debug method & flow method
+      var detailsReflexion = Reflexion.onBound(details);
+      var listPacketsMethod = detailsReflexion.findMethod(LIST_PACKETS_MATCHER).orElseThrow();
+      var flow = detailsReflexion
         .findMethod("flow")
         .flatMap(accessor -> accessor.invoke().asOptional())
         .orElseThrow();
 
-      // try to convert the packet flow
+      // convert the packet flow, register the packet type to id mapping for the flow
       var convertedFlow = this.resolveFlowFromBuilderType(flow.toString());
       if (convertedFlow != null) {
-        var typeToId = Reflexion.onBound(codec)
-          .findField("toId")
-          .flatMap(accessor -> accessor.<Map<Object, Integer>>getValue().asOptional())
-          .orElseThrow();
+        var typeToId = new HashMap<Object, Integer>();
+        var packetVisitorImplGenerator = new PacketVisitorImplGenerator(packetVisitorClass, typeToId);
+        var packetVisitorImpl = packetVisitorImplGenerator.generatePacketVisitor();
+        var ignored = listPacketsMethod.invokeWithArgs(packetVisitorImpl).getOrThrow();
         result.put(convertedFlow, typeToId);
       }
     }
@@ -300,6 +326,15 @@ private void doMapTypes(
     @NonNull Class<T> type,
     @NonNull MethodNode methodNode,
     @Nullable AbstractInsnNode after
+  ) {
+    return this.findNodeOfType(type, methodNode, after, ignored -> true);
+  }
+
+  private @Nullable <T extends AbstractInsnNode> T findNodeOfType(
+    @NonNull Class<T> type,
+    @NonNull MethodNode methodNode,
+    @Nullable AbstractInsnNode after,
+    @NonNull Predicate<T> filter
   ) {
     // either start from the top or from the next node of the insn we want to start after
     AbstractInsnNode start;
@@ -317,7 +352,10 @@ private void doMapTypes(
       }
 
       if (current.getClass() == type) {
-        return type.cast(current);
+        var node = type.cast(current);
+        if (filter.test(node)) {
+          return node;
+        }
       }
 
       current = current.getNext();

src/main/java/dev/derklaro/protocolgenerator/protocol/McUnboundProtocolBinder.java

@@ -0,0 +1,65 @@
+/*
+ * This file is part of mc-protocol-generator, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2023 - 2024 Pasqual K. and contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package dev.derklaro.protocolgenerator.protocol;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Proxy;
+import java.util.Map;
+import lombok.NonNull;
+
+final class PacketVisitorImplGenerator {
+
+  private final Class<?> packetVisitorClass;
+  private final Map<Object, Integer> targetPacketTypeToIdMap;
+
+  public PacketVisitorImplGenerator(
+    @NonNull Class<?> packetVisitorClass,
+    @NonNull Map<Object, Integer> targetPacketTypeToIdMap
+  ) {
+    this.packetVisitorClass = packetVisitorClass;
+    this.targetPacketTypeToIdMap = targetPacketTypeToIdMap;
+  }
+
+  public @NonNull Object generatePacketVisitor() {
+    InvocationHandler packetVisitorHandler = (ignored, method, args) -> {
+      // accept(PacketType, int)
+      if (method.getName().equals("accept")
+        && method.getParameterCount() == 2
+        && method.getReturnType() == void.class) {
+        var packetType = args[0];
+        var packetId = (Integer) args[1];
+        this.targetPacketTypeToIdMap.put(packetType, packetId);
+        return null;
+      }
+
+      // unknown method
+      throw new UnsupportedOperationException("Called unsupported method: " + method.getName() + " on proxy");
+    };
+    return Proxy.newProxyInstance(
+      this.packetVisitorClass.getClassLoader(),
+      new Class<?>[]{this.packetVisitorClass},
+      packetVisitorHandler);
+  }
+}

src/main/java/dev/derklaro/protocolgenerator/protocol/PacketVisitorImplGenerator.java

@@ -52,6 +52,7 @@ public ProtocolInfoCollector(@NonNull Path remappedJarPath, @NonNull ClassLoader
     // build a class loader which is aware of the remapped jar
     var remappedJarUrl = this.remappedJarPath.toUri().toURL();
     var loader = new URLClassLoader(new URL[]{remappedJarUrl}, this.libraryClassLoader);
+    var packetVisitorClass = Class.forName(McClassNames.PACKET_VISITOR, true, loader);
 
     // important first step: initialize minecraft registries
     var mcInitializer = new McInit(loader);
@@ -66,7 +67,7 @@ public ProtocolInfoCollector(@NonNull Path remappedJarPath, @NonNull ClassLoader
       var protocolsClass = Class.forName(state.protocolsClass(), true, loader);
       var infoDecoder = new McProtocolsDecoder(protocolsClass, loader);
       var typeToPacketClassPerFlow = infoDecoder.decodeAssociatedPackets();
-      var typeToPacketIdPerFlow = infoDecoder.resolvePacketIds();
+      var typeToPacketIdPerFlow = infoDecoder.resolvePacketIds(packetVisitorClass);
 
       // register for each flow (if present)
       for (var flow : McPacketFlow.values()) {