Use a stateful login system for the initial handler.

proxy/src/main/java/net/md_5/bungee/InitialHandler.java

@@ -1,12 +1,11 @@
 package net.md_5.bungee;
 
-import net.md_5.bungee.config.Configuration;
+import com.google.common.base.Preconditions;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.util.ArrayList;
-import java.util.List;
 import javax.crypto.SecretKey;
 import lombok.Getter;
 import net.md_5.bungee.api.ChatColor;
@@ -18,15 +17,19 @@ import net.md_5.bungee.api.connection.PendingConnection;
 import net.md_5.bungee.api.connection.ProxiedPlayer;
 import net.md_5.bungee.api.event.LoginEvent;
 import net.md_5.bungee.api.event.ProxyPingEvent;
+import net.md_5.bungee.packet.DefinedPacket;
 import net.md_5.bungee.packet.Packet2Handshake;
+import net.md_5.bungee.packet.PacketCDClientStatus;
 import net.md_5.bungee.packet.PacketFCEncryptionResponse;
 import net.md_5.bungee.packet.PacketFDEncryptionRequest;
+import net.md_5.bungee.packet.PacketFEPing;
 import net.md_5.bungee.packet.PacketFFKick;
+import net.md_5.bungee.packet.PacketHandler;
 import net.md_5.bungee.packet.PacketInputStream;
 import org.bouncycastle.crypto.io.CipherInputStream;
 import org.bouncycastle.crypto.io.CipherOutputStream;
 
-public class InitialHandler implements Runnable, PendingConnection
+public class InitialHandler extends PacketHandler implements Runnable, PendingConnection
 {
 
     private final Socket socket;
@@ -35,6 +38,8 @@ public class InitialHandler implements Runnable, PendingConnection
     private PacketInputStream in;
     private OutputStream out;
     private Packet2Handshake handshake;
+    private PacketFDEncryptionRequest request;
+    private State thisState = State.HANDSHAKE;
 
     public InitialHandler(Socket socket, ListenerInfo info) throws IOException
     {
@@ -44,22 +49,56 @@ public class InitialHandler implements Runnable, PendingConnection
         out = socket.getOutputStream();
     }
 
-    @Override
+    private enum State
-    public void run()
     {
+
+        HANDSHAKE, ENCRYPT, LOGIN, FINISHED;
+    }
+
+    @Override
+    public void handle(PacketFEPing ping) throws Exception
+    {
+        socket.setSoTimeout(100);
+        boolean newPing = false;
         try
         {
-            byte[] packet = in.readPacket();
+            socket.getInputStream().read();
-            int id = Util.getId(packet);
+            newPing = true;
-            switch (id)
+        } catch (IOException ex)
         {
-                case 0x02:
+        }
-                    handshake = new Packet2Handshake(packet);
-                    PacketFDEncryptionRequest request = EncryptionUtil.encryptRequest();
-                    out.write(request.getPacket());
-                    PacketFCEncryptionResponse response = new PacketFCEncryptionResponse(in.readPacket());
 
-                    SecretKey shared = EncryptionUtil.getSecret(response, request);
+        ServerPing pingevent = new ServerPing(BungeeCord.PROTOCOL_VERSION, BungeeCord.GAME_VERSION,
+                listener.getMotd(), ProxyServer.getInstance().getPlayers().size(), listener.getMaxPlayers());
+
+        pingevent = ProxyServer.getInstance().getPluginManager().callEvent(new ProxyPingEvent(this, pingevent)).getResponse();
+
+        String response = (newPing) ? ChatColor.COLOR_CHAR + "1"
+                + "\00" + pingevent.getProtocolVersion()
+                + "\00" + pingevent.getGameVersion()
+                + "\00" + pingevent.getMotd()
+                + "\00" + pingevent.getCurrentPlayers()
+                + "\00" + pingevent.getMaxPlayers()
+                : pingevent.getMotd() + ChatColor.COLOR_CHAR + pingevent.getCurrentPlayers() + ChatColor.COLOR_CHAR + pingevent.getMaxPlayers();
+        disconnect(response);
+    }
+
+    @Override
+    public void handle(Packet2Handshake handshake) throws Exception
+    {
+        Preconditions.checkState(thisState == State.HANDSHAKE);
+        this.handshake = handshake;
+        request = EncryptionUtil.encryptRequest();
+        out.write(request.getPacket());
+        thisState = State.ENCRYPT;
+    }
+
+    @Override
+    public void handle(PacketFCEncryptionResponse encryptResponse) throws Exception
+    {
+        Preconditions.checkState(thisState == State.ENCRYPT);
+
+        SecretKey shared = EncryptionUtil.getSecret(encryptResponse, request);
         if (!EncryptionUtil.isAuthenticated(handshake.username, request.serverId, shared))
         {
             throw new KickException("Not authenticated with minecraft.net");
@@ -76,60 +115,40 @@ public class InitialHandler implements Runnable, PendingConnection
         LoginEvent event = new LoginEvent(this);
         if (event.isCancelled())
         {
-                        throw new KickException(event.getCancelReason());
+            disconnect(event.getCancelReason());
         }
 
         out.write(new PacketFCEncryptionResponse().getPacket());
         in = new PacketInputStream(new CipherInputStream(socket.getInputStream(), EncryptionUtil.getCipher(false, shared)));
         out = new CipherOutputStream(socket.getOutputStream(), EncryptionUtil.getCipher(true, shared));
-                    List<byte[]> customPackets = new ArrayList<>();
+
-                    byte[] custom;
+        thisState = State.LOGIN;
-                    while (Util.getId((custom = in.readPacket())) != 0xCD)
-                    {
-                        customPackets.add(custom);
     }
 
-                    UserConnection userCon = new UserConnection(socket, this, in, out, handshake, customPackets);
+    @Override
+    public void handle(PacketCDClientStatus clientStatus) throws Exception
+    {
+        Preconditions.checkState(thisState == State.LOGIN);
+
+        UserConnection userCon = new UserConnection(socket, this, in, out, handshake, new ArrayList<byte[]>());
         String server = ProxyServer.getInstance().getReconnectHandler().getServer(userCon);
         ServerInfo s = BungeeCord.getInstance().config.getServers().get(server);
         userCon.connect(s);
-                    break;
+
-                case 0xFE:
+        thisState = State.FINISHED;
-                    socket.setSoTimeout(100);
+    }
-                    boolean newPing = false;
+
+    @Override
+    public void run()
+    {
         try
         {
-                        socket.getInputStream().read();
+            while (thisState != State.FINISHED)
-                        newPing = true;
-                    } catch (IOException ex)
             {
+                byte[] buf = in.readPacket();
+                DefinedPacket packet = DefinedPacket.packet(buf);
+                packet.handle(this);
             }
-
-                    ServerPing pingevent = new ServerPing(BungeeCord.PROTOCOL_VERSION, BungeeCord.GAME_VERSION,
-                            listener.getMotd(), ProxyServer.getInstance().getPlayers().size(), listener.getMaxPlayers());
-
-                    pingevent = ProxyServer.getInstance().getPluginManager().callEvent(new ProxyPingEvent(this, pingevent)).getResponse();
-
-                    String ping = (newPing) ? ChatColor.COLOR_CHAR + "1"
-                            + "\00" + pingevent.getProtocolVersion()
-                            + "\00" + pingevent.getGameVersion()
-                            + "\00" + pingevent.getMotd()
-                            + "\00" + pingevent.getCurrentPlayers()
-                            + "\00" + pingevent.getMaxPlayers()
-                            : pingevent.getMotd() + ChatColor.COLOR_CHAR + pingevent.getCurrentPlayers() + ChatColor.COLOR_CHAR + pingevent.getMaxPlayers();
-                    throw new KickException(ping);
-                default:
-                    if (id == 0xFA)
-                    {
-                        run(); // WTF Spoutcraft
-                    } else
-                    {
-                        // throw new IllegalArgumentException("Wasn't ready for packet id " + Util.hex(id));
-                    }
-            }
-        } catch (KickException ex)
-        {
-            disconnect(ex.getMessage());
         } catch (Exception ex)
         {
             disconnect("[Proxy Error] " + Util.exception(ex));
@@ -140,6 +159,7 @@ public class InitialHandler implements Runnable, PendingConnection
     @Override
     public void disconnect(String reason)
     {
+        thisState = State.FINISHED;
         try
         {
             out.write(new PacketFFKick(reason).getPacket());
@@ -149,7 +169,7 @@ public class InitialHandler implements Runnable, PendingConnection
         {
             try
             {
-                out.flush();
+                socket.shutdownOutput();
                 socket.close();
             } catch (IOException ioe2)
             {

proxy/src/main/java/net/md_5/bungee/packet/DefinedPacket.java

@@ -5,6 +5,8 @@ import com.google.common.io.ByteArrayDataOutput;
 import com.google.common.io.ByteStreams;
 import java.io.DataInput;
 import java.io.DataOutput;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
 import lombok.Delegate;
 import net.md_5.bungee.Util;
 
@@ -106,4 +108,46 @@ public abstract class DefinedPacket implements DataInput, DataOutput
 
     @Override
     public abstract String toString();
+
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
+    private static Class<? extends DefinedPacket>[] classes = new Class[256];
+
+    public static DefinedPacket packet(byte[] buf)
+    {
+        int id = Util.getId(buf);
+        Class<? extends DefinedPacket> clazz = classes[id];
+        DefinedPacket ret = null;
+        if (clazz != null)
+        {
+            try
+            {
+                Constructor<? extends DefinedPacket> constructor = clazz.getDeclaredConstructor(byte[].class);
+                if (constructor != null)
+                {
+                    ret = constructor.newInstance(buf);
+                }
+            } catch (IllegalAccessException | InstantiationException | InvocationTargetException | NoSuchMethodException ex)
+            {
+            }
+        }
+        return ret;
+    }
+
+    static
+    {
+        classes[0x00] = Packet0KeepAlive.class;
+        classes[0x01] = Packet1Login.class;
+        classes[0x02] = Packet2Handshake.class;
+        classes[0x03] = Packet3Chat.class;
+        classes[0x09] = Packet9Respawn.class;
+        classes[0xC9] = PacketC9PlayerListItem.class;
+        classes[0xCD] = PacketCDClientStatus.class;
+        classes[0xFA] = PacketFAPluginMessage.class;
+        classes[0xFC] = PacketFCEncryptionResponse.class;
+        classes[0xFE] = PacketFEPing.class;
+        classes[0xFF] = PacketFFKick.class;
+    }
 }

proxy/src/main/java/net/md_5/bungee/packet/Packet2Handshake.java

@@ -30,4 +30,10 @@ public class Packet2Handshake extends DefinedPacket
         this.host = readUTF();
         this.port = readInt();
     }
+
+    @Override
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
 }

proxy/src/main/java/net/md_5/bungee/packet/PacketCDClientStatus.java

@@ -18,4 +18,15 @@ public class PacketCDClientStatus extends DefinedPacket
         super(0xCD);
         writeByte(payload);
     }
+
+    public PacketCDClientStatus(byte[] buf)
+    {
+        super(0xCD, buf);
+    }
+
+    @Override
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
 }

proxy/src/main/java/net/md_5/bungee/packet/PacketFCEncryptionResponse.java

@@ -31,4 +31,10 @@ public class PacketFCEncryptionResponse extends DefinedPacket
         this.sharedSecret = readArray();
         this.verifyToken = readArray();
     }
+
+    @Override
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
 }

proxy/src/main/java/net/md_5/bungee/packet/PacketFDEncryptionRequest.java

@@ -30,4 +30,10 @@ public class PacketFDEncryptionRequest extends DefinedPacket
         publicKey = readArray();
         verifyToken = readArray();
     }
+
+    @Override
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
 }

proxy/src/main/java/net/md_5/bungee/packet/PacketFEPing.java

@@ -0,0 +1,21 @@
+package net.md_5.bungee.packet;
+
+import lombok.EqualsAndHashCode;
+import lombok.ToString;
+
+@ToString
+@EqualsAndHashCode(callSuper = false)
+public class PacketFEPing extends DefinedPacket
+{
+
+    public PacketFEPing(byte[] buffer)
+    {
+        super(0xFE, buffer);
+    }
+
+    @Override
+    public void handle(PacketHandler handler) throws Exception
+    {
+        handler.handle(this);
+    }
+}

proxy/src/main/java/net/md_5/bungee/packet/PacketHandler.java

@@ -0,0 +1,30 @@
+package net.md_5.bungee.packet;
+
+public abstract class PacketHandler
+{
+
+    public void handle(DefinedPacket packet) throws Exception
+    {
+        throw new UnsupportedOperationException("No handler defined for packet " + packet.getClass());
+    }
+
+    public void handle(Packet2Handshake handshake) throws Exception
+    {
+    }
+
+    public void handle(PacketCDClientStatus clientStatus) throws Exception
+    {
+    }
+
+    public void handle(PacketFCEncryptionResponse encryptResponse) throws Exception
+    {
+    }
+
+    public void handle(PacketFDEncryptionRequest encryptRequest) throws Exception
+    {
+    }
+
+    public void handle(PacketFEPing ping) throws Exception
+    {
+    }
+}