From d3d047b5b1da1eca8d773d6f26476bdf716342c3 Mon Sep 17 00:00:00 2001
From: David Berdik <berdikgit@theprogrammersworld.net>
Date: Sat, 19 Nov 2016 16:32:23 -0500
Subject: [PATCH] *Fixed issue with custom entities not working properly with
 1.11

---
 .../nms/entity/CustomEntityRegistry.java      | 97 +++++++++++++++++++
 .../herobrine/nms/entity/EntityInjector.java  | 24 +----
 2 files changed, 100 insertions(+), 21 deletions(-)
 create mode 100644 src/net/theprogrammersworld/herobrine/nms/entity/CustomEntityRegistry.java

diff --git a/src/net/theprogrammersworld/herobrine/nms/entity/CustomEntityRegistry.java b/src/net/theprogrammersworld/herobrine/nms/entity/CustomEntityRegistry.java
new file mode 100644
index 0000000..fe94c08
--- /dev/null
+++ b/src/net/theprogrammersworld/herobrine/nms/entity/CustomEntityRegistry.java
@@ -0,0 +1,97 @@
+// This class was originally written by Scyntrus - https://github.com/Scyntrus/misc/blob/master/CustomEntityRegistry.java
+
+package net.theprogrammersworld.herobrine.nms.entity;
+
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
+import java.util.HashMap;
+import java.util.Map;
+import net.minecraft.server.v1_11_R1.RegistryMaterials;
+import net.minecraft.server.v1_11_R1.Entity;
+import net.minecraft.server.v1_11_R1.EntityTypes;
+import net.minecraft.server.v1_11_R1.MinecraftKey;
+
+@SuppressWarnings("rawtypes")
+public class CustomEntityRegistry extends RegistryMaterials {
+
+    private static CustomEntityRegistry instance = null;
+
+    private final BiMap<MinecraftKey, Class<? extends Entity>> customEntities = HashBiMap.create();
+    private final BiMap<Class<? extends Entity>, MinecraftKey> customEntityClasses = this.customEntities.inverse();
+    private final Map<Class<? extends Entity>, Integer> customEntityIds = new HashMap<>();
+
+    private final RegistryMaterials wrapped;
+
+    private CustomEntityRegistry(RegistryMaterials original) {
+        this.wrapped = original;
+    }
+
+    public static CustomEntityRegistry getInstance() {
+        if (instance != null) {
+            return instance;
+        }
+
+        instance = new CustomEntityRegistry(EntityTypes.b);
+
+        try {
+            //TODO: Update name on version change (RegistryMaterials)
+            Field registryMaterialsField = EntityTypes.class.getDeclaredField("b");
+            registryMaterialsField.setAccessible(true);
+
+            Field modifiersField = Field.class.getDeclaredField("modifiers");
+            modifiersField.setAccessible(true);
+            modifiersField.setInt(registryMaterialsField, registryMaterialsField.getModifiers() & ~Modifier.FINAL);
+
+            registryMaterialsField.set(null, instance);
+        } catch (Exception e) {
+            instance = null;
+
+            throw new RuntimeException("Unable to override the old entity RegistryMaterials", e);
+        }
+
+        return instance;
+    }
+
+    public static void addCustomEntity(int entityId, String entityName, Class<? extends Entity> entityClass) {
+        getInstance().putCustomEntity(entityId, entityName, entityClass);
+    }
+
+    public void putCustomEntity(int entityId, String entityName, Class<? extends Entity> entityClass) {
+        MinecraftKey minecraftKey = new MinecraftKey(entityName);
+
+        this.customEntities.put(minecraftKey, entityClass);
+        this.customEntityIds.put(entityClass, entityId);
+    }
+
+    @SuppressWarnings("unchecked")
+	@Override
+    public Class<? extends Entity> get(Object key) {
+        if (this.customEntities.containsKey(key)) {
+            return this.customEntities.get(key);
+        }
+
+        return (Class<? extends Entity>) wrapped.get(key);
+    }
+
+    @SuppressWarnings("unchecked")
+	@Override
+    public int a(Object key) { //TODO: Update name on version change (getId)
+        if (this.customEntityIds.containsKey(key)) {
+            return this.customEntityIds.get(key);
+        }
+
+        return this.wrapped.a(key);
+    }
+
+    @SuppressWarnings("unchecked")
+	@Override
+    public MinecraftKey b(Object value) { //TODO: Update name on version change (getKey)
+        if (this.customEntityClasses.containsKey(value)) {
+            return this.customEntityClasses.get(value);
+        }
+
+        return (MinecraftKey) wrapped.b(value);
+    }
+}
\ No newline at end of file
diff --git a/src/net/theprogrammersworld/herobrine/nms/entity/EntityInjector.java b/src/net/theprogrammersworld/herobrine/nms/entity/EntityInjector.java
index 2989f47..6b71ab0 100644
--- a/src/net/theprogrammersworld/herobrine/nms/entity/EntityInjector.java
+++ b/src/net/theprogrammersworld/herobrine/nms/entity/EntityInjector.java
@@ -1,31 +1,13 @@
 package net.theprogrammersworld.herobrine.nms.entity;
 
-import java.lang.reflect.Field;
-import java.lang.reflect.Method;
-import java.util.Map;
-
-import net.minecraft.server.v1_11_R1.EntityTypes;
 import net.theprogrammersworld.herobrine.Herobrine;
 
 public class EntityInjector {
 
-	@SuppressWarnings("unchecked")
 	public static void inject() {
-		try {
-			final Field idMapField = EntityTypes.class.getDeclaredField("e");
-			idMapField.setAccessible(true);
-			Map<Integer, Class<?>> idMap = (Map<Integer, Class<?>>) idMapField.get(null);
-			final Field nameMapField = EntityTypes.class.getDeclaredField("c");
-			nameMapField.setAccessible(true);
-			Map<String, Class<?>> nameMap = (Map<String, Class<?>>) nameMapField.get(null);
-			nameMap.remove("Zombie");
-			nameMap.remove("Skeleton");
-			idMap.remove(54);
-			idMap.remove(51);
-			final Method a = EntityTypes.class.getDeclaredMethod("a", Class.class, String.class, Integer.TYPE);
-			a.setAccessible(true);
-			a.invoke(null, CustomZombie.class, "Zombie", 54);
-			a.invoke(null, CustomSkeleton.class, "Skeleton", 51);
+		try {			
+			CustomEntityRegistry.addCustomEntity(51, "Skeleton", CustomSkeleton.class);
+			CustomEntityRegistry.addCustomEntity(54, "Zombie", CustomZombie.class);
 		} catch (Throwable t) {
 			t.printStackTrace();
 			Herobrine.isNPCDisabled = true;