Refresh the Netty registry when packet listeners are added. FIXES #40

ProtocolLib/src/main/java/com/comphenix/protocol/injector/netty/NettyProtocolRegistry.java

@@ -9,10 +9,12 @@ import java.util.Set;
 import com.comphenix.protocol.PacketType;
 import com.comphenix.protocol.PacketType.Protocol;
 import com.comphenix.protocol.PacketType.Sender;
+import com.comphenix.protocol.injector.packet.MapContainer;
 import com.comphenix.protocol.reflect.StructureModifier;
 import com.comphenix.protocol.utility.MinecraftReflection;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 
@@ -22,12 +24,35 @@ import com.google.common.collect.Sets;
  */
 // TODO: Handle modifications to the BiMap
 public class NettyProtocolRegistry {
+	/**
+	 * Represents a register we are currently building.
+	 * @author Kristian
+	 */
+	private static class Register {
+		// The main lookup table
+		public BiMap<PacketType, Class<?>> typeToClass = HashBiMap.create();
+		public volatile Set<PacketType> serverPackets = Sets.newHashSet();
+		public volatile Set<PacketType> clientPackets = Sets.newHashSet();
+		public List<MapContainer> containers = Lists.newArrayList();
+		
+		/**
+		 * Determine if the current register is outdated.
+		 * @return TRUE if it is, FALSE otherwise.
+		 */
+		public boolean isOutdated() {
+			for (MapContainer container : containers) {
+				if (container.hasChanged()) {
+					return true;
+				}
+			}
+			return false;
+		}
+	}
+	
 	private Class<?> enumProtocol;
 	
-	// The main lookup table
+	// Current register
-	private BiMap<PacketType, Class<?>> typeToClass = HashBiMap.create();
+	private volatile Register register;
-	private Set<PacketType> serverPackets = Sets.newHashSet();
-	private Set<PacketType> clientPackets = Sets.newHashSet();
 	
 	public NettyProtocolRegistry() {
 		enumProtocol = MinecraftReflection.getEnumProtocolClass();
@@ -39,7 +64,7 @@ public class NettyProtocolRegistry {
 	 * @return The packet type lookup.
 	 */
 	public Map<PacketType, Class<?>> getPacketTypeLookup() {
-		return Collections.unmodifiableMap(typeToClass);
+		return Collections.unmodifiableMap(register.typeToClass);
 	}
 	
 	/**
@@ -47,7 +72,7 @@ public class NettyProtocolRegistry {
 	 * @return The packet type lookup.
 	 */
 	public Map<Class<?>, PacketType> getPacketClassLookup() {
-		return Collections.unmodifiableMap(typeToClass.inverse());
+		return Collections.unmodifiableMap(register.typeToClass.inverse());
 	}
 	
 	/**
@@ -55,7 +80,7 @@ public class NettyProtocolRegistry {
 	 * @return Every client packet.
 	 */
 	public Set<PacketType> getClientPackets() {
-		return Collections.unmodifiableSet(clientPackets);
+		return Collections.unmodifiableSet(register.clientPackets);
 	}
 	
 	/**
@@ -63,33 +88,52 @@ public class NettyProtocolRegistry {
 	 * @return Every server packet.
 	 */
 	public Set<PacketType> getServerPackets() {
-		return Collections.unmodifiableSet(serverPackets);
+		return Collections.unmodifiableSet(register.serverPackets);
+	}
+	
+	/**
+	 * Ensure that our local register is up-to-date with Minecraft.
+	 * <p>
+	 * This operation may block the calling thread.
+	 */
+	public synchronized void synchronize() {
+		// See if the register is outdated
+		if (register.isOutdated()) {
+			initialize();
+		}
 	}
 	
 	/**
 	 * Load the packet lookup tables in each protocol.
 	 */
-	private void initialize() {
+	private synchronized void initialize() {
 		final Object[] protocols = enumProtocol.getEnumConstants();
-		List<Map<Integer, Class<?>>> serverPackets = Lists.newArrayList();
+		List<Map<Integer, Class<?>>> serverMaps = Lists.newArrayList();
-		List<Map<Integer, Class<?>>> clientPackets = Lists.newArrayList();
+		List<Map<Integer, Class<?>>> clientMaps = Lists.newArrayList();
 		StructureModifier<Object> modifier = null;
 		
+		// Result
+		Register result = new Register();
+		
 		for (Object protocol : protocols) {
 			if (modifier == null)
 				modifier = new StructureModifier<Object>(protocol.getClass().getSuperclass(), false);
 			StructureModifier<Map<Integer, Class<?>>> maps = modifier.withTarget(protocol).withType(Map.class);
 			
-			serverPackets.add(maps.read(0));
+			serverMaps.add(maps.read(0));
-			clientPackets.add(maps.read(1));
+			clientMaps.add(maps.read(1));
 		}
-		
+		// Maps we have to occationally check have changed
+		for (Map<Integer, Class<?>> map : Iterables.concat(serverMaps, clientMaps)) {
+			result.containers.add(new MapContainer(map));
+		}
+ 		
 		// Heuristic - there are more server packets than client packets
-		if (sum(clientPackets) > sum(serverPackets)) {
+		if (sum(clientMaps) > sum(serverMaps)) {
 			// Swap if this is violated
-			List<Map<Integer, Class<?>>> temp = serverPackets;
+			List<Map<Integer, Class<?>>> temp = serverMaps;
-			serverPackets = clientPackets;
+			serverMaps = clientMaps;
-			clientPackets = temp;
+			clientMaps = temp;
 		}
 		
 		for (int i = 0; i < protocols.length; i++) {
@@ -97,20 +141,22 @@ public class NettyProtocolRegistry {
 			Protocol equivalent = Protocol.fromVanilla(enumProtocol);
 			
 			// Associate known types
-			associatePackets(serverPackets.get(i), equivalent, Sender.SERVER);
+			associatePackets(result, serverMaps.get(i), equivalent, Sender.SERVER);
-			associatePackets(clientPackets.get(i), equivalent, Sender.CLIENT);
+			associatePackets(result, clientMaps.get(i), equivalent, Sender.CLIENT);
 		}
+		// Exchange (thread safe, as we have only one writer)
+		this.register = result;
 	}
 	
-	private void associatePackets(Map<Integer, Class<?>> lookup, Protocol protocol, Sender sender) {
+	private void associatePackets(Register register, Map<Integer, Class<?>> lookup, Protocol protocol, Sender sender) {
 		for (Entry<Integer, Class<?>> entry : lookup.entrySet()) {
 			PacketType type = PacketType.fromCurrent(protocol, sender, entry.getKey(), PacketType.UNKNOWN_PACKET);
-			typeToClass.put(type, entry.getValue());
+			register.typeToClass.put(type, entry.getValue());
 			
 			if (sender == Sender.SERVER)
-				serverPackets.add(type);
+				register.serverPackets.add(type);
 			if (sender == Sender.CLIENT)
-				clientPackets.add(type);
+				register.clientPackets.add(type);
 		}
 	}
 	

ProtocolLib/src/main/java/com/comphenix/protocol/injector/packet/InverseMaps.java

@@ -1,9 +1,7 @@
 package com.comphenix.protocol.injector.packet;
 
-import java.lang.reflect.Field;
 import java.util.Map;
 
-import com.comphenix.protocol.reflect.FieldUtils;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ForwardingMap;
 import com.google.common.collect.ForwardingMultimap;
@@ -65,66 +63,4 @@ public class InverseMaps {
 			}
 		};
 	}
-	
-	/**
-	 * Represents a class that can detect if a map has changed.
-	 * @author Kristian
-	 */
-	private static class MapContainer {
-		// For detecting changes
-		private Field modCountField;
-		private int lastModCount;
-		
-		// The object along with whether or not this is the initial run
-		private Object source;
-		private boolean changed;
-		
-		public MapContainer(Object source) {
-			this.source = source;
-			this.changed = true;
-			this.modCountField = FieldUtils.getField(source.getClass(), "modCount", true);
-		}
-		
-		/**
-		 * Determine if the map has changed.
-		 * @return TRUE if it has, FALSE otherwise.
-		 */
-		public boolean hasChanged() {
-			// Check if unchanged
-			checkChanged();
-			return changed;
-		}
-		
-		/**
-		 * Mark the map as changed or unchanged.
-		 * @param changed - TRUE if the map has changed, FALSE otherwise.
-		 */
-		public void setChanged(boolean changed) {
-			this.changed = changed;
-		}
-		
-		/**
-		 * Check for modifications to the current map.
-		 */
-		protected void checkChanged() {
-			if (!changed) {
-				if (getModificationCount() != lastModCount) {
-					lastModCount = getModificationCount();
-					changed = true;
-				}
-			}
-		}
-		
-		/**
-		 * Retrieve the current modification count.
-		 * @return The current count, or something different than lastModCount if not accessible.
-		 */
-		private int getModificationCount() {
-			try {
-				return modCountField != null ? modCountField.getInt(source) : lastModCount + 1;
-			} catch (Exception e) {
-				throw new RuntimeException("Unable to retrieve modCount.", e);
-			}
-		}
-	}
 }

ProtocolLib/src/main/java/com/comphenix/protocol/injector/packet/MapContainer.java

@@ -0,0 +1,67 @@
+package com.comphenix.protocol.injector.packet;
+
+import java.lang.reflect.Field;
+
+import com.comphenix.protocol.reflect.FieldUtils;
+
+/**
+ * Represents a class that can detect if a map has changed.
+ * @author Kristian
+ */
+public class MapContainer {
+	// For detecting changes
+	private Field modCountField;
+	private int lastModCount;
+	
+	// The object along with whether or not this is the initial run
+	private Object source;
+	private boolean changed;
+	
+	public MapContainer(Object source) {
+		this.source = source;
+		this.changed = true;
+		this.modCountField = FieldUtils.getField(source.getClass(), "modCount", true);
+	}
+	
+	/**
+	 * Determine if the map has changed.
+	 * @return TRUE if it has, FALSE otherwise.
+	 */
+	public boolean hasChanged() {
+		// Check if unchanged
+		checkChanged();
+		return changed;
+	}
+	
+	/**
+	 * Mark the map as changed or unchanged.
+	 * @param changed - TRUE if the map has changed, FALSE otherwise.
+	 */
+	public void setChanged(boolean changed) {
+		this.changed = changed;
+	}
+	
+	/**
+	 * Check for modifications to the current map.
+	 */
+	protected void checkChanged() {
+		if (!changed) {
+			if (getModificationCount() != lastModCount) {
+				lastModCount = getModificationCount();
+				changed = true;
+			}
+		}
+	}
+	
+	/**
+	 * Retrieve the current modification count.
+	 * @return The current count, or something different than lastModCount if not accessible.
+	 */
+	private int getModificationCount() {
+		try {
+			return modCountField != null ? modCountField.getInt(source) : lastModCount + 1;
+		} catch (Exception e) {
+			throw new RuntimeException("Unable to retrieve modCount.", e);
+		}
+	}
+}

ProtocolLib/src/main/java/com/comphenix/protocol/injector/packet/PacketRegistry.java

@@ -232,8 +232,10 @@ public class PacketRegistry {
 	public static Set<PacketType> getServerPacketTypes() {
 		initialize();
 		
-		if (NETTY != null)
+		if (NETTY != null) {
+			NETTY.synchronize();
 			return NETTY.getServerPackets();
+		}
 		
 		// Handle legacy
 		if (NETTY_SERVER_PACKETS == null) {
@@ -269,8 +271,10 @@ public class PacketRegistry {
 	public static Set<PacketType> getClientPacketTypes() {
 		initialize();
 		
-		if (NETTY != null)
+		if (NETTY != null) {
+			NETTY.synchronize();
 			return NETTY.getClientPackets();
+		}
 		
 		// Handle legacy
 		if (NETTY_CLIENT_PACKETS == null) {

ProtocolLib/src/test/java/com/comphenix/integration/protocol/TestPingPacket.java

@@ -8,9 +8,12 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
+import org.bukkit.entity.Player;
 import org.bukkit.plugin.Plugin;
 
 import com.comphenix.protocol.PacketType;
+import com.comphenix.protocol.PacketType.Protocol;
+import com.comphenix.protocol.PacketType.Sender;
 import com.comphenix.protocol.ProtocolLibrary;
 import com.comphenix.protocol.events.PacketAdapter;
 import com.comphenix.protocol.events.PacketEvent;
@@ -55,6 +58,8 @@ public class TestPingPacket {
 	}
 	
 	private Future<String> testInterception(Plugin test) {
+		PacketType customPacket = PacketType.fromCurrent(Protocol.STATUS, Sender.CLIENT, 3, -1);
+		
 		ProtocolLibrary.getProtocolManager().addPacketListener(
 		  new PacketAdapter(test, PacketType.Status.Server.OUT_SERVER_INFO) {
 			@Override