Fix entity trackers in 1.17

Addresses #1217

Author: dmulloy2

src/main/java/com/comphenix/protocol/injector/EntityUtilities.java

@@ -17,6 +17,7 @@
 
 package com.comphenix.protocol.injector;
 
+import java.lang.reflect.Modifier;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -157,9 +158,21 @@ private Object getNewEntityTracker(Object worldServer, int entityId) {
 		Object playerChunkMap = chunkMapField.get(chunkProvider);
 
 		if (trackedEntitiesField == null) {
-			trackedEntitiesField = Accessors.getFieldAccessor(
+			if (MinecraftVersion.CAVES_CLIFFS_1.atOrAbove()) {
+				trackedEntitiesField = Accessors.getFieldAccessor(
+					FuzzyReflection.fromClass(playerChunkMap.getClass(), true).getField(
+						FuzzyFieldContract.newBuilder()
+							.banModifier(Modifier.STATIC)
+							.requirePublic()
+							.typeExact(org.bukkit.craftbukkit.libs.it.unimi.dsi.fastutil.ints.Int2ObjectMap.class)
+							.build()
+					)
+				);
+			} else {
+				trackedEntitiesField = Accessors.getFieldAccessor(
 					FuzzyReflection.fromClass(playerChunkMap.getClass(), false).getField(
 							FuzzyFieldContract.newBuilder().typeDerivedOf(Map.class).nameExact("trackedEntities").build()));
+			}
 		}
 
 		Map<Integer, Object> trackedEntities = (Map<Integer, Object>) trackedEntitiesField.get(playerChunkMap);
@@ -229,14 +242,17 @@ public Entity getEntityFromID(World world, int entityID) {
 			Object tracker = null;
 
 			if (trackerEntry != null) {
-				// plugins like citizens will use their own tracker
+				// plugins like citizens will use their own tracker class, so cache the result
 				FieldAccessor trackerField = trackerFields.computeIfAbsent(trackerEntry.getClass(), x -> {
+					// get the first entity field
 					try {
-						return Accessors.getFieldAccessor(trackerEntry.getClass(), "tracker", true);
-					} catch (Exception e) {
-						// Assume it's the first entity field then
-						return Accessors.getFieldAccessor(FuzzyReflection.fromObject(trackerEntry, true)
-								.getFieldByType("tracker", MinecraftReflection.getEntityClass()));
+						return Accessors.getFieldAccessor(FuzzyReflection.fromClass(trackerEntry.getClass(), true)
+								.getField(FuzzyFieldContract.newBuilder().typeExact(MinecraftReflection.getEntityClass()).build()));
+					} catch (Exception ex) {
+						// try with the default class
+						Class<?> trackerEntryClass = MinecraftReflection.getEntityTrackerClass();
+						return Accessors.getFieldAccessor(FuzzyReflection.fromClass(trackerEntryClass, true)
+								.getField(FuzzyFieldContract.newBuilder().typeExact(MinecraftReflection.getEntityClass()).build()));
 					}
 				});
 

src/test/java/com/comphenix/protocol/injector/EntityUtilitiesTest.java

@@ -7,6 +7,7 @@
 import com.comphenix.protocol.reflect.StructureModifier;
 import com.comphenix.protocol.reflect.accessors.Accessors;
 
+import com.comphenix.protocol.reflect.accessors.FieldAccessor;
 import com.comphenix.protocol.reflect.fuzzy.FuzzyFieldContract;
 import net.minecraft.server.level.ChunkProviderServer;
 import net.minecraft.server.level.EntityTrackerEntry;
@@ -24,6 +25,8 @@
 
 import java.lang.reflect.Field;
 
+import static com.comphenix.protocol.utility.TestUtils.setFinalField;
+
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -35,7 +38,7 @@ public static void beforeClass() {
 		BukkitInitialization.initializeItemMeta();
 	}
 
-	// @Test
+	@Test
 	public void testReflection() throws ReflectiveOperationException {
 		CraftWorld bukkit = mock(CraftWorld.class);
 		WorldServer world = mock(WorldServer.class);
@@ -44,29 +47,25 @@ public void testReflection() throws ReflectiveOperationException {
 		ChunkProviderServer provider = mock(ChunkProviderServer.class);
 		when(world.getChunkProvider()).thenReturn(provider);
 
-		// TODO unsetting final doesn't work anymore
 		PlayerChunkMap chunkMap = mock(PlayerChunkMap.class);
 		Field chunkMapField = FuzzyReflection.fromClass(ChunkProviderServer.class, true)
 				.getField(FuzzyFieldContract.newBuilder().typeExact(PlayerChunkMap.class).build());
-		chunkMapField.setAccessible(true);
-		chunkMapField.set(provider, chunkMap);
+		setFinalField(provider, chunkMapField, chunkMap);
 
 		CraftEntity bukkitEntity = mock(CraftEntity.class);
 		Entity fakeEntity = mock(Entity.class);
 		when(fakeEntity.getBukkitEntity()).thenReturn(bukkitEntity);
 
-		PlayerChunkMap.EntityTracker tracker = mock(PlayerChunkMap.EntityTracker.class);
-		FuzzyReflection.fromClass(EntityTracker.class, true)
-				.getField(FuzzyFieldContract.newBuilder().typeExact(EntityTrackerEntry.class).build())
-				.set(tracker, fakeEntity);
+		EntityTracker tracker = mock(EntityTracker.class);
+		Field trackerField = FuzzyReflection.fromClass(EntityTracker.class, true)
+				.getField(FuzzyFieldContract.newBuilder().typeExact(Entity.class).build());
+		setFinalField(tracker, trackerField, fakeEntity);
 
-		Int2ObjectMap<PlayerChunkMap.EntityTracker> trackerMap = new Int2ObjectOpenHashMap<>();
+		Int2ObjectMap<EntityTracker> trackerMap = new Int2ObjectOpenHashMap<>();
 		trackerMap.put(1, tracker);
-
-		new StructureModifier<>(PlayerChunkMap.class, true)
-				.withTarget(chunkMap)
-				.withParamType(Int2ObjectMap.class, null, EntityTracker.class)
-				.write(0, trackerMap);
+		Field trackedEntitiesField = FuzzyReflection.fromClass(PlayerChunkMap.class, true)
+				.getField(FuzzyFieldContract.newBuilder().typeExact(Int2ObjectMap.class).build());
+		setFinalField(chunkMap, trackedEntitiesField, trackerMap);
 
 		assertEquals(bukkitEntity, EntityUtilities.getInstance().getEntityFromID(bukkit, 1));
 	}

src/test/java/com/comphenix/protocol/utility/TestUtils.java

@@ -1,7 +1,9 @@
 package com.comphenix.protocol.utility;
 
+import java.lang.reflect.Field;
 import java.util.List;
 
+import sun.misc.Unsafe;
 import org.bukkit.Bukkit;
 import org.bukkit.inventory.ItemStack;
 
@@ -42,4 +44,13 @@ public static boolean equivalentItem(ItemStack first, ItemStack second) {
 			return first.getType().equals(second.getType());
 		}
 	}
+
+	public static void setFinalField(Object obj, Field field, Object newValue) throws ReflectiveOperationException {
+		Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe");
+		unsafeField.setAccessible(true);
+		Unsafe unsafe = (Unsafe) unsafeField.get(null);
+
+		long offset = unsafe.objectFieldOffset(field);
+		unsafe.putObject(obj, offset, newValue);
+	}
 }