From 8d0e8139de024f7fc705d8b972d52f215d50040c Mon Sep 17 00:00:00 2001
From: "Kristian S. Stangeland" <kr_stang@hotmail.com>
Date: Sun, 23 Jun 2013 21:37:37 +0200
Subject: [PATCH] Add support for the Lilypad Bukkit Connector. FIXES
 Ticket-101

---
 .../injector/player/PlayerInjector.java         |  2 +-
 .../player/ProxyPlayerInjectionHandler.java     | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/PlayerInjector.java b/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/PlayerInjector.java
index bf5a0cd9..79bb4c00 100644
--- a/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/PlayerInjector.java
+++ b/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/PlayerInjector.java
@@ -284,7 +284,7 @@ public abstract class PlayerInjector implements SocketInjector {
 	
 	/**
 	 * Retrieve the associated remote address of a player.
-	 * @return The associated remote address..
+	 * @return The associated remote address.
 	 * @throws IllegalAccessException If we're unable to read the socket address field.
 	 */
 	@Override
diff --git a/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/ProxyPlayerInjectionHandler.java b/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/ProxyPlayerInjectionHandler.java
index 1f3fa663..15570caf 100644
--- a/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/ProxyPlayerInjectionHandler.java
+++ b/ProtocolLib/src/main/java/com/comphenix/protocol/injector/player/ProxyPlayerInjectionHandler.java
@@ -21,6 +21,7 @@ import java.io.DataInputStream;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.InvocationTargetException;
 import java.net.InetSocketAddress;
+import java.net.Socket;
 import java.net.SocketAddress;
 import java.util.Map;
 import java.util.Set;
@@ -51,6 +52,7 @@ import com.comphenix.protocol.injector.server.SocketInjector;
 import com.comphenix.protocol.utility.MinecraftReflection;
 import com.comphenix.protocol.utility.MinecraftVersion;
 
+import com.google.common.base.Objects;
 import com.google.common.base.Predicate;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
@@ -350,6 +352,7 @@ class ProxyPlayerInjectionHandler implements PlayerInjectionHandler {
 							return null;
 						
 						SocketInjector previous = inputStreamLookup.peekSocketInjector(address);
+						Socket socket = injector.getSocket();
 
 						// Close any previously associated hooks before we proceed
 						if (previous != null && !(player instanceof Factory)) {
@@ -363,8 +366,7 @@ class ProxyPlayerInjectionHandler implements PlayerInjectionHandler {
 						}
 						injector.injectManager();
 						
-						// Save injector
-						inputStreamLookup.setSocketInjector(address, injector);
+						saveAddressLookup(address, socket, injector);
 						break;
 					}
 					
@@ -413,6 +415,17 @@ class ProxyPlayerInjectionHandler implements PlayerInjectionHandler {
 		return injector;
 	}
 
+	private void saveAddressLookup(SocketAddress address, Socket socket, SocketInjector injector) {
+		SocketAddress socketAddress = socket != null ? socket.getRemoteSocketAddress() : null;
+		
+		if (socketAddress != null && !Objects.equal(socketAddress, address)) {
+			// Save this version as well
+			inputStreamLookup.setSocketInjector(socketAddress, injector);
+		}
+		// Save injector
+		inputStreamLookup.setSocketInjector(address, injector);
+	}
+	
 	private void cleanupHook(PlayerInjector injector) {
 		// Clean up as much as possible
 		try {