Refactored connection handling

src/main/java/io/minimum/minecraft/velocity/Velocity.java

@@ -2,6 +2,7 @@ package io.minimum.minecraft.velocity;
 
 import io.minimum.minecraft.velocity.protocol.ProtocolConstants;
 import io.minimum.minecraft.velocity.protocol.netty.*;
+import io.minimum.minecraft.velocity.proxy.InboundMinecraftConnection;
 import io.minimum.minecraft.velocity.proxy.MinecraftClientSessionHandler;
 import io.netty.bootstrap.ServerBootstrap;
 import io.netty.channel.Channel;
@@ -17,6 +18,7 @@ public class Velocity {
                 .childHandler(new ChannelInitializer<Channel>() {
                     @Override
                     protected void initChannel(Channel ch) throws Exception {
+                        ch.attr(InboundMinecraftConnection.CONNECTION).set(new InboundMinecraftConnection(ch));
                         ch.pipeline().addLast("legacy-ping-decode", new LegacyPingDecoder());
                         ch.pipeline().addLast("frame-decoder", new MinecraftVarintFrameDecoder());
                         ch.pipeline().addLast("legacy-ping-encode", new LegacyPingEncoder());

src/main/java/io/minimum/minecraft/velocity/proxy/InboundMinecraftConnection.java

@@ -0,0 +1,83 @@
+package io.minimum.minecraft.velocity.proxy;
+
+import com.google.common.base.Preconditions;
+import io.minimum.minecraft.velocity.protocol.StateRegistry;
+import io.minimum.minecraft.velocity.protocol.netty.MinecraftDecoder;
+import io.minimum.minecraft.velocity.protocol.netty.MinecraftEncoder;
+import io.minimum.minecraft.velocity.protocol.packets.Handshake;
+import io.minimum.minecraft.velocity.proxy.handler.HandshakeSessionHandler;
+import io.minimum.minecraft.velocity.proxy.handler.StatusSessionHandler;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelFutureListener;
+import io.netty.util.AttributeKey;
+
+public class InboundMinecraftConnection {
+    public static final AttributeKey<InboundMinecraftConnection> CONNECTION = AttributeKey.newInstance("velocity-connection");
+
+    private final Channel channel;
+    private boolean closed;
+    private Handshake handshake;
+    private StateRegistry state;
+    private MinecraftSessionHandler sessionHandler;
+
+    public InboundMinecraftConnection(Channel channel) {
+        this.channel = channel;
+        this.closed = false;
+        this.state = StateRegistry.HANDSHAKE;
+        this.sessionHandler = new HandshakeSessionHandler(this);
+    }
+
+    public void write(Object msg) {
+        ensureOpen();
+        channel.writeAndFlush(msg, channel.voidPromise());
+    }
+
+    public void closeWith(Object msg) {
+        ensureOpen();
+        closed = true;
+        channel.writeAndFlush(msg).addListener(new ChannelFutureListener() {
+            @Override
+            public void operationComplete(ChannelFuture future) throws Exception {
+                future.channel().close();
+            }
+        });
+    }
+
+    public void close() {
+        ensureOpen();
+        channel.close();
+        closed = true;
+    }
+
+    public MinecraftSessionHandler getSessionHandler() {
+        return sessionHandler;
+    }
+
+    public void handleHandshake(Handshake handshake) {
+        ensureOpen();
+        Preconditions.checkNotNull(handshake, "handshake");
+        Preconditions.checkState(this.handshake == null, "Already handled a handshake for this connection!");
+        this.handshake = handshake;
+        switch (handshake.getNextStatus()) {
+            case 1:
+                // Status protocol
+                this.setStatus(StateRegistry.STATUS);
+                this.sessionHandler = new StatusSessionHandler(this);
+                break;
+            default:
+                throw new UnsupportedOperationException("Unsupported next protocol state " + handshake.getNextStatus());
+        }
+    }
+
+    private void ensureOpen() {
+        Preconditions.checkState(!closed, "Connection is closed.");
+    }
+
+    private void setStatus(StateRegistry state) {
+        Preconditions.checkNotNull(state, "state");
+        this.state = state;
+        channel.pipeline().get(MinecraftEncoder.class).setState(state);
+        channel.pipeline().get(MinecraftDecoder.class).setState(state);
+    }
+}

src/main/java/io/minimum/minecraft/velocity/proxy/MinecraftClientSessionHandler.java

@@ -24,15 +24,22 @@ public class MinecraftClientSessionHandler extends ChannelInboundHandlerAdapter
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+        InboundMinecraftConnection connection = ctx.channel().attr(InboundMinecraftConnection.CONNECTION).get();
         if (msg instanceof PacketWrapper) {
+            PacketWrapper pw = (PacketWrapper) msg;
             try {
-                handle(ctx, (PacketWrapper) msg);
+                if (pw.getPacket() == null) {
+                    connection.getSessionHandler().handleUnknown(pw.getBuffer());
+                } else {
+                    connection.getSessionHandler().handle(pw.getPacket());
+                }
             } finally {
                 ((PacketWrapper) msg).getBuffer().release();
             }
         }
 
         if (msg instanceof LegacyPing) {
+            // TODO: port this
             System.out.println("Got LEGACY status request!");
             ServerPing ping = new ServerPing(
                     new ServerPing.Version(340, "1.12"),
@@ -50,44 +57,4 @@ public class MinecraftClientSessionHandler extends ChannelInboundHandlerAdapter
         cause.printStackTrace();
         ctx.close();
     }
-
-    private void handle(ChannelHandlerContext ctx, PacketWrapper msg) {
-        MinecraftPacket packet = msg.getPacket();
-        if (packet == null) {
-            System.out.println("no packet!");
-            return;
-        }
-
-        if (packet instanceof Handshake) {
-            System.out.println("Handshake: " + packet);
-            switch (((Handshake) packet).getNextStatus()) {
-                case 1:
-                    // status
-                    ctx.pipeline().get(MinecraftDecoder.class).setState(StateRegistry.STATUS);
-                    ctx.pipeline().get(MinecraftEncoder.class).setState(StateRegistry.STATUS);
-                    break;
-                case 2:
-                    // login
-                    throw new UnsupportedOperationException("Login not supported yet");
-            }
-        }
-
-        if (packet instanceof StatusRequest) {
-            System.out.println("Got status request!");
-            ServerPing ping = new ServerPing(
-                    new ServerPing.Version(340, "1.12.2"),
-                    new ServerPing.Players(0, 0),
-                    TextComponent.of("test"),
-                    null
-            );
-            StatusResponse response = new StatusResponse();
-            response.setStatus(GSON.toJson(ping));
-            ctx.writeAndFlush(response, ctx.voidPromise());
-        }
-
-        if (packet instanceof Ping) {
-            System.out.println("Ping: " + packet);
-            ctx.writeAndFlush(packet).addListener(ChannelFutureListener.CLOSE);
-        }
-    }
 }

src/main/java/io/minimum/minecraft/velocity/proxy/MinecraftSessionHandler.java

@@ -0,0 +1,12 @@
+package io.minimum.minecraft.velocity.proxy;
+
+import io.minimum.minecraft.velocity.protocol.MinecraftPacket;
+import io.netty.buffer.ByteBuf;
+
+public interface MinecraftSessionHandler {
+    void handle(MinecraftPacket packet);
+
+    default void handleUnknown(ByteBuf buf) {
+        // No-op: we'll release the buffer later.
+    }
+}

src/main/java/io/minimum/minecraft/velocity/proxy/handler/HandshakeSessionHandler.java

@@ -0,0 +1,25 @@
+package io.minimum.minecraft.velocity.proxy.handler;
+
+import com.google.common.base.Preconditions;
+import io.minimum.minecraft.velocity.protocol.MinecraftPacket;
+import io.minimum.minecraft.velocity.protocol.packets.Handshake;
+import io.minimum.minecraft.velocity.proxy.InboundMinecraftConnection;
+import io.minimum.minecraft.velocity.proxy.MinecraftSessionHandler;
+
+public class HandshakeSessionHandler implements MinecraftSessionHandler {
+    private final InboundMinecraftConnection connection;
+
+    public HandshakeSessionHandler(InboundMinecraftConnection connection) {
+        this.connection = Preconditions.checkNotNull(connection, "connection");
+    }
+
+    @Override
+    public void handle(MinecraftPacket packet) {
+        if (!(packet instanceof Handshake)) {
+            throw new IllegalArgumentException("Did not expect packet " + packet.getClass().getName());
+        }
+
+        Handshake handshake = (Handshake) packet;
+        connection.handleHandshake(handshake);
+    }
+}

src/main/java/io/minimum/minecraft/velocity/proxy/handler/StatusSessionHandler.java

@@ -0,0 +1,51 @@
+package io.minimum.minecraft.velocity.proxy.handler;
+
+import com.google.common.base.Preconditions;
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import io.minimum.minecraft.velocity.data.ServerPing;
+import io.minimum.minecraft.velocity.protocol.MinecraftPacket;
+import io.minimum.minecraft.velocity.protocol.packets.LegacyPing;
+import io.minimum.minecraft.velocity.protocol.packets.Ping;
+import io.minimum.minecraft.velocity.protocol.packets.StatusRequest;
+import io.minimum.minecraft.velocity.protocol.packets.StatusResponse;
+import io.minimum.minecraft.velocity.proxy.InboundMinecraftConnection;
+import io.minimum.minecraft.velocity.proxy.MinecraftSessionHandler;
+import net.kyori.text.Component;
+import net.kyori.text.TextComponent;
+import net.kyori.text.serializer.GsonComponentSerializer;
+
+public class StatusSessionHandler implements MinecraftSessionHandler {
+    private static final Gson GSON = new GsonBuilder()
+            .registerTypeHierarchyAdapter(Component.class, new GsonComponentSerializer())
+            .create();
+    private final InboundMinecraftConnection connection;
+
+    public StatusSessionHandler(InboundMinecraftConnection connection) {
+        this.connection = connection;
+    }
+
+    @Override
+    public void handle(MinecraftPacket packet) {
+        Preconditions.checkArgument(packet instanceof Ping || packet instanceof StatusRequest,
+                "Unrecognized packet type " + packet.getClass().getName());
+
+        if (packet instanceof Ping) {
+            // Just send back the client's packet, no processing to do here.
+            connection.closeWith(packet);
+            return;
+        }
+
+        // Status request
+        System.out.println("Got status request!");
+        ServerPing ping = new ServerPing(
+                new ServerPing.Version(340, "1.12.2"),
+                new ServerPing.Players(0, 0),
+                TextComponent.of("test"),
+                null
+        );
+        StatusResponse response = new StatusResponse();
+        response.setStatus(GSON.toJson(ping));
+        connection.write(response);
+    }
+}