Use socket as key instead of input stream.

ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/ProxyPlayerInjectionHandler.java

@@ -284,7 +284,6 @@ class ProxyPlayerInjectionHandler implements PlayerInjectionHandler {
 	
 	// Unsafe variant of the above
 	private PlayerInjector injectPlayerInternal(Player player, Object injectionPoint, GamePhase phase) {
-		
 		PlayerInjector injector = playerInjection.get(player);
 		PlayerInjectHooks tempHook = getPlayerHook(phase);
 		PlayerInjectHooks permanentHook = tempHook;
@@ -310,21 +309,24 @@ class ProxyPlayerInjectionHandler implements PlayerInjectionHandler {
 					if (injector.canInject(phase)) {
 						injector.initialize(injectionPoint);
 						
-						DataInputStream inputStream = injector.getInputStream(false);
+						// Get socket and socket injector
 						Socket socket = injector.getSocket();
+						SocketInjector previous = null;
 						
-						// Guard against NPE here too
-						SocketInjector previous = socket != null ? inputStreamLookup.getSocketInjector(socket) : null;
+						// Due to a race condition, the main server "accept connections" thread may 
+						// get a closed network manager with a NULL input stream, 
+						if (socket == null) {
 							
-						// Close any previously associated hooks before we proceed
-						if (previous != null) {
-							uninjectPlayer(previous.getPlayer(), true);
 						}
 						
+						// Close any previously associated hooks before we proceed
+						if (previous != null && previous instanceof PlayerInjector) {
+							uninjectPlayer(previous.getPlayer(), true);
+						}
 						injector.injectManager();
 						
 						// Save injector
-						inputStreamLookup.setSocketInjector(inputStream, injector);
+						inputStreamLookup.setSocketInjector(socket, injector);
 						break;
 					}
 					

ProtocolLib/src/main/java/com/comphenix/protocol/injector/server/AbstractInputStreamLookup.java

@@ -5,8 +5,6 @@ import java.io.InputStream;
 import java.lang.reflect.Field;
 import java.net.Socket;
 import java.net.SocketAddress;
-import java.util.concurrent.ConcurrentMap;
-
 import org.bukkit.Server;
 import org.bukkit.entity.Player;
 
@@ -14,16 +12,11 @@ import com.comphenix.protocol.error.ErrorReporter;
 import com.comphenix.protocol.reflect.FieldAccessException;
 import com.comphenix.protocol.reflect.FieldUtils;
 import com.comphenix.protocol.reflect.FuzzyReflection;
-import com.google.common.collect.MapMaker;
 
 public abstract class AbstractInputStreamLookup {
 	// Used to access the inner input stream of a filtered input stream
 	private static Field filteredInputField;
 
-	// Using weak keys and values ensures that we will not hold up garbage collection
-	protected ConcurrentMap<InputStream, SocketInjector> ownerSocket = new MapMaker().weakKeys().makeMap();
-	protected ConcurrentMap<SocketAddress, InputStream> addressLookup = new MapMaker().weakValues().makeMap();
-
 	// Error reporter
 	protected final ErrorReporter reporter;
 	
@@ -82,63 +75,45 @@ public abstract class AbstractInputStreamLookup {
 
 	/**
 	 * Retrieve the associated socket injector for a player.
-	 * @param filtered - the indentifying filtered input stream.
+	 * @param input - the indentifying filtered input stream.
 	 * @return The socket injector we have associated with this player.
 	 */
 	public abstract SocketInjector getSocketInjector(InputStream input);
 	
-	/**
-	 * Retrieve a injector by its address.
-	 * @param address - the address of the socket.
-	 * @return The socket injector.
-	 */
-	public abstract SocketInjector getSocketInjector(SocketAddress address);
-	
 	/**
 	 * Retrieve an injector by its socket.
 	 * @param socket - the socket.
 	 * @return The socket injector.
 	 */
-	public SocketInjector getSocketInjector(Socket socket) {
-		if (socket == null)
-			throw new IllegalArgumentException("The socket cannot be NULL.");
-		return getSocketInjector(socket.getRemoteSocketAddress());
-	}
+	public abstract SocketInjector getSocketInjector(Socket socket);
 	
 	/**
-	 * Associate a given input stream with the provided socket injector.
-	 * @param input - the filtered input stream to associate.
-	 * @param injector - the injector.
-	 * @throws FieldAccessException Unable to access input stream.
+	 * Retrieve a injector by its address.
+	 * @param address - the address of the socket.
+	 * @return The socket injector, or NULL if not found.
 	 */
-	public void setSocketInjector(FilterInputStream input, SocketInjector injector) {
-		setSocketInjector(getInputStream(input), injector);
-	}
+	public abstract SocketInjector getSocketInjector(SocketAddress address);
 
 	/**
-	 * Associate a given input stream with the provided socket injector.
-	 * @param input - the input stream to associate.
+	 * Associate a given socket the provided socket injector.
+	 * @param input - the socket to associate.
 	 * @param injector - the injector.
 	 */
-	public void setSocketInjector(InputStream input, SocketInjector injector) {
-		SocketInjector previous = ownerSocket.put(input, injector);
+	public abstract void setSocketInjector(Socket socket, SocketInjector injector);
 
-		// Any previous temporary players will also be associated
-		if (previous != null) {
+	/**
+	 * If a player can hold a reference to its parent injector, this method will update that reference.
+	 * @param previous - the previous injector.
+	 * @param current - the new injector.
+	 */
+	protected void onPreviousSocketOverwritten(SocketInjector previous, SocketInjector current) {
 		Player player = previous.getPlayer();
 		
+		// Default implementation
 		if (player instanceof InjectContainer) {
 			InjectContainer container = (InjectContainer) player;
-				container.setInjector(injector);
+			container.setInjector(current);
 		}
-			
-			// Update the reference to any previous injector
-			onPreviousSocketOverwritten(previous, injector);
-		}
-	}
-	
-	protected void onPreviousSocketOverwritten(SocketInjector previous, SocketInjector current) {
-		// Do nothing
 	}
 	
 	/**

ProtocolLib/src/main/java/com/comphenix/protocol/injector/server/InjectContainer.java

@@ -7,7 +7,7 @@ package com.comphenix.protocol.injector.server;
  * @author Kristian
  */
 class InjectContainer {
-	private SocketInjector injector;
+	private volatile SocketInjector injector;
 
 	public SocketInjector getInjector() {
 		return injector;

ProtocolLib/src/main/java/com/comphenix/protocol/injector/server/InputStreamProxyLookup.java

@@ -12,6 +12,7 @@ import java.net.SocketAddress;
 import java.nio.charset.Charset;
 import java.util.Collections;
 import java.util.Set;
+import java.util.concurrent.ConcurrentMap;
 
 import org.bukkit.Server;
 import org.bukkit.entity.Player;
@@ -34,6 +35,11 @@ class InputStreamProxyLookup extends AbstractInputStreamLookup {
 	private static final int READ_TIMEOUT = 5000;
 	private static final int CONNECT_TIMEOUT = 1000;
 		
+	// Using weak keys and values ensures that we will not hold up garbage collection
+	protected ConcurrentMap<InputStream, SocketInjector> ownerSocket = new MapMaker().weakKeys().makeMap();
+	protected ConcurrentMap<SocketAddress, InputStream> addressLookup = new MapMaker().weakValues().makeMap();
+	protected ConcurrentMap<Socket, InputStream> socketLookup = new MapMaker().weakKeys().makeMap();
+	
 	// Fake connections
 	private Set<SocketAddress> fakeConnections = Collections.newSetFromMap(
 			new MapMaker().weakKeys().<SocketAddress, Boolean>makeMap()
@@ -100,6 +106,9 @@ class InputStreamProxyLookup extends AbstractInputStreamLookup {
 									// Update it
 									TemporaryPlayerFactory.setInjectorInPlayer(temporaryPlayer, socketInjector);
 								
+									// Socket lookup
+									socketLookup.put(this, input);
+									
 									// Associate the socket with a given input stream
 									setSocketInjector(input, socketInjector);
 								}
@@ -115,6 +124,54 @@ class InputStreamProxyLookup extends AbstractInputStreamLookup {
 		}
 	}
 	
+	@Override
+	public SocketInjector getSocketInjector(Socket socket) {
+		InputStream stream = getStream(socket);
+		
+		if (stream != null)
+			return getSocketInjector(stream);
+		else
+			return null;
+	}
+
+	@Override
+	public void setSocketInjector(Socket socket, SocketInjector injector) {
+		InputStream stream = getStream(socket);
+		
+		if (stream != null) {
+			socketLookup.put(socket, stream);
+			setSocketInjector(stream, injector);
+		}
+	}
+	
+	/**
+	 * Set the referenced socket injector by input stream.
+	 * @param stream - the input stream.
+	 * @param injector - the injector to reference.
+	 */
+	public void setSocketInjector(InputStream stream, SocketInjector injector) {
+		SocketInjector previous = ownerSocket.put(stream, injector);
+		
+		// Handle overwrite
+		if (previous != null) {
+			onPreviousSocketOverwritten(previous, injector);
+		}
+	}
+
+	private InputStream getStream(Socket socket) {
+		InputStream result = socketLookup.get(socket);
+		
+		// Use the socket as well
+		if (result == null) {
+			try {
+				result = socket.getInputStream();
+			} catch (IOException e) {
+				throw new RuntimeException("Unable to retrieve input stream from socket " + socket, e);
+			}
+		}
+		return result;
+	}
+	
 	@Override
 	public SocketInjector getSocketInjector(InputStream input) {
 		return ownerSocket.get(input);
@@ -213,6 +270,9 @@ class InputStreamProxyLookup extends AbstractInputStreamLookup {
 	
 	@Override
 	protected void onPreviousSocketOverwritten(SocketInjector previous, SocketInjector current) {
+		// Don't forget this
+		super.onPreviousSocketOverwritten(previous, current);
+		
 		if (previous instanceof DelegatedSocketInjector) {
 			DelegatedSocketInjector delegated = (DelegatedSocketInjector) previous;
 			

ProtocolLib/src/main/java/com/comphenix/protocol/injector/server/InputStreamReflectLookup.java

@@ -5,6 +5,7 @@ import java.io.InputStream;
 import java.lang.reflect.Field;
 import java.net.Socket;
 import java.net.SocketAddress;
+import java.util.concurrent.ConcurrentMap;
 
 import org.bukkit.Server;
 import org.bukkit.entity.Player;
@@ -13,8 +14,14 @@ import com.comphenix.protocol.error.ErrorReporter;
 import com.comphenix.protocol.reflect.FieldAccessException;
 import com.comphenix.protocol.reflect.FieldUtils;
 import com.comphenix.protocol.reflect.FuzzyReflection;
+import com.google.common.collect.MapMaker;
 
 class InputStreamReflectLookup extends AbstractInputStreamLookup {
+	// Using weak keys and values ensures that we will not hold up garbage collection
+	protected ConcurrentMap<Socket, SocketInjector> ownerSocket = new MapMaker().weakKeys().makeMap();
+	protected ConcurrentMap<SocketAddress, Socket> addressLookup = new MapMaker().weakValues().makeMap();
+	protected ConcurrentMap<InputStream, Socket> inputLookup = new MapMaker().weakValues().makeMap();
+	
 	// Used to create fake players
 	private TemporaryPlayerFactory tempPlayerFactory = new TemporaryPlayerFactory();
 
@@ -33,39 +40,86 @@ class InputStreamReflectLookup extends AbstractInputStreamLookup {
 	}
 
 	@Override
-	public SocketInjector getSocketInjector(InputStream input) {
-		SocketInjector injector = ownerSocket.get(input);
+	public SocketInjector getSocketInjector(Socket socket) {
+		SocketInjector result = ownerSocket.get(socket);
 		
-		if (injector != null) {
-			return injector;
-		} else {
-			try {
-				Socket socket = getSocket(input);
+		if (result == null) {
 			Player player = tempPlayerFactory.createTemporaryPlayer(server);
 			SocketInjector created = new TemporarySocketInjector(player, socket);
 							
-				// Update injector
+			result = ownerSocket.putIfAbsent(socket, created);
+			
+			if (result == null) {
+				// We won - use our created injector
 				TemporaryPlayerFactory.setInjectorInPlayer(player, created);
+				result = created;
+			}
+		}
+		return result;		
+	}
 	
-				// Save address too
-				addressLookup.put(socket.getRemoteSocketAddress(), input);
-				
-				// Associate the socket with a given input stream
-				setSocketInjector(input, created);
-				return created;
+	@Override
+	public SocketInjector getSocketInjector(InputStream input) {
+		try {
+			Socket socket = getSocket(input);
 			
+			// Guard against NPE
+			if (socket != null)
+				return getSocketInjector(socket);
+			else
+				return null;
 		} catch (IllegalAccessException e) {
 			throw new FieldAccessException("Cannot find or access socket field for " + input, e);
 		}
 	}
+
+	/**
+	 * Use reflection to get the underlying socket from an input stream.
+	 * @param stream - the socket stream to lookup.
+	 * @return The underlying socket, or NULL if not found.
+	 * @throws IllegalAccessException Unable to access socket field.
+	 */
+	private Socket getSocket(InputStream stream) throws IllegalAccessException {
+		// Extra check, just in case
+		if (stream instanceof FilterInputStream)
+			return getSocket(getInputStream((FilterInputStream) stream));
+		
+		Socket result = inputLookup.get(stream);
+		
+		if (result == null) {
+			result = lookupSocket(stream);
+			
+			// Save it
+			inputLookup.put(stream, result);
+		}
+		return result;
+	}
+	
+	@Override
+	public void setSocketInjector(Socket socket, SocketInjector injector) {
+		if (socket == null)
+			throw new IllegalArgumentException("socket cannot be NULL");
+		if (injector == null)
+			throw new IllegalArgumentException("injector cannot be NULL.");
+		
+		SocketInjector previous = ownerSocket.put(socket, injector);
+		
+		// Save the address lookup too
+		addressLookup.put(socket.getRemoteSocketAddress(), socket);
+		
+		// Any previous temporary players will also be associated
+		if (previous != null) {
+			// Update the reference to any previous injector
+			onPreviousSocketOverwritten(previous, injector);
+		}
 	}
 	
 	@Override
 	public SocketInjector getSocketInjector(SocketAddress address) {
-		InputStream input = addressLookup.get(address);
+		Socket socket = addressLookup.get(address);
 		
-		if (input != null)
-			return getSocketInjector(input);
+		if (socket != null)
+			return getSocketInjector(socket);
 		else
 			return null;
 	}
@@ -75,9 +129,15 @@ class InputStreamReflectLookup extends AbstractInputStreamLookup {
 		// Do nothing
 	}
 	
-	private static Socket getSocket(InputStream stream) throws IllegalAccessException {
+	/**
+	 * Lookup the underlying socket of a stream through reflection.
+	 * @param stream - the socket stream.
+	 * @return The underlying socket. 
+	 * @throws IllegalAccessException If reflection failed.
+	 */
+	private static Socket lookupSocket(InputStream stream) throws IllegalAccessException {
 		if (stream instanceof FilterInputStream) {
-			return getSocket(getInputStream((FilterInputStream) stream));
+			return lookupSocket(getInputStream((FilterInputStream) stream));
 		} else {
 			// Just do it
 			Field socketField = FuzzyReflection.fromObject(stream, true).