Fix entity tracking in 1.14

Fixes #600
2019-05-04 22:41:32 -04:00 · 2019-05-04 22:41:32 -04:00 · b14b4cc345
Commit b14b4cc345
--- a/src/main/java/com/comphenix/protocol/injector/DelayedPacketManager.java
+++ b/src/main/java/com/comphenix/protocol/injector/DelayedPacketManager.java
@ -5,28 +5,15 @@ import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
-
 import javax.annotation.Nonnull;

-import org.bukkit.Location;
-import org.bukkit.World;
-import org.bukkit.entity.Entity;
-import org.bukkit.entity.Player;
-import org.bukkit.plugin.Plugin;
-import org.bukkit.plugin.PluginManager;
-
 import com.comphenix.protocol.AsynchronousManager;
 import com.comphenix.protocol.PacketType;
 import com.comphenix.protocol.PacketType.Sender;
 import com.comphenix.protocol.error.ErrorReporter;
 import com.comphenix.protocol.error.Report;
 import com.comphenix.protocol.error.ReportType;
-import com.comphenix.protocol.events.ConnectionSide;
-import com.comphenix.protocol.events.ListeningWhitelist;
-import com.comphenix.protocol.events.NetworkMarker;
-import com.comphenix.protocol.events.PacketAdapter;
-import com.comphenix.protocol.events.PacketContainer;
-import com.comphenix.protocol.events.PacketListener;
+import com.comphenix.protocol.events.*;
 import com.comphenix.protocol.injector.netty.WirePacket;
 import com.comphenix.protocol.injector.packet.PacketRegistry;
 import com.comphenix.protocol.reflect.FieldAccessException;
@ -37,6 +24,13 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;

+import org.bukkit.Location;
+import org.bukkit.World;
+import org.bukkit.entity.Entity;
+import org.bukkit.entity.Player;
+import org.bukkit.plugin.Plugin;
+import org.bukkit.plugin.PluginManager;
+
 /**
 * A protocol manager that delays all packet listener registrations and unregistrations until
 * an underlying protocol manager can be constructed.
@ -440,7 +434,7 @@ public class DelayedPacketManager implements InternalManager {
 		if (delegate != null)
 			delegate.updateEntity(entity, observers);
 		else
-			EntityUtilities.updateEntity(entity, observers);
+			EntityUtilities.getInstance().updateEntity(entity, observers);
 	}
 	
 	@Override
@ -448,7 +442,7 @@ public class DelayedPacketManager implements InternalManager {
 		if (delegate != null)
 			return delegate.getEntityFromID(container, id);
 		else
-			return EntityUtilities.getEntityFromID(container, id);
+			return EntityUtilities.getInstance().getEntityFromID(container, id);
 	}
 	
 	@Override
@ -456,7 +450,7 @@ public class DelayedPacketManager implements InternalManager {
 		if (delegate != null)
 			return delegate.getEntityTrackers(entity);
 		else
-			return EntityUtilities.getEntityTrackers(entity);
+			return EntityUtilities.getInstance().getEntityTrackers(entity);
 	}

 	@Override
--- a/src/main/java/com/comphenix/protocol/injector/EntityUtilities.java
+++ b/src/main/java/com/comphenix/protocol/injector/EntityUtilities.java
@ -17,25 +17,27 @@

 package com.comphenix.protocol.injector;

-import java.lang.reflect.Field;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
+import com.comphenix.protocol.reflect.FieldAccessException;
+import com.comphenix.protocol.reflect.FuzzyReflection;
+import com.comphenix.protocol.reflect.accessors.Accessors;
+import com.comphenix.protocol.reflect.accessors.FieldAccessor;
+import com.comphenix.protocol.reflect.accessors.MethodAccessor;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyFieldContract;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyMethodContract;
+import com.comphenix.protocol.utility.MinecraftReflection;
+import com.comphenix.protocol.utility.MinecraftVersion;
+import com.comphenix.protocol.wrappers.WrappedIntHashMap;
+import com.google.common.collect.Lists;

 import org.apache.commons.lang.Validate;
 import org.bukkit.World;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;

-import com.comphenix.protocol.reflect.FieldAccessException;
-import com.comphenix.protocol.reflect.FieldUtils;
-import com.comphenix.protocol.reflect.FuzzyReflection;
-import com.comphenix.protocol.utility.MinecraftReflection;
-import com.comphenix.protocol.wrappers.WrappedIntHashMap;
-import com.google.common.collect.Lists;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;

 /**
 * Used to perform certain operations on entities.
@ -43,64 +45,41 @@ import com.google.common.collect.Lists;
 * @author Kristian
 */
 class EntityUtilities {
+	private static final boolean NEW_TRACKER = MinecraftVersion.atOrAbove(MinecraftVersion.VILLAGE_UPDATE);
+	private static final EntityUtilities INSTANCE = new EntityUtilities();
 	
-	private static Field entityTrackerField;
-	private static Field trackedEntitiesField;
-	private static Field trackedPlayersField;
-	private static Field trackerField;
-	
-	private static Method scanPlayersMethod;
-
-	/*
-	public static void updateEntity2(Entity entity, List<Player> observers) {
-		EntityTrackerEntry entry = getEntityTrackerEntry(entity.getWorld(), entity.getEntityId());
-
-		List<EntityPlayer> nmsPlayers = getNmsPlayers(observers);
-
-		entry.trackedPlayers.removeAll(nmsPlayers);
-		entry.scanPlayers(nmsPlayers);
+	public static EntityUtilities getInstance() {
+		return INSTANCE;
 	}
-	*/

-	public static void updateEntity(Entity entity, List<Player> observers) throws FieldAccessException {
+	private EntityUtilities() { }
+
+	private FieldAccessor entityTrackerField;
+	private FieldAccessor trackedEntitiesField;
+	private FieldAccessor trackedPlayersField;
+	private FieldAccessor trackerField;
+	
+	private MethodAccessor scanPlayersMethod;
+
+	public void updateEntity(Entity entity, List<Player> observers) {
 		if (entity == null || !entity.isValid()) {
 			return;
 		}

-		try {
-			Object trackerEntry = getEntityTrackerEntry(entity.getWorld(), entity.getEntityId());
-			if (trackerEntry == null) {
-				throw new IllegalArgumentException("Cannot find entity trackers for " + entity + ".");
-			}
-
-			if (trackedPlayersField == null) {
-				trackedPlayersField = FuzzyReflection.fromObject(trackerEntry).getFieldByType("java\\.util\\..*");
-			}
-			
-			// Phew, finally there.
-			Collection<?> trackedPlayers = getTrackedPlayers(trackedPlayersField, trackerEntry);
+		Collection<?> trackedPlayers = getTrackedPlayers(entity);
 		List<Object> nmsPlayers = unwrapBukkit(observers);

 		trackedPlayers.removeAll(nmsPlayers);

-			// We have to rely on a NAME once again. Damn it.
-			// TODO: Make sure this stays up to date with version changes - 1.8 - 1.10
+		// TODO Find equivalent method for newer versions
+
 		if (scanPlayersMethod == null) {
-				scanPlayersMethod = trackerEntry.getClass().getMethod("scanPlayers", List.class);
+			Class<?> trackerEntry = MinecraftReflection.getMinecraftClass("EntityTrackerEntry");
+			scanPlayersMethod = Accessors.getMethodAccessor(trackerEntry, "scanPlayers");
 		}

+		Object trackerEntry = getEntityTrackerEntry(entity.getWorld(), entity.getEntityId());
 		scanPlayersMethod.invoke(trackerEntry, nmsPlayers);
-		} catch (IllegalArgumentException e) {
-			throw e;
-		} catch (IllegalAccessException e) {
-			throw new FieldAccessException("Security limitation prevents access to 'get' method in IntHashMap", e);
-		} catch (InvocationTargetException e) {
-			throw new RuntimeException("Exception occurred in Minecraft.", e);
-		} catch (SecurityException e) {
-			throw new FieldAccessException("Security limitation prevents access to 'scanPlayers' method in trackerEntry.", e);
-		} catch (NoSuchMethodException e) {
-			throw new FieldAccessException("Cannot find 'scanPlayers' method. Is ProtocolLib up to date?", e);
-		}
 	}

 	/**
@ -109,24 +88,13 @@ class EntityUtilities {
 	 * @return Every client/player that is tracking the given entity.
 	 * @throws FieldAccessException If reflection failed.
 	 */
-	public static List<Player> getEntityTrackers(Entity entity) {
+	public List<Player> getEntityTrackers(Entity entity) {
 		if (entity == null || !entity.isValid()) {
 			return new ArrayList<>();
 		}

-		try {
-			List<Player> result = new ArrayList<Player>();
-
-			Object trackerEntry = getEntityTrackerEntry(entity.getWorld(), entity.getEntityId());
-			if (trackerEntry == null) {
-				throw new IllegalArgumentException("Cannot find entity trackers for " + entity + ".");
-			}
-
-			if (trackedPlayersField == null) {
-				trackedPlayersField = FuzzyReflection.fromObject(trackerEntry).getFieldByType("java\\.util\\..*");
-			}
-
-			Collection<?> trackedPlayers = getTrackedPlayers(trackedPlayersField, trackerEntry);
+		List<Player> result = new ArrayList<>();
+		Collection<?> trackedPlayers = getTrackedPlayers(entity);

 		// Wrap every player - we also ensure that the underlying tracker list is immutable
 		for (Object tracker : trackedPlayers) {
@ -136,18 +104,21 @@ class EntityUtilities {
 		}

 		return result;
-		} catch (IllegalAccessException e) {
-			throw new FieldAccessException("Security limitation prevented access to the list of tracked players.", e);
-		}
 	}

-	// Damn you, Paper
-	private static Collection<?> getTrackedPlayers(Field field, Object entry) throws IllegalAccessException {
-		Validate.notNull(field, "Cannot find 'trackedPlayers' field.");
-		Validate.notNull(entry, "entry cannot be null!");
+	private Collection<?> getTrackedPlayers(Entity entity) {
+		Validate.notNull(entity, "entity cannot be null");

-		Object value = FieldUtils.readField(field, entry, false);
+		Object trackerEntry = getEntityTrackerEntry(entity.getWorld(), entity.getEntityId());
+		Validate.notNull(trackerEntry, "Could not find entity trackers for " + entity);

+		if (trackedPlayersField == null) {
+			trackedPlayersField = Accessors.getFieldAccessor(FuzzyReflection.fromObject(trackerEntry).getFieldByType("java\\.util\\..*"));
+		}
+
+		Validate.notNull(trackedPlayersField, "Could not find trackedPlayers field");
+
+		Object value = trackedPlayersField.get(trackerEntry);
 		if (value instanceof Collection) {
 			return (Collection<?>) value;
 		} else if (value instanceof Map) {
@ -158,46 +129,62 @@ class EntityUtilities {
 		}
 	}
 	
-	/*
-	private static EntityTrackerEntry getEntityTrackerEntry2(World world, int entityID) {
-		WorldServer worldServer = ((CraftWorld) world).getHandle();
-		EntityTracker tracker = worldServer.tracker;
-		return tracker.trackedEntities.get(entityID);
-	}
-	*/
+	private MethodAccessor getChunkProvider;
+	private FieldAccessor chunkMapField;

-	private static Object getEntityTrackerEntry(World world, int entityID) throws FieldAccessException, IllegalArgumentException {
+	@SuppressWarnings("unchecked")
+	private Object getNewEntityTracker(Object worldServer, int entityId) {
+		if (getChunkProvider == null) {
+			Class<?> chunkProviderClass = MinecraftReflection.getMinecraftClass("ChunkProviderServer");
+			getChunkProvider = Accessors.getMethodAccessor(
+					FuzzyReflection.fromClass(worldServer.getClass(), false).getMethod(
+							FuzzyMethodContract.newBuilder().parameterCount(0).returnTypeExact(chunkProviderClass).build()));
+		}
+		
+		Object chunkProvider = getChunkProvider.invoke(worldServer);
+
+		if (chunkMapField == null) {
+			Class<?> chunkMapClass = MinecraftReflection.getMinecraftClass("PlayerChunkMap");
+			chunkMapField = Accessors.getFieldAccessor(
+					FuzzyReflection.fromClass(chunkProvider.getClass(), false).getField(
+							FuzzyFieldContract.newBuilder().typeExact(chunkMapClass).build()));
+		}
+		
+		Object playerChunkMap = chunkMapField.get(chunkProvider);
+
+		if (trackedEntitiesField == null) {
+			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);
+		return trackedEntities.get(entityId);
+	}
+
+	private Object getEntityTrackerEntry(World world, int entityID) {
 		BukkitUnwrapper unwrapper = new BukkitUnwrapper();
 		Object worldServer = unwrapper.unwrapItem(world);

+		if (NEW_TRACKER) {
+			return getNewEntityTracker(worldServer, entityID);
+		}
+
 		if (entityTrackerField == null)
-			entityTrackerField = FuzzyReflection.fromObject(worldServer).
-									getFieldByType("tracker", MinecraftReflection.getEntityTrackerClass());
+			entityTrackerField = Accessors.getFieldAccessor(FuzzyReflection.fromObject(worldServer).
+									getFieldByType("tracker", MinecraftReflection.getEntityTrackerClass()));
 		
 		// Get the tracker
-		Object tracker = null;
-		
-		try {
-			tracker = FieldUtils.readField(entityTrackerField, worldServer, false);
-		} catch (IllegalAccessException e) {
-			throw new FieldAccessException("Cannot access 'tracker' field due to security limitations.", e);
-		}
+		Object tracker = entityTrackerField.get(worldServer);

 		// Looking for an IntHashMap in the tracker entry
 		if (trackedEntitiesField == null) {
-			trackedEntitiesField = FuzzyReflection.fromObject(tracker, false)
-					.getFieldByType("trackedEntities", MinecraftReflection.getIntHashMapClass());
+			trackedEntitiesField = Accessors.getFieldAccessor(FuzzyReflection.fromObject(tracker, false)
+					.getFieldByType("trackedEntities", MinecraftReflection.getIntHashMapClass()));
 		}

 		// Read the map
-		Object trackedEntities = null;
-
-		try {
-			trackedEntities = FieldUtils.readField(trackedEntitiesField, tracker, false);
-		} catch (IllegalAccessException e) {
-			throw new FieldAccessException("Cannot access 'trackedEntities' field due to security limitations.", e);
-		}
-
+		Object trackedEntities = trackedEntitiesField.get(tracker);
 		return WrappedIntHashMap.fromHandle(trackedEntities).get(entityID);
 	}

@ -206,7 +193,10 @@ class EntityUtilities {
 	 * @return The associated entity.
 	 * @throws FieldAccessException Reflection error.
 	 */
-	public static Entity getEntityFromID(World world, int entityID) throws FieldAccessException {
+	public Entity getEntityFromID(World world, int entityID) {
+		Validate.notNull(world, "world cannot be null");
+		Validate.isTrue(entityID >= 0, "entityID cannot be negative");
+
 		try {
 			Object trackerEntry = getEntityTrackerEntry(world, entityID);
 			Object tracker = null;
@ -215,29 +205,25 @@ class EntityUtilities {
 			if (trackerEntry != null) {
 				if (trackerField == null) {
 					try {
-						Class<?> entryClass = MinecraftReflection.getMinecraftClass("EntityTrackerEntry");
-						trackerField = entryClass.getDeclaredField("tracker");
-					} catch (NoSuchFieldException e) {
+						trackerField = Accessors.getFieldAccessor(trackerEntry.getClass(), "tracker", true);
+					} catch (Exception e) {
 						// Assume it's the first entity field then
-						trackerField = FuzzyReflection.fromObject(trackerEntry, true)
-								.getFieldByType("tracker", MinecraftReflection.getEntityClass());
+						trackerField = Accessors.getFieldAccessor(FuzzyReflection.fromObject(trackerEntry, true)
+								.getFieldByType("tracker", MinecraftReflection.getEntityClass()));
 					}
 				}

-				tracker = FieldUtils.readField(trackerField, trackerEntry, true);
+				tracker = trackerField.get(trackerEntry);
 			}

 			// If the tracker is NULL, we'll just assume this entity doesn't exist
-			if (tracker != null)
-				return (Entity) MinecraftReflection.getBukkitEntity(tracker);
-			else
-				return null;
+			return tracker != null ? (Entity) MinecraftReflection.getBukkitEntity(tracker) : null;
 		} catch (Exception e) {
 			throw new FieldAccessException("Cannot find entity from ID " + entityID + ".", e);
 		}
 	}

-	private static List<Object> unwrapBukkit(List<Player> players) {
+	private List<Object> unwrapBukkit(List<Player> players) {
 		List<Object> output = Lists.newArrayList();
 		BukkitUnwrapper unwrapper = new BukkitUnwrapper();
 		
--- a/src/main/java/com/comphenix/protocol/injector/PacketFilterManager.java
+++ b/src/main/java/com/comphenix/protocol/injector/PacketFilterManager.java
@ -56,8 +56,6 @@ import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;

-import io.netty.channel.Channel;
-
 import org.bukkit.Bukkit;
 import org.bukkit.Location;
 import org.bukkit.Server;
@ -74,6 +72,8 @@ import org.bukkit.event.server.PluginDisableEvent;
 import org.bukkit.plugin.Plugin;
 import org.bukkit.plugin.PluginManager;

+import io.netty.channel.Channel;
+
 public final class PacketFilterManager implements ListenerInvoker, InternalManager {

 	public static final ReportType REPORT_CANNOT_LOAD_PACKET_LIST = new ReportType("Cannot load server and client packet list.");
@ -924,17 +924,17 @@ public final class PacketFilterManager implements ListenerInvoker, InternalManag

 	@Override
 	public void updateEntity(Entity entity, List<Player> observers) throws FieldAccessException {
-		EntityUtilities.updateEntity(entity, observers);
+		EntityUtilities.getInstance().updateEntity(entity, observers);
 	}

 	@Override
 	public Entity getEntityFromID(World container, int id) throws FieldAccessException {
-		return EntityUtilities.getEntityFromID(container, id);
+		return EntityUtilities.getInstance().getEntityFromID(container, id);
 	}

 	@Override
 	public List<Player> getEntityTrackers(Entity entity) throws FieldAccessException {
-		return EntityUtilities.getEntityTrackers(entity);
+		return EntityUtilities.getInstance().getEntityTrackers(entity);
 	}

 	/**
--- a/src/main/java/com/comphenix/protocol/utility/MinecraftReflection.java
+++ b/src/main/java/com/comphenix/protocol/utility/MinecraftReflection.java
@ -1527,7 +1527,7 @@ public class MinecraftReflection {
 	 */
 	public static Class<?> getEntityTrackerClass() {
 		try {
-			return getMinecraftClass("EntityTracker");
+			return getMinecraftClass("EntityTracker", "PlayerChunkMap$EntityTracker");
 		} catch (RuntimeException e) {
 			FuzzyClassContract entityTrackerContract = FuzzyClassContract.newBuilder().
 					field(FuzzyFieldContract.newBuilder().
--- a/src/test/java/com/comphenix/protocol/injector/EntityUtilitiesTest.java
+++ b/src/test/java/com/comphenix/protocol/injector/EntityUtilitiesTest.java
@ -0,0 +1,55 @@
+package com.comphenix.protocol.injector;
+
+import com.comphenix.protocol.BukkitInitialization;
+import com.comphenix.protocol.reflect.accessors.Accessors;
+
+import net.minecraft.server.v1_14_R1.ChunkProviderServer;
+import net.minecraft.server.v1_14_R1.Entity;
+import net.minecraft.server.v1_14_R1.PlayerChunkMap;
+import net.minecraft.server.v1_14_R1.PlayerChunkMap.EntityTracker;
+import net.minecraft.server.v1_14_R1.WorldServer;
+
+import org.bukkit.craftbukkit.libs.it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import org.bukkit.craftbukkit.libs.it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import org.bukkit.craftbukkit.v1_14_R1.CraftWorld;
+import org.bukkit.craftbukkit.v1_14_R1.entity.CraftEntity;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class EntityUtilitiesTest {
+
+	@BeforeClass
+	public static void beforeClass() {
+		BukkitInitialization.initializeItemMeta();
+	}
+
+	@Test
+	public void testReflection() {
+		CraftWorld bukkit = mock(CraftWorld.class);
+		WorldServer world = mock(WorldServer.class);
+		when(bukkit.getHandle()).thenReturn(world);
+
+		ChunkProviderServer provider = mock(ChunkProviderServer.class);
+		when(world.getChunkProvider()).thenReturn(provider);
+
+		PlayerChunkMap chunkMap = mock(PlayerChunkMap.class);
+		Accessors.getFieldAccessor(ChunkProviderServer.class, "playerChunkMap", true).set(provider, chunkMap);
+
+		CraftEntity bukkitEntity = mock(CraftEntity.class);
+		Entity fakeEntity = mock(Entity.class);
+		when(fakeEntity.getBukkitEntity()).thenReturn(bukkitEntity);
+
+		EntityTracker tracker = mock(EntityTracker.class);
+		Accessors.getFieldAccessor(EntityTracker.class, "tracker", true).set(tracker, fakeEntity);
+
+		Int2ObjectMap<EntityTracker> trackerMap = new Int2ObjectOpenHashMap<>();
+		trackerMap.put(1, tracker);
+		Accessors.getFieldAccessor(PlayerChunkMap.class, "trackedEntities", true).set(chunkMap, trackerMap);
+
+		assertEquals(bukkitEntity, EntityUtilities.getInstance().getEntityFromID(bukkit, 1));
+	}
+}