1.1.16: Override registered packets

Author: JNNGL

README.md

@@ -53,7 +53,7 @@ Test server: [``ely.su``](https://hotmc.ru/minecraft-server-203216)
         <dependency>
             <groupId>net.elytrium.limboapi</groupId>
             <artifactId>api</artifactId>
-            <version>1.1.15</version>
+            <version>1.1.16</version>
             <scope>provided</scope>
         </dependency>
     </dependencies>
@@ -70,7 +70,7 @@ Test server: [``ely.su``](https://hotmc.ru/minecraft-server-203216)
     }
 
     dependencies {
-        compileOnly("net.elytrium.limboapi:api:1.1.15")
+        compileOnly("net.elytrium.limboapi:api:1.1.16")
     }
 ```
 

VERSION

@@ -1 +1 @@
-1.1.15
+1.1.16

api/src/main/java/net/elytrium/limboapi/api/utils/OverlayMap.java

@@ -13,6 +13,7 @@
 
 public abstract class OverlayMap<K, V> implements Map<K, V> {
 
+  protected boolean override = false;
   protected final Map<K, V> parent;
   protected final Map<K, V> overlay;
 
@@ -33,11 +34,19 @@ public boolean isEmpty() {
 
   @Override
   public boolean containsKey(Object o) {
+    if (this.override) {
+      return this.overlay.containsKey(o);
+    }
+
     return this.parent.containsKey(o) || this.overlay.containsKey(o);
   }
 
   @Override
   public boolean containsValue(Object o) {
+    if (this.override) {
+      return this.overlay.containsValue(o);
+    }
+
     return this.parent.containsValue(o) || this.overlay.containsValue(o);
   }
 
@@ -71,4 +80,12 @@ public void putAll(@NotNull Map<? extends K, ? extends V> map) {
   public void clear() {
     this.overlay.clear();
   }
+
+  public boolean isOverride() {
+    return this.override;
+  }
+
+  public void setOverride(boolean override) {
+    this.override = override;
+  }
 }

build.gradle

@@ -13,7 +13,7 @@ allprojects {
   apply(plugin: "org.cadixdev.licenser")
 
   setGroup("net.elytrium.limboapi")
-  setVersion("1.1.15")
+  setVersion("1.1.16")
 
   compileJava {
     getOptions().setEncoding("UTF-8")

plugin/src/main/java/net/elytrium/limboapi/protocol/LimboProtocol.java

@@ -31,11 +31,15 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumMap;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import net.elytrium.commons.utils.reflection.ReflectionException;
 import net.elytrium.limboapi.api.protocol.PacketDirection;
 import net.elytrium.limboapi.api.protocol.packets.PacketMapping;
+import net.elytrium.limboapi.api.utils.OverlayMap;
 import net.elytrium.limboapi.protocol.packets.c2s.MoveOnGroundOnlyPacket;
 import net.elytrium.limboapi.protocol.packets.c2s.MovePacket;
 import net.elytrium.limboapi.protocol.packets.c2s.MovePositionOnlyPacket;
@@ -67,8 +71,11 @@ public class LimboProtocol {
 
   public static final String READ_TIMEOUT = "limboapi-read-timeout";
 
+  public static final MethodHandle VERSIONS_GETTER;
   public static final Field VERSIONS_FIELD;
+  public static final MethodHandle PACKET_ID_TO_SUPPLIER_GETTER;
   public static final Field PACKET_ID_TO_SUPPLIER_FIELD;
+  public static final MethodHandle PACKET_CLASS_TO_ID_GETTER;
   public static final Field PACKET_CLASS_TO_ID_FIELD;
   public static final StateRegistry.PacketRegistry PLAY_CLIENTBOUND_REGISTRY;
   public static final StateRegistry.PacketRegistry PLAY_SERVERBOUND_REGISTRY;
@@ -85,12 +92,18 @@ public class LimboProtocol {
 
       LIMBO_STATE_REGISTRY = (StateRegistry) UNSAFE.allocateInstance(StateRegistry.class);
 
+      VERSIONS_GETTER = MethodHandles.privateLookupIn(StateRegistry.PacketRegistry.class, MethodHandles.lookup())
+          .findGetter(StateRegistry.PacketRegistry.class, "versions", Map.class);
       VERSIONS_FIELD = StateRegistry.PacketRegistry.class.getDeclaredField("versions");
       VERSIONS_FIELD.setAccessible(true);
 
+      PACKET_ID_TO_SUPPLIER_GETTER = MethodHandles.privateLookupIn(StateRegistry.PacketRegistry.ProtocolRegistry.class, MethodHandles.lookup())
+          .findGetter(StateRegistry.PacketRegistry.ProtocolRegistry.class, "packetIdToSupplier", IntObjectMap.class);
       PACKET_ID_TO_SUPPLIER_FIELD = StateRegistry.PacketRegistry.ProtocolRegistry.class.getDeclaredField("packetIdToSupplier");
       PACKET_ID_TO_SUPPLIER_FIELD.setAccessible(true);
 
+      PACKET_CLASS_TO_ID_GETTER = MethodHandles.privateLookupIn(StateRegistry.PacketRegistry.ProtocolRegistry.class, MethodHandles.lookup())
+          .findGetter(StateRegistry.PacketRegistry.ProtocolRegistry.class, "packetClassToId", Object2IntMap.class);
       PACKET_CLASS_TO_ID_FIELD = StateRegistry.PacketRegistry.ProtocolRegistry.class.getDeclaredField("packetClassToId");
       PACKET_CLASS_TO_ID_FIELD.setAccessible(true);
 
@@ -121,7 +134,7 @@ public class LimboProtocol {
   }
 
   private static void overlayRegistry(StateRegistry stateRegistry,
-                                      String registryName, StateRegistry.PacketRegistry playRegistry) throws ReflectiveOperationException {
+                                      String registryName, StateRegistry.PacketRegistry playRegistry) throws Throwable {
     StateRegistry.PacketRegistry registry = (StateRegistry.PacketRegistry) UNSAFE.allocateInstance(StateRegistry.PacketRegistry.class);
 
     Field directionField = StateRegistry.PacketRegistry.class.getDeclaredField("direction");
@@ -133,7 +146,8 @@ private static void overlayRegistry(StateRegistry stateRegistry,
 
     // Overlay packets from PLAY state registry.
     // P.S. I hate it when someone uses var in code, but there I had no choice.
-    var playProtocolRegistryVersions = (Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry>) VERSIONS_FIELD.get(playRegistry);
+    var playProtocolRegistryVersions =
+        (Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry>) VERSIONS_GETTER.invokeExact(playRegistry);
     Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry> versions = new EnumMap<>(ProtocolVersion.class);
     for (ProtocolVersion version : ProtocolVersion.values()) {
       if (!version.isLegacy() && !version.isUnknown()) {
@@ -142,10 +156,10 @@ private static void overlayRegistry(StateRegistry stateRegistry,
 
         versionField.set(protoRegistry, version);
 
-        var playPacketIDToSupplier = (IntObjectMap<Supplier<? extends MinecraftPacket>>) PACKET_ID_TO_SUPPLIER_FIELD.get(playProtoRegistry);
+        var playPacketIDToSupplier = (IntObjectMap<Supplier<? extends MinecraftPacket>>) PACKET_ID_TO_SUPPLIER_GETTER.invokeExact(playProtoRegistry);
         PACKET_ID_TO_SUPPLIER_FIELD.set(protoRegistry, new OverlayIntObjectMap<>(playPacketIDToSupplier, new IntObjectHashMap<>(16, 0.5F)));
 
-        var playPacketClassToID = (Object2IntMap<Class<? extends MinecraftPacket>>) PACKET_CLASS_TO_ID_FIELD.get(playProtoRegistry);
+        var playPacketClassToID = (Object2IntMap<Class<? extends MinecraftPacket>>) PACKET_CLASS_TO_ID_GETTER.invokeExact(playProtoRegistry);
         Object2IntMap<Class<? extends MinecraftPacket>> packetClassToID = new Object2IntOpenHashMap<>(16, 0.5F);
         packetClassToID.defaultReturnValue(playPacketClassToID.defaultReturnValue());
         PACKET_CLASS_TO_ID_FIELD.set(protoRegistry, new OverlayObject2IntMap<>(playPacketClassToID, packetClassToID));
@@ -448,7 +462,27 @@ public static void register(StateRegistry stateRegistry,
   public static void register(StateRegistry.PacketRegistry registry,
                               Class<?> packetClass, Supplier<?> packetSupplier, StateRegistry.PacketMapping... mappings) {
     try {
+      var versions = (Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry>) VERSIONS_GETTER.invokeExact(registry);
+      List<OverlayMap<?, ?>> overlayMaps = versions.values().stream().flatMap(protocolRegistry -> {
+        try {
+          var idToSupplier = (IntObjectMap<Supplier<? extends MinecraftPacket>>) PACKET_ID_TO_SUPPLIER_GETTER.invokeExact(protocolRegistry);
+          var classToId = (Object2IntMap<Class<? extends MinecraftPacket>>) PACKET_CLASS_TO_ID_GETTER.invokeExact(protocolRegistry);
+          if (idToSupplier instanceof OverlayMap<?, ?> && classToId instanceof OverlayMap<?, ?>) {
+            return Stream.of(
+                (OverlayMap<?, ?>) idToSupplier,
+                (OverlayMap<?, ?>) classToId
+            );
+          } else {
+            return Stream.empty();
+          }
+        } catch (Throwable e) {
+          throw new ReflectionException(e);
+        }
+      }).collect(Collectors.toList());
+
+      overlayMaps.forEach(overlayMap -> overlayMap.setOverride(true));
       REGISTER_METHOD.invokeExact(registry, packetClass, packetSupplier, mappings);
+      overlayMaps.forEach(overlayMap -> overlayMap.setOverride(false));
     } catch (Throwable e) {
       throw new ReflectionException(e);
     }