Finished the port to new session handler stuff

proxy/src/main/java/com/velocitypowered/proxy/command/VelocityCommandManager.java

@@ -6,7 +6,6 @@ import com.velocitypowered.api.command.CommandManager;
 import com.velocitypowered.api.command.CommandSource;
 
 import java.util.*;
-import java.util.stream.Collectors;
 
 public class VelocityCommandManager implements CommandManager {
     private final Map<String, Command> commands = new HashMap<>();
@@ -57,6 +56,10 @@ public class VelocityCommandManager implements CommandManager {
         }
     }
 
+    public boolean hasCommand(String command) {
+        return commands.containsKey(command);
+    }
+
     public Optional<List<String>> offerSuggestions(CommandSource source, String cmdLine) {
         Preconditions.checkNotNull(source, "source");
         Preconditions.checkNotNull(cmdLine, "cmdLine");

proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java

@@ -69,12 +69,19 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
         if (msg instanceof MinecraftPacket) {
+            if (sessionHandler.beforeHandle()) {
+                return;
+            }
+
             MinecraftPacket pkt = (MinecraftPacket) msg;
             if (!pkt.handle(sessionHandler)) {
                 sessionHandler.handleGeneric((MinecraftPacket) msg);
             }
         } else if (msg instanceof ByteBuf) {
             try {
+                if (sessionHandler.beforeHandle()) {
+                    return;
+                }
                 sessionHandler.handleUnknown((ByteBuf) msg);
             } finally {
                 ReferenceCountUtil.release(msg);

proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftSessionHandler.java

@@ -5,10 +5,16 @@ import com.velocitypowered.proxy.protocol.packet.*;
 import io.netty.buffer.ByteBuf;
 
 public interface MinecraftSessionHandler {
-    void handleGeneric(MinecraftPacket packet);
+    default boolean beforeHandle() {
+        return false;
+    }
+
+    default void handleGeneric(MinecraftPacket packet) {
+
+    }
 
     default void handleUnknown(ByteBuf buf) {
-        // No-op: we'll release the buffer later.
+
     }
 
     default void connected() {

proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java

@@ -17,10 +17,12 @@ import io.netty.buffer.ByteBuf;
 public class BackendPlaySessionHandler implements MinecraftSessionHandler {
     private final VelocityServer server;
     private final VelocityServerConnection serverConn;
+    private final ClientPlaySessionHandler playerSessionHandler;
 
     public BackendPlaySessionHandler(VelocityServer server, VelocityServerConnection serverConn) {
         this.server = server;
         this.serverConn = serverConn;
+        this.playerSessionHandler = (ClientPlaySessionHandler) serverConn.getPlayer().getConnection().getSessionHandler();
     }
 
     @Override
@@ -29,6 +31,93 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
         serverConn.getServer().addPlayer(serverConn.getPlayer());
     }
 
+    @Override
+    public boolean beforeHandle() {
+        if (!serverConn.getPlayer().isActive()) {
+            // Obsolete connection
+            serverConn.disconnect();
+            return true;
+        }
+        return false;
+    }
+
+    @Override
+    public boolean handle(KeepAlive packet) {
+        serverConn.setLastPingId(packet.getRandomId());
+        return false; // forwards on
+    }
+
+    @Override
+    public boolean handle(Disconnect packet) {
+        serverConn.disconnect();
+        serverConn.getPlayer().handleConnectionException(serverConn.getServer(), packet);
+        return true;
+    }
+
+    @Override
+    public boolean handle(JoinGame packet) {
+        playerSessionHandler.handleBackendJoinGame(packet);
+        return true;
+    }
+
+    @Override
+    public boolean handle(BossBar packet) {
+        switch (packet.getAction()) {
+            case 0: // add
+                playerSessionHandler.getServerBossBars().add(packet.getUuid());
+                break;
+            case 1: // remove
+                playerSessionHandler.getServerBossBars().remove(packet.getUuid());
+                break;
+        }
+        return false; // forward
+    }
+
+    @Override
+    public boolean handle(PluginMessage packet) {
+        if (!canForwardPluginMessage(packet)) {
+            return true;
+        }
+
+        if (PluginMessageUtil.isMCBrand(packet)) {
+            serverConn.getPlayer().getConnection().write(PluginMessageUtil.rewriteMCBrand(packet));
+            return true;
+        }
+
+        if (!serverConn.hasCompletedJoin() && packet.getChannel().equals(VelocityConstants.FORGE_LEGACY_HANDSHAKE_CHANNEL)) {
+            if (!serverConn.isLegacyForge()) {
+                serverConn.setLegacyForge(true);
+
+                // We must always reset the handshake before a modded connection is established if
+                // we haven't done so already.
+                serverConn.getPlayer().sendLegacyForgeHandshakeResetPacket();
+            }
+
+            // Always forward these messages during login
+            return false;
+        }
+
+        ChannelIdentifier id = server.getChannelRegistrar().getFromId(packet.getChannel());
+        if (id == null) {
+            serverConn.getPlayer().getConnection().write(packet);
+        } else {
+            PluginMessageEvent event = new PluginMessageEvent(serverConn, serverConn.getPlayer(), id, packet.getData());
+            server.getEventManager().fire(event)
+                    .thenAcceptAsync(pme -> {
+                        if (pme.getResult().isAllowed()) {
+                            serverConn.getPlayer().getConnection().write(packet);
+                        }
+                    }, serverConn.getConnection().eventLoop());
+        }
+        return true;
+    }
+
+    @Override
+    public boolean handle(TabCompleteResponse packet) {
+        playerSessionHandler.handleTabCompleteResponse(packet);
+        return true;
+    }
+
     @Override
     public void handleGeneric(MinecraftPacket packet) {
         if (!serverConn.getPlayer().isActive()) {
@@ -38,69 +127,7 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
             return;
         }
 
-        ClientPlaySessionHandler playerHandler =
+        if (serverConn.hasCompletedJoin()) {
-                (ClientPlaySessionHandler) serverConn.getPlayer().getConnection().getSessionHandler();
-        if (packet instanceof KeepAlive) {
-            // Forward onto the player
-            serverConn.setLastPingId(((KeepAlive) packet).getRandomId());
-            serverConn.getPlayer().getConnection().write(packet);
-        } else if (packet instanceof Disconnect) {
-            Disconnect original = (Disconnect) packet;
-            serverConn.disconnect();
-            serverConn.getPlayer().handleConnectionException(serverConn.getServer(), original);
-        } else if (packet instanceof JoinGame) {
-            playerHandler.handleBackendJoinGame((JoinGame) packet);
-        } else if (packet instanceof BossBar) {
-            BossBar bossBar = (BossBar) packet;
-            switch (bossBar.getAction()) {
-                case 0: // add
-                    playerHandler.getServerBossBars().add(bossBar.getUuid());
-                    break;
-                case 1: // remove
-                    playerHandler.getServerBossBars().remove(bossBar.getUuid());
-                    break;
-            }
-            serverConn.getPlayer().getConnection().write(packet);
-        } else if (packet instanceof PluginMessage) {
-            PluginMessage pm = (PluginMessage) packet;
-            if (!canForwardPluginMessage(pm)) {
-                return;
-            }
-
-            if (PluginMessageUtil.isMCBrand(pm)) {
-                serverConn.getPlayer().getConnection().write(PluginMessageUtil.rewriteMCBrand(pm));
-                return;
-            }
-
-            if (!serverConn.hasCompletedJoin() && pm.getChannel().equals(VelocityConstants.FORGE_LEGACY_HANDSHAKE_CHANNEL)) {
-                if (!serverConn.isLegacyForge()) {
-                    serverConn.setLegacyForge(true);
-
-                    // We must always reset the handshake before a modded connection is established if
-                    // we haven't done so already.
-                    serverConn.getPlayer().sendLegacyForgeHandshakeResetPacket();
-                }
-
-                // Always forward these messages during login
-                serverConn.getPlayer().getConnection().write(pm);
-                return;
-            }
-
-            ChannelIdentifier id = server.getChannelRegistrar().getFromId(pm.getChannel());
-            if (id == null) {
-                serverConn.getPlayer().getConnection().write(pm);
-            } else {
-                PluginMessageEvent event = new PluginMessageEvent(serverConn, serverConn.getPlayer(), id, pm.getData());
-                server.getEventManager().fire(event)
-                        .thenAcceptAsync(pme -> {
-                            if (pme.getResult().isAllowed()) {
-                                serverConn.getPlayer().getConnection().write(pm);
-                            }
-                        }, serverConn.getConnection().eventLoop());
-            }
-        } else if (packet instanceof TabCompleteResponse) {
-            playerHandler.handleTabCompleteResponse((TabCompleteResponse) packet);
-        } else if (serverConn.hasCompletedJoin()) {
             // Just forward the packet on. We don't have anything to handle at this time.
             serverConn.getPlayer().getConnection().write(packet);
         }

proxy/src/main/java/com/velocitypowered/proxy/connection/backend/LoginSessionHandler.java

@@ -36,17 +36,18 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
     }
 
     @Override
-    public void handleGeneric(MinecraftPacket packet) {
+    public boolean handle(EncryptionRequest packet) {
-        if (packet instanceof EncryptionRequest) {
         throw new IllegalStateException("Backend server is online-mode!");
-        } else if (packet instanceof LoginPluginMessage) {
+    }
-            LoginPluginMessage message = (LoginPluginMessage) packet;
+
+    @Override
+    public boolean handle(LoginPluginMessage packet) {
         VelocityConfiguration configuration = server.getConfiguration();
-            if (configuration.getPlayerInfoForwardingMode() == PlayerInfoForwarding.MODERN &&
+        if (configuration.getPlayerInfoForwardingMode() == PlayerInfoForwarding.MODERN && packet.getChannel()
-                    message.getChannel().equals(VelocityConstants.VELOCITY_IP_FORWARDING_CHANNEL)) {
+                .equals(VelocityConstants.VELOCITY_IP_FORWARDING_CHANNEL)) {
             LoginPluginResponse response = new LoginPluginResponse();
             response.setSuccess(true);
-                response.setId(message.getId());
+            response.setId(packet.getId());
             response.setData(createForwardingData(configuration.getForwardingSecret(),
                     serverConn.getPlayer().getRemoteAddress().getHostString(),
                     serverConn.getPlayer().getProfile()));
@@ -56,24 +57,35 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
             // Don't understand
             LoginPluginResponse response = new LoginPluginResponse();
             response.setSuccess(false);
-                response.setId(message.getId());
+            response.setId(packet.getId());
             response.setData(Unpooled.EMPTY_BUFFER);
             serverConn.getConnection().write(response);
         }
-        } else if (packet instanceof Disconnect) {
+        return true;
+    }
+
+    @Override
+    public boolean handle(Disconnect packet) {
         Disconnect disconnect = (Disconnect) packet;
         // Do we have an outstanding notification? If so, fulfill it.
         doNotify(ConnectionRequestResults.forDisconnect(disconnect));
         serverConn.disconnect();
-        } else if (packet instanceof SetCompression) {
+        return true;
-            SetCompression sc = (SetCompression) packet;
+    }
-            serverConn.getConnection().setCompressionThreshold(sc.getThreshold());
+
-        } else if (packet instanceof ServerLoginSuccess) {
+    @Override
+    public boolean handle(SetCompression packet) {
+        serverConn.getConnection().setCompressionThreshold(packet.getThreshold());
+        return true;
+    }
+
+    @Override
+    public boolean handle(ServerLoginSuccess packet) {
         if (server.getConfiguration().getPlayerInfoForwardingMode() == PlayerInfoForwarding.MODERN && !informationForwarded) {
             doNotify(ConnectionRequestResults.forDisconnect(
                     TextComponent.of("Your server did not send a forwarding request to the proxy. Is it set up correctly?")));
             serverConn.disconnect();
-                return;
+            return true;
         }
 
         // The player has been logged on to the backend server.
@@ -94,7 +106,7 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
         doNotify(ConnectionRequestResults.SUCCESSFUL);
         serverConn.getConnection().setSessionHandler(new BackendPlaySessionHandler(server, serverConn));
         serverConn.getPlayer().setConnectedServer(serverConn);
-        }
+        return true;
     }
 
     @Override

proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java

@@ -48,6 +48,128 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
         player.getConnection().write(register);
     }
 
+    @Override
+    public boolean handle(KeepAlive packet) {
+        VelocityServerConnection serverConnection = player.getConnectedServer();
+        if (serverConnection != null && packet.getRandomId() == serverConnection.getLastPingId()) {
+            player.setPing(System.currentTimeMillis() - serverConnection.getLastPingSent());
+            serverConnection.getConnection().write(packet);
+            serverConnection.resetLastPingId();
+        }
+        return true;
+    }
+
+    @Override
+    public boolean handle(ClientSettings packet) {
+        player.setPlayerSettings(packet);
+        return false; // will forward onto the handleGeneric below, which will write the packet to the remote server
+    }
+
+    @Override
+    public boolean handle(Chat packet) {
+        String msg = packet.getMessage();
+        if (msg.startsWith("/")) {
+            try {
+                if (!server.getCommandManager().execute(player, msg.substring(1))) {
+                    return false;
+                }
+            } catch (Exception e) {
+                logger.info("Exception occurred while running command for {}", player.getProfile().getName(), e);
+                player.sendMessage(TextComponent.of("An error occurred while running this command.", TextColor.RED));
+                return true;
+            }
+        } else {
+            VelocityServerConnection serverConnection = player.getConnectedServer();
+            if (serverConnection == null) {
+                return true;
+            }
+            PlayerChatEvent event = new PlayerChatEvent(player, msg);
+            server.getEventManager().fire(event)
+                    .thenAcceptAsync(pme -> {
+                        if (pme.getResult().equals(PlayerChatEvent.ChatResult.allowed())){
+                            serverConnection.getConnection().write(packet);
+                        } else if (pme.getResult().isAllowed() && pme.getResult().getMessage().isPresent()){
+                            serverConnection.getConnection().write(Chat.createServerbound(pme.getResult().getMessage().get()));
+                        }
+                    }, serverConnection.getConnection().eventLoop());
+        }
+        return true;
+    }
+
+    @Override
+    public boolean handle(TabCompleteRequest packet) {
+        // Record the request so that the outstanding request can be augmented later.
+        if (!packet.isAssumeCommand() && packet.getCommand().startsWith("/")) {
+            int spacePos = packet.getCommand().indexOf(' ');
+            if (spacePos > 0) {
+                String cmd = packet.getCommand().substring(1, spacePos);
+                if (server.getCommandManager().hasCommand(cmd)) {
+                    Optional<List<String>> suggestions = server.getCommandManager().offerSuggestions(player, packet.getCommand().substring(1));
+                    if (suggestions.isPresent()) {
+                        TabCompleteResponse resp = new TabCompleteResponse();
+                        resp.getOffers().addAll(suggestions.get());
+                        player.getConnection().write(resp);
+                        return true;
+                    }
+                }
+            }
+        }
+        outstandingTabComplete = packet;
+        return false;
+    }
+
+    @Override
+    public boolean handle(PluginMessage packet) {
+        if (PluginMessageUtil.isMCRegister(packet)) {
+            List<String> actuallyRegistered = new ArrayList<>();
+            List<String> channels = PluginMessageUtil.getChannels(packet);
+            for (String channel : channels) {
+                if (clientPluginMsgChannels.size() >= MAX_PLUGIN_CHANNELS &&
+                        !clientPluginMsgChannels.contains(channel)) {
+                    throw new IllegalStateException("Too many plugin message channels registered");
+                }
+                if (clientPluginMsgChannels.add(channel)) {
+                    actuallyRegistered.add(channel);
+                }
+            }
+
+            if (actuallyRegistered.size() > 0) {
+                PluginMessage newRegisterPacket = PluginMessageUtil.constructChannelsPacket(player.getProtocolVersion(), actuallyRegistered);
+                player.getConnectedServer().getConnection().write(newRegisterPacket);
+            }
+        } else if (PluginMessageUtil.isMCUnregister(packet)) {
+            List<String> channels = PluginMessageUtil.getChannels(packet);
+            clientPluginMsgChannels.removeAll(channels);
+            player.getConnectedServer().getConnection().write(packet);
+        } else if (PluginMessageUtil.isMCBrand(packet)) {
+            player.getConnectedServer().getConnection().write(PluginMessageUtil.rewriteMCBrand(packet));
+        } else if (player.getConnectedServer().isLegacyForge() && !player.getConnectedServer().hasCompletedJoin()) {
+            if (packet.getChannel().equals(VelocityConstants.FORGE_LEGACY_HANDSHAKE_CHANNEL)) {
+                // Always forward the FML handshake to the remote server.
+                player.getConnectedServer().getConnection().write(packet);
+            } else {
+                // The client is trying to send messages too early. This is primarily caused by mods, but it's further
+                // aggravated by Velocity. To work around these issues, we will queue any non-FML handshake messages to
+                // be sent once the JoinGame packet has been received by the proxy.
+                loginPluginMessages.add(packet);
+            }
+        } else {
+            ChannelIdentifier id = server.getChannelRegistrar().getFromId(packet.getChannel());
+            if (id == null) {
+                player.getConnectedServer().getConnection().write(packet);
+            } else {
+                PluginMessageEvent event = new PluginMessageEvent(player, player.getConnectedServer(), id, packet.getData());
+                server.getEventManager().fire(event)
+                        .thenAcceptAsync(pme -> {
+                            if (pme.getResult().isAllowed()) {
+                                player.getConnectedServer().getConnection().write(packet);
+                            }
+                        }, player.getConnectedServer().getConnection().eventLoop());
+            }
+        }
+        return true;
+    }
+
     @Override
     public void handleGeneric(MinecraftPacket packet) {
         VelocityServerConnection serverConnection = player.getConnectedServer();
@@ -56,65 +178,9 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
             return;
         }
 
-        if (packet instanceof KeepAlive) {
-            KeepAlive keepAlive = (KeepAlive) packet;
-            if (keepAlive.getRandomId() != serverConnection.getLastPingId()) {
-                // The last keep alive we got was probably from a different server. Let's ignore it, and hope the next
-                // ping is alright.
-                return;
-            }
-            player.setPing(System.currentTimeMillis() - serverConnection.getLastPingSent());
-            serverConnection.getConnection().write(packet);
-            serverConnection.resetLastPingId();
-            return;
-        }
-
-        if (packet instanceof ClientSettings) {
-            player.setPlayerSettings((ClientSettings) packet);
-            // forward it on
-        }
-
-        if (packet instanceof Chat) {
-            // Try to handle any commands on the proxy. If that fails, send it onto the client.
-            Chat chat = (Chat) packet;
-            String msg = ((Chat) packet).getMessage();
-            if (msg.startsWith("/")) {
-                try {
-                    if (!server.getCommandManager().execute(player, msg.substring(1))) {
-                        player.getConnectedServer().getConnection().write(chat);
-                    }
-                } catch (Exception e) {
-                    logger.info("Exception occurred while running command for {}", player.getProfile().getName(), e);
-                    player.sendMessage(TextComponent.of("An error occurred while running this command.", TextColor.RED));
-                    return;
-                }
-            } else {
-                PlayerChatEvent event = new PlayerChatEvent(player, msg);
-                server.getEventManager().fire(event)
-                        .thenAcceptAsync(pme -> {
-                            if (pme.getResult().equals(PlayerChatEvent.ChatResult.allowed())){
-                                serverConnection.getConnection().write(chat);
-                            } else if (pme.getResult().isAllowed() && pme.getResult().getMessage().isPresent()){
-                                serverConnection.getConnection().write(Chat.createServerbound(pme.getResult().getMessage().get()));
-                            }
-                        }, serverConnection.getConnection().eventLoop());
-            }
-            return;
-        }
-
-        if (packet instanceof TabCompleteRequest) {
-            // Record the request so that the outstanding request can be augmented later.
-            outstandingTabComplete = (TabCompleteRequest) packet;
-            serverConnection.getConnection().write(packet);
-        }
-
-        if (packet instanceof PluginMessage) {
-            handleClientPluginMessage((PluginMessage) packet);
-            return;
-        }
-
         // If we don't want to handle this packet, just forward it on.
         if (serverConnection.hasCompletedJoin()) {
+            logger.info("Will write {}", packet);
             serverConnection.getConnection().write(packet);
         }
     }
@@ -245,62 +311,15 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
         return serverBossBars;
     }
 
-    private void handleClientPluginMessage(PluginMessage packet) {
-        if (PluginMessageUtil.isMCRegister(packet)) {
-            List<String> actuallyRegistered = new ArrayList<>();
-            List<String> channels = PluginMessageUtil.getChannels(packet);
-            for (String channel : channels) {
-                if (clientPluginMsgChannels.size() >= MAX_PLUGIN_CHANNELS &&
-                        !clientPluginMsgChannels.contains(channel)) {
-                    throw new IllegalStateException("Too many plugin message channels registered");
-                }
-                if (clientPluginMsgChannels.add(channel)) {
-                    actuallyRegistered.add(channel);
-                }
-            }
-
-            if (actuallyRegistered.size() > 0) {
-                PluginMessage newRegisterPacket = PluginMessageUtil.constructChannelsPacket(player.getProtocolVersion(), actuallyRegistered);
-                player.getConnectedServer().getConnection().write(newRegisterPacket);
-            }
-        } else if (PluginMessageUtil.isMCUnregister(packet)) {
-            List<String> channels = PluginMessageUtil.getChannels(packet);
-            clientPluginMsgChannels.removeAll(channels);
-            player.getConnectedServer().getConnection().write(packet);
-        } else if (PluginMessageUtil.isMCBrand(packet)) {
-            player.getConnectedServer().getConnection().write(PluginMessageUtil.rewriteMCBrand(packet));
-        } else if (player.getConnectedServer().isLegacyForge() && !player.getConnectedServer().hasCompletedJoin()) {
-            if (packet.getChannel().equals(VelocityConstants.FORGE_LEGACY_HANDSHAKE_CHANNEL)) {
-                // Always forward the FML handshake to the remote server.
-                player.getConnectedServer().getConnection().write(packet);
-            } else {
-                // The client is trying to send messages too early. This is primarily caused by mods, but it's further
-                // aggravated by Velocity. To work around these issues, we will queue any non-FML handshake messages to
-                // be sent once the JoinGame packet has been received by the proxy.
-                loginPluginMessages.add(packet);
-            }
-        } else {
-            ChannelIdentifier id = server.getChannelRegistrar().getFromId(packet.getChannel());
-            if (id == null) {
-                player.getConnectedServer().getConnection().write(packet);
-            } else {
-                PluginMessageEvent event = new PluginMessageEvent(player, player.getConnectedServer(), id, packet.getData());
-                server.getEventManager().fire(event)
-                        .thenAcceptAsync(pme -> {
-                            if (pme.getResult().isAllowed()) {
-                                player.getConnectedServer().getConnection().write(packet);
-                            }
-                        }, player.getConnectedServer().getConnection().eventLoop());
-            }
-        }
-    }
-
     public Set<String> getClientPluginMsgChannels() {
         return clientPluginMsgChannels;
     }
 
     public void handleTabCompleteResponse(TabCompleteResponse response) {
+        logger.info("Got {}", response);
+        logger.info("Request {}", outstandingTabComplete);
         if (outstandingTabComplete != null) {
+            logger.info("HANDLING");
             if (!outstandingTabComplete.isAssumeCommand()) {
                 String command = outstandingTabComplete.getCommand().substring(1);
                 try {

proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java

@@ -117,7 +117,8 @@ public class HandshakeSessionHandler implements MinecraftSessionHandler {
 
     @Override
     public void handleUnknown(ByteBuf buf) {
-        throw new IllegalStateException("Unknown data " + ByteBufUtil.hexDump(buf));
+        // what even is going on?
+        connection.close();
     }
 
     private static class LegacyInboundConnection implements InboundConnection {

proxy/src/main/java/com/velocitypowered/proxy/connection/client/LoginSessionHandler.java

@@ -59,23 +59,8 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
     }
 
     @Override
-    public void handleGeneric(MinecraftPacket packet) {
+    public boolean handle(ServerLogin packet) {
-        if (packet instanceof LoginPluginResponse) {
+        this.login = packet;
-            LoginPluginResponse lpr = (LoginPluginResponse) packet;
-            if (lpr.getId() == playerInfoId) {
-                if (lpr.isSuccess()) {
-                    // Uh oh, someone's trying to run Velocity behind Velocity. We don't want that happening.
-                    inbound.closeWith(Disconnect.create(
-                            TextComponent.of("Running Velocity behind Velocity isn't supported.", TextColor.RED)
-                    ));
-                } else {
-                    // Proceed with the regular login process.
-                    beginPreLogin();
-                }
-            }
-        } else if (packet instanceof ServerLogin) {
-            this.login = (ServerLogin) packet;
-
         if (inbound.getProtocolVersion() >= ProtocolConstants.MINECRAFT_1_13) {
             LoginPluginMessage message = new LoginPluginMessage();
             playerInfoId = ThreadLocalRandom.current().nextInt();
@@ -86,7 +71,27 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
         } else {
             beginPreLogin();
         }
-        } else if (packet instanceof EncryptionResponse) {
+        return true;
+    }
+
+    @Override
+    public boolean handle(LoginPluginResponse packet) {
+        if (packet.getId() == playerInfoId) {
+            if (packet.isSuccess()) {
+                // Uh oh, someone's trying to run Velocity behind Velocity. We don't want that happening.
+                inbound.closeWith(Disconnect.create(
+                        TextComponent.of("Running Velocity behind Velocity isn't supported.", TextColor.RED)
+                ));
+            } else {
+                // Proceed with the regular login process.
+                beginPreLogin();
+            }
+        }
+        return true;
+    }
+
+    @Override
+    public boolean handle(EncryptionResponse packet) {
         try {
             KeyPair serverKeyPair = server.getServerKeyPair();
             EncryptionResponse response = (EncryptionResponse) packet;
@@ -141,7 +146,7 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
         } catch (MalformedURLException e) {
             throw new AssertionError(e);
         }
-        }
+        return true;
     }
 
     private void beginPreLogin() {

proxy/src/main/java/com/velocitypowered/proxy/connection/client/StatusSessionHandler.java

@@ -1,6 +1,5 @@
 package com.velocitypowered.proxy.connection.client;
 
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.velocitypowered.api.event.proxy.ProxyPingEvent;
 import com.velocitypowered.api.proxy.InboundConnection;
@@ -9,13 +8,11 @@ import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.config.VelocityConfiguration;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
-import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.ProtocolConstants;
 import com.velocitypowered.proxy.protocol.packet.StatusPing;
 import com.velocitypowered.proxy.protocol.packet.StatusRequest;
 import com.velocitypowered.proxy.protocol.packet.StatusResponse;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.ByteBufUtil;
 
 public class StatusSessionHandler implements MinecraftSessionHandler {
     private final VelocityServer server;
@@ -29,19 +26,15 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
     }
 
     @Override
-    public void handleGeneric(MinecraftPacket packet) {
+    public boolean handle(StatusPing packet) {
-        Preconditions.checkArgument(packet instanceof StatusPing || packet instanceof StatusRequest,
-                "Unrecognized packet type " + packet.getClass().getName());
-
-        if (packet instanceof StatusPing) {
-            // Just send back the client's packet, no processing to do here.
         connection.closeWith(packet);
-            return;
+        return true;
     }
 
+    @Override
+    public boolean handle(StatusRequest packet) {
         VelocityConfiguration configuration = server.getConfiguration();
 
-        // Status request
         int shownVersion = ProtocolConstants.isSupported(connection.getProtocolVersion()) ? connection.getProtocolVersion() :
                 ProtocolConstants.MAXIMUM_GENERIC_VERSION;
         ServerPing initialPing = new ServerPing(
@@ -59,10 +52,12 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
                     response.setStatus(VelocityServer.GSON.toJson(event.getPing()));
                     connection.write(response);
                 }, connection.eventLoop());
+        return true;
     }
 
     @Override
     public void handleUnknown(ByteBuf buf) {
-        throw new IllegalStateException("Unknown data " + ByteBufUtil.hexDump(buf));
+        // what even is going on?
+        connection.close();
     }
 }

proxy/src/main/java/com/velocitypowered/proxy/server/ping/PingSessionHandler.java

@@ -1,12 +1,10 @@
 package com.velocitypowered.proxy.server.ping;
 
-import com.google.common.base.Preconditions;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.api.proxy.server.ServerPing;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
-import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.ProtocolConstants;
 import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.packet.Handshake;
@@ -42,15 +40,14 @@ public class PingSessionHandler implements MinecraftSessionHandler {
     }
 
     @Override
-    public void handleGeneric(MinecraftPacket packet) {
+    public boolean handle(StatusResponse packet) {
-        Preconditions.checkState(packet instanceof StatusResponse, "Did not get status response back from connection");
-
         // All good!
         completed = true;
         connection.close();
 
         ServerPing ping = VelocityServer.GSON.fromJson(((StatusResponse) packet).getStatus(), ServerPing.class);
         result.complete(ping);
+        return true;
     }
 
     @Override