[TinyProtocol] Add support for intercepting login and status packets.

Examples/TinyProtocol/src/main/java/com/comphenix/tinyprotocol/ExamplePlugin.java

@@ -4,8 +4,11 @@ import org.bukkit.Location;
 import org.bukkit.entity.Player;
 import org.bukkit.plugin.java.JavaPlugin;
 
+import com.comphenix.tinyprotocol.Reflection.ConstructorInvoker;
 import com.comphenix.tinyprotocol.Reflection.FieldAccessor;
 
+import net.minecraft.util.io.netty.channel.Channel;
+
 public class ExamplePlugin extends JavaPlugin {
 	// Chat packets
 	private FieldAccessor<String> CHAT_MESSAGE = Reflection.getField(
@@ -19,13 +22,21 @@ public class ExamplePlugin extends JavaPlugin {
 	private FieldAccessor<Float> particleZ = Reflection.getField(particleClass, float.class, 2);
 	private FieldAccessor<Integer> particleCount = Reflection.getField(particleClass, int.class, 0);
 	
+	// Server info packet
+	private Class<?> serverInfoClass = Reflection.getClass("{nms}.PacketStatusOutServerInfo");
+	private Class<Object> serverPingClass = Reflection.getUntypedClass("{nms}.ServerPing");
+	private Class<Object> playerSampleClass = Reflection.getUntypedClass("{nms}.ServerPingPlayerSample");
+	private FieldAccessor<Object> serverPing = Reflection.getField(serverInfoClass, serverPingClass, 0);
+	private FieldAccessor<Object> playerSample = Reflection.getField(serverPingClass, playerSampleClass, 0);
+	private ConstructorInvoker playerSampleInvoker = Reflection.getConstructor(playerSampleClass, int.class, int.class);
+	
 	private TinyProtocol protocol;
 	
 	@Override
 	public void onEnable() {
 		protocol = new TinyProtocol(this) {
 			@Override
-			public Object onPacketInAsync(Player sender, Object packet) {
+			public Object onPacketInAsync(Player sender, Channel channel, Object packet) {
 				// Cancel chat packets
 				if (CHAT_MESSAGE.hasField(packet)) {
 					if (CHAT_MESSAGE.get(packet).contains("dirty")) {
@@ -36,7 +47,20 @@ public class ExamplePlugin extends JavaPlugin {
 				if (particleName.hasField(packet)) {
 					System.out.println("Sending particle field:" + packet);
 				}
-				return super.onPacketInAsync(sender, packet);
+				return super.onPacketInAsync(sender, channel, packet);
+			}
+			
+			@Override
+			public Object onPacketOutAsync(Player reciever, Channel channel, Object packet) {
+				if (serverInfoClass.isInstance(packet)) {
+					Object ping = serverPing.get(packet);
+					playerSample.set(ping, playerSampleInvoker.invoke(1000, 0));
+					
+					// Which is equivalent to:
+					//serverPing.get(packet).setPlayerSample(new ServerPingPlayerSample(1000, 0));
+					return packet;
+				}
+				return super.onPacketOutAsync(reciever, channel, packet);
 			}
 		};
 	}

Examples/TinyProtocol/src/main/java/com/comphenix/tinyprotocol/Reflection.java

@@ -1,5 +1,6 @@
 package com.comphenix.tinyprotocol;
 
+import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.Arrays;
@@ -8,11 +9,25 @@ import java.util.regex.Pattern;
 
 import org.bukkit.Bukkit;
 
+import com.comphenix.tinyprotocol.Reflection.FieldAccessor;
+
 /**
  * An utility class that simplifies reflection in Bukkit plugins.
  * @author Kristian
  */
 public final class Reflection {
+	/**
+	 * An interface for invoking a specific constructor. 
+	 */
+	public interface ConstructorInvoker {
+		/**
+		 * Invoke a constructor for a specific class.
+		 * @param arguments - the arguments to pass to the constructor.
+		 * @return The constructed object.
+		 */
+		public Object invoke(Object... arguments); 
+	}
+	
 	/**
 	 * An interface for invoking a specific method. 
 	 */
@@ -161,7 +176,7 @@ public final class Reflection {
      * @throws IllegalStateException If we cannot find this method.
      */
 	public static MethodInvoker getMethod(String className, String methodName, Class<?>... params) {
-		return getMethod(getClass(className), methodName, params);
+		return getTypedMethod(getClass(className), methodName, null, params);
 	}
 	
     /**
@@ -173,10 +188,24 @@ public final class Reflection {
      * @throws IllegalStateException If we cannot find this method.
      */
     public static MethodInvoker getMethod(Class<?> clazz, String methodName, Class<?>... params) {
+    	return getTypedMethod(clazz, methodName, null, params);
+    }
+	
+    /**
+     * Search for the first publically and privately defined method of the given name and parameter count.
+     * @param clazz - a class to start with.
+     * @param methodName - the method name, or NULL to skip.
+     * @param returnType - the expected return type, or NULL to ignore.
+     * @param params - the expected parameters.
+     * @return An object that invokes this specific method.
+     * @throws IllegalStateException If we cannot find this method.
+     */
+    public static MethodInvoker getTypedMethod(Class<?> clazz, String methodName, Class<?> returnType, Class<?>... params) {
         for (final Method method : clazz.getDeclaredMethods()) {
             if ((methodName == null || method.getName().equals(methodName)) &&
+            	(returnType == null) || method.getReturnType().equals(returnType) &&
                  Arrays.equals(method.getParameterTypes(), params)) {
-                
+
                 method.setAccessible(true);
                 return new MethodInvoker() {
                 	@Override
@@ -196,7 +225,61 @@ public final class Reflection {
         throw new IllegalStateException(String.format(
             "Unable to find method %s (%s).", methodName, Arrays.asList(params)));
     }
-	 
+    
+    /**
+     * Search for the first publically and privately defined constructor of the given name and parameter count.
+     * @param className - lookup name of the class, see {@link #getClass(String)}.
+     * @param params - the expected parameters.
+     * @return An object that invokes this constructor.
+     * @throws IllegalStateException If we cannot find this method.
+     */
+    public static ConstructorInvoker getConstructor(String className, Class<?>... params) {
+    	return getConstructor(getClass(className), params);
+    }
+    
+    /**
+     * Search for the first publically and privately defined constructor of the given name and parameter count.
+     * @param clazz - a class to start with.
+     * @param params - the expected parameters.
+     * @return An object that invokes this constructor.
+     * @throws IllegalStateException If we cannot find this method.
+     */
+    public static ConstructorInvoker getConstructor(Class<?> clazz, Class<?>... params) {
+        for (final Constructor<?> constructor : clazz.getDeclaredConstructors()) {
+            if (Arrays.equals(constructor.getParameterTypes(), params)) {
+
+            	constructor.setAccessible(true);
+                return new ConstructorInvoker() {
+                	@Override
+                	public Object invoke(Object... arguments) {
+                		try {
+							return constructor.newInstance(arguments);
+						} catch (Exception e) {
+							throw new RuntimeException("Cannot invoke constructor " + constructor, e);
+						}
+                	}
+                };
+            }
+        }
+        throw new IllegalStateException(String.format(
+            "Unable to find constructor for %s (%s).", clazz, Arrays.asList(params)));
+    }
+     
+    /**
+     * Retrieve a class from its full name, without knowing its type on compile time.
+     * <p>
+     * This is useful when looking up fields by a NMS or OBC type.
+     * <p>
+     * @see {@link #getClass()} for more information.
+     * @param lookupName - the class name with variables.
+     * @return The class.
+     */
+    public static Class<Object> getUntypedClass(String lookupName) {
+    	@SuppressWarnings({"rawtypes", "unchecked"})
+		Class<Object> clazz = (Class<Object>)(Class)getClass(lookupName);
+    	return clazz;
+    }
+    
     /**
      * Retrieve a class from its full name.
      * <p>

Examples/TinyProtocol/src/main/java/com/comphenix/tinyprotocol/TinyProtocol.java

@@ -1,53 +1,169 @@
 package com.comphenix.tinyprotocol;
 
+import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
 
+import net.minecraft.util.com.mojang.authlib.GameProfile;
 // These are not versioned, but they require CraftBukkit
 import net.minecraft.util.io.netty.channel.Channel;
 import net.minecraft.util.io.netty.channel.ChannelDuplexHandler;
+import net.minecraft.util.io.netty.channel.ChannelFuture;
 import net.minecraft.util.io.netty.channel.ChannelHandlerContext;
+import net.minecraft.util.io.netty.channel.ChannelInboundHandlerAdapter;
+import net.minecraft.util.io.netty.channel.ChannelInitializer;
+import net.minecraft.util.io.netty.channel.ChannelPipeline;
 import net.minecraft.util.io.netty.channel.ChannelPromise;
 
+import org.bukkit.Bukkit;
 import org.bukkit.entity.Player;
 import org.bukkit.event.EventHandler;
 import org.bukkit.event.EventPriority;
 import org.bukkit.event.HandlerList;
 import org.bukkit.event.Listener;
-import org.bukkit.event.player.PlayerJoinEvent;
+import org.bukkit.event.player.PlayerLoginEvent;
 import org.bukkit.event.server.PluginDisableEvent;
 import org.bukkit.plugin.Plugin;
 
 import com.comphenix.tinyprotocol.Reflection.FieldAccessor;
 import com.comphenix.tinyprotocol.Reflection.MethodInvoker;
+import com.google.common.collect.Lists;
 import com.google.common.collect.MapMaker;
 
 /**
  * Represents a very tiny alternative to ProtocolLib in 1.7.2.
  * <p>
- * Note that it does not support intercepting packets sent during login (such as OUT_SERVER_PING).
+ * It now supports intercepting packets during login and status ping (such as OUT_SERVER_PING)!
  * @author Kristian
  */
 public abstract class TinyProtocol {	
 	// Used in order to lookup a channel
-	private MethodInvoker getPlayerHandle = Reflection.getMethod("{obc}.entity.CraftPlayer", "getHandle");
-	private FieldAccessor<Object> getConnection = Reflection.getField("{nms}.EntityPlayer", "playerConnection", Object.class);
-	private FieldAccessor<Object> getManager = Reflection.getField("{nms}.PlayerConnection", "networkManager", Object.class);
-	private FieldAccessor<Channel> getChannel = Reflection.getField("{nms}.NetworkManager", Channel.class, 0);
+	private static final MethodInvoker getPlayerHandle = Reflection.getMethod("{obc}.entity.CraftPlayer", "getHandle");
+	private static final FieldAccessor<Object> getConnection = Reflection.getField("{nms}.EntityPlayer", "playerConnection", Object.class);
+	private static final FieldAccessor<Object> getManager = Reflection.getField("{nms}.PlayerConnection", "networkManager", Object.class);
+	private static final FieldAccessor<Channel> getChannel = Reflection.getField("{nms}.NetworkManager", Channel.class, 0);
+	
+	// Looking up ServerConnection
+	private static final Class<Object> minecraftServerClass = Reflection.getUntypedClass("{nms}.MinecraftServer");
+	private static final Class<Object> serverConnectionClass = Reflection.getUntypedClass("{nms}.ServerConnection");
+	private static final FieldAccessor<Object> getMinecraftServer = Reflection.getField("{obc}.CraftServer", minecraftServerClass, 0);
+	private static final FieldAccessor<Object> getServerConnection = Reflection.getField(minecraftServerClass, serverConnectionClass, 0);
+	private static final MethodInvoker getNetworkMarkers = Reflection.getTypedMethod(serverConnectionClass, null, List.class, serverConnectionClass);
+	
+	// Packets we have to intercept
+	private static final Class<?> PACKET_LOGIN_IN_START = Reflection.getMinecraftClass("PacketLoginInStart");
+	private static final FieldAccessor<GameProfile> getGameProfile = Reflection.getField(PACKET_LOGIN_IN_START, GameProfile.class, 0);
 	
 	// Speedup channel lookup
-	private Map<Player, Channel> channelLookup = new MapMaker().weakKeys().makeMap();
+	private Map<String, Channel> channelLookup = new MapMaker().weakKeys().makeMap();
 	private Listener listener;
 	
-	protected boolean closed;
+	// List of network markers
+	private List<Object> networkManagers;
+	
+	// Injected channel handlers
+	private List<Channel> serverChannels = Lists.newArrayList();
+	private ChannelInboundHandlerAdapter serverChannelHandler;
+	private ChannelInitializer<Channel> beginInitProtocol;
+	private ChannelInitializer<Channel> endInitProtocol;
+	
+	protected volatile boolean closed;
 	protected Plugin plugin;
 	
 	public TinyProtocol(Plugin plugin) {
 		this.plugin = plugin;
 		this.plugin.getServer().getPluginManager().registerEvents(
-				listener = createListener(), plugin);
+			listener = createListener(), plugin);
 		
 		// Prepare existing players
+		registerChannelHandler();
+		registerPlayers(plugin);	
+	}
+
+	private void createServerChannelHandler() {
+        // Handle connected channels
+        endInitProtocol = new ChannelInitializer<Channel>() {
+            @Override
+            protected void initChannel(Channel channel) throws Exception {
+            	try {
+                    // This can take a while, so we need to stop the main thread from interfering
+                    synchronized (networkManagers) {
+                    	// Stop injecting channels
+                    	if (closed)
+                    		return;
+                    	injectChannel(channel);
+					}
+            	} catch (Exception e) {
+            		plugin.getLogger().log(Level.SEVERE, "Cannot inject incomming channel " + channel, e);
+            	}
+            }
+        };
+        
+        // This is executed before Minecraft's channel handler
+        beginInitProtocol = new ChannelInitializer<Channel>() {
+        	@Override
+        	protected void initChannel(Channel channel) throws Exception {
+        		channel.pipeline().addLast(endInitProtocol);
+        	}
+        };
+		
+        serverChannelHandler = new ChannelInboundHandlerAdapter() {
+        	@Override
+        	public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+                Channel channel = (Channel) msg;
+
+                // Prepare to initialize ths channel
+                channel.pipeline().addFirst(beginInitProtocol);
+                ctx.fireChannelRead(msg);
+        	}
+        };
+	}
+	
+	@SuppressWarnings("unchecked")
+	private void registerChannelHandler() {
+		Object mcServer = getMinecraftServer.get(Bukkit.getServer());
+		Object serverConnection = getServerConnection.get(mcServer);
+		boolean looking = true;
+		
+		// We need to synchronize against this list
+		networkManagers = (List<Object>) getNetworkMarkers.invoke(null, serverConnection);
+		createServerChannelHandler();
+		
+		// Find the correct list, or implicitly throw an exception
+		for (int i = 0; looking; i++) {
+			List<Object> list = Reflection.getField(serverConnection.getClass(), List.class, i).get(serverConnection);
+			
+			for (Object item : list) {
+				if (!ChannelFuture.class.isInstance(item)) 
+					break;
+				
+				// Channel future that contains the server connection
+				Channel serverChannel = ((ChannelFuture)item).channel();
+	
+				serverChannels.add(serverChannel);
+				serverChannel.pipeline().addFirst(serverChannelHandler);
+				looking = false;
+			}
+		}
+	}
+	
+	private void unregisterChannelHandler() {
+		if (serverChannelHandler == null)
+			return;
+		
+		for (Channel serverChannel : serverChannels) {
+			final ChannelPipeline pipeline = serverChannel.pipeline();
+			
+			// Remove channel handler
+			serverChannel.eventLoop().execute(new Runnable() {
+				public void run() {
+					pipeline.remove(serverChannelHandler);
+				}
+			});
+		}
+	}
+	
+	private void registerPlayers(Plugin plugin) {
 		for (Player player : plugin.getServer().getOnlinePlayers()) {
 			injectPlayer(player);
 		}
@@ -57,21 +173,25 @@ public abstract class TinyProtocol {
 	 * Invoked when the server is starting to send a packet to a player.
 	 * <p>
 	 * Note that this is not executed on the main thread.
-	 * @param reciever - the receiving player.
+	 * @param reciever - the receiving player, NULL for early login/status packets.
+	 * @param remoteAddress - remote address of the sending client. Never NULL.
 	 * @param packet - the packet being sent.
 	 * @return The packet to send instead, or NULL to cancel the transmission.
 	 */
-	public Object onPacketOutAsync(Player reciever, Object packet) {
+	public Object onPacketOutAsync(Player reciever, Channel channel, Object packet) {
 		return packet;
 	}
 
 	/**
 	 * Invoked when the server has received a packet from a given player.
-	 * @param sender - the player that sent the packet.
-	 * @param packet - the packet being sent.
+	 * <p>
+	 * Use {@link Channel#remoteAddress()} to get the remote address of the client.
+	 * @param sender - the player that sent the packet, NULL for early login/status packets.
+	 * @param channel - channel that received the packet. Never NULL.
+	 * @param packet - the packet being received.
 	 * @return The packet to recieve instead, or NULL to cancel.
 	 */
-	public Object onPacketInAsync(Player sender, Object packet) {
+	public Object onPacketInAsync(Player sender, Channel channel, Object packet) {
 		return packet;
 	}
 	
@@ -83,9 +203,20 @@ public abstract class TinyProtocol {
 	 * @param packet - the packet to send.
 	 */
 	public void sendPacket(Player player, Object packet) {
-		getChannel(player).pipeline().writeAndFlush(packet);
+		sendPacket(getChannel(player), packet);
 	}
 	
+	/**
+	 * Send a packet to a particular client.
+	 * <p>
+	 * Note that {@link #onPacketOutAsync(Player, Object)} will be invoked with this packet.
+	 * @param channel - client identified by a channel.
+	 * @param packet - the packet to send.
+	 */
+	public void sendPacket(Channel channel, Object packet) {
+		channel.pipeline().writeAndFlush(packet);
+	}
+		
 	/**
 	 * Pretend that a given packet has been received from a player.
 	 * <p>
@@ -94,7 +225,18 @@ public abstract class TinyProtocol {
 	 * @param packet - the packet that will be received by the server.
 	 */
 	public void receivePacket(Player player, Object packet) {
-		getChannel(player).pipeline().context("encoder").fireChannelRead(packet);
+		receivePacket(getChannel(player), packet);
+	}
+	
+	/**
+	 * Pretend that a given packet has been received from a given client.
+	 * <p>
+	 * Note that {@link #onPacketInAsync(Player, Object)} will be invoked with this packet.
+	 * @param channel - client identified by a channel.
+	 * @param packet - the packet that will be received by the server.
+	 */
+	public void receivePacket(Channel channel, Object packet) {
+		channel.pipeline().context("encoder").fireChannelRead(packet);
 	}
 	
 	/**
@@ -106,40 +248,36 @@ public abstract class TinyProtocol {
 	protected String getHandlerName() {
 		return "tiny-" + plugin.getName();
 	}
-
+	
 	/**
-	 * Add a custom channel handler to the given player's channel pipeline, allowing us to intercept sent and received packets.
+	 * Add a custom channel handler to the given player's channel pipeline, 
+	 * allowing us to intercept sent and received packets.
 	 * @param player - the player to inject.
 	 */
-	private void injectPlayer(final Player player) {
-		// Inject our packet interceptor
-		getChannel(player).pipeline().addBefore("packet_handler", getHandlerName(), new ChannelDuplexHandler() {
-			@Override
-			public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-				try {
-					msg = onPacketInAsync(player, msg);
-				} catch (Exception e) {
-					plugin.getLogger().log(Level.SEVERE, "Error in onPacketInAsync().", e);
-				}
-				
-				if (msg != null) {
-					super.channelRead(ctx, msg);
-				}
-			}
+	private void injectPlayer(Player player) {
+		injectChannel(getChannel(player)).player = player;
+	}
+	
+	/**
+	 * Add a custom channel handler to the given channel.
+	 * @param player - the channel to inject.
+	 */
+	private PacketInterceptor injectChannel(Channel channel) {
+		String handlerName = getHandlerName();
+		
+		try {
+			PacketInterceptor interceptor = (PacketInterceptor) channel.pipeline().get(handlerName);
 			
-			@Override
-			public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
-				try {
-					msg = onPacketOutAsync(player, msg);
-				} catch (Exception e) {
-					plugin.getLogger().log(Level.SEVERE, "Error in onPacketOutAsync().", e);
-				}
-				
-				if (msg != null) {
-					super.write(ctx, msg, promise);
-				}
+			// Inject our packet interceptor
+			if (interceptor == null) {
+				interceptor = new PacketInterceptor();
+				channel.pipeline().addBefore("packet_handler", handlerName, interceptor);
 			}
-		});
+			return interceptor;	
+		} catch (IllegalArgumentException e) {
+			// Try again
+			return (PacketInterceptor) channel.pipeline().get(handlerName);
+		}
 	}
 	
 	/**
@@ -148,14 +286,14 @@ public abstract class TinyProtocol {
 	 * @return The Netty channel.
 	 */
 	private Channel getChannel(Player player) {
-		Channel channel = channelLookup.get(player);
+		Channel channel = channelLookup.get(player.getName());
 		
 		// Lookup channel again
 		if (channel == null) {
 			Object connection = getConnection.get(getPlayerHandle.invoke(player));
 			Object manager = getManager.get(connection);
 			
-			channelLookup.put(player, channel = getChannel.get(manager));
+			channelLookup.put(player.getName(), channel = getChannel.get(manager));
 		}
 		return channel;
 	}
@@ -166,7 +304,7 @@ public abstract class TinyProtocol {
 	private Listener createListener() {
 		return new Listener() {
 			@EventHandler(priority = EventPriority.LOWEST)
-			public final void onPlayerJoin(PlayerJoinEvent e) {
+			public final void onPlayerLogin(PlayerLoginEvent e) {
 				if (closed) 
 					return;
 				injectPlayer(e.getPlayer());
@@ -205,6 +343,52 @@ public abstract class TinyProtocol {
 			
 			// Clean up Bukkit
 			HandlerList.unregisterAll(listener);
+			unregisterChannelHandler();
+		}
+	}
+	
+	/**
+	 * Channel handler that is inserted into the player's channel pipeline, allowing us to intercept sent and received packets.
+	 * @author Kristian
+	 */
+	private final class PacketInterceptor extends ChannelDuplexHandler {		
+		// Updated by the login event
+		public volatile Player player;
+		
+		@Override
+		public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+			// Intercept channel
+			final Channel channel = ctx.channel();
+			handleLoginStart(channel, msg);
+			
+			try {
+				msg = onPacketInAsync(player, channel, msg);
+			} catch (Exception e) {
+				plugin.getLogger().log(Level.SEVERE, "Error in onPacketInAsync().", e);
+			}
+			
+			if (msg != null) {
+				super.channelRead(ctx, msg);
+			}
+		}
+		@Override
+		public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {			
+			try {
+				msg = onPacketOutAsync(player, ctx.channel(), msg);
+			} catch (Exception e) {
+				plugin.getLogger().log(Level.SEVERE, "Error in onPacketOutAsync().", e);
+			}
+			
+			if (msg != null) {
+				super.write(ctx, msg, promise);
+			}
+		}
+		
+		private void handleLoginStart(Channel channel, Object packet) {
+			if (PACKET_LOGIN_IN_START.isInstance(packet)) {
+				GameProfile profile = getGameProfile.get(packet);
+				channelLookup.put(profile.getName(), channel);
+			}
 		}
 	}
 }