Split incoming/outgoing packet registry, transition protocol states correctly

example/src/main/java/org/geysermc/mcprotocollib/network/example/TestProtocol.java

@@ -82,7 +82,12 @@ public void newServerSession(Server server, Session session) {
     }
 
     @Override
-    public PacketRegistry getPacketRegistry() {
+    public PacketRegistry getInboundPacketRegistry() {
+        return registry;
+    }
+
+    @Override
+    public PacketRegistry getOutboundPacketRegistry() {
         return registry;
     }
 }

example/src/main/java/org/geysermc/mcprotocollib/protocol/example/MinecraftProtocolTest.java

@@ -133,7 +133,7 @@ public void packetReceived(Session session, Packet packet) {
                 @Override
                 public void sessionRemoved(SessionRemovedEvent event) {
                     MinecraftProtocol protocol = (MinecraftProtocol) event.getSession().getPacketProtocol();
-                    if (protocol.getState() == ProtocolState.GAME) {
+                    if (protocol.getOutboundState() == ProtocolState.GAME) {
                         log.info("Closing server.");
                         event.getServer().close(false);
                     }

protocol/src/main/java/org/geysermc/mcprotocollib/network/Session.java

@@ -258,7 +258,17 @@ public interface Session {
      *
      * @param packet Packet to send.
      */
-    void send(Packet packet);
+    default void send(Packet packet) {
+        this.send(packet, null);
+    }
+
+    /**
+     * Sends a packet and runs the specified callback when the packet has been sent.
+     *
+     * @param packet Packet to send.
+     * @param onSent Callback to run when the packet has been sent.
+     */
+    void send(Packet packet, @Nullable Runnable onSent);
 
     /**
      * Disconnects the session.
@@ -301,4 +311,29 @@ default void disconnect(@NonNull Component reason) {
      * @param cause Throwable responsible for disconnecting.
      */
     void disconnect(@NonNull Component reason, @Nullable Throwable cause);
+
+    /**
+     * Auto read in netty means that the server is automatically reading from the channel.
+     * Turning it off means that we won't get more packets being decoded until we turn it back on.
+     * We use this to hold off on reading packets until we are ready to process them.
+     * For example this is used for switching inbound states with {@link #switchInboundState(Runnable)}.
+     *
+     * @param autoRead Whether to enable auto read or not.
+     *                 Default is true.
+     */
+    void setAutoRead(boolean autoRead);
+
+    /**
+     * Changes the inbound state of the session and then re-enables auto read.
+     * This is used after a terminal packet was handled and the session is ready to receive more packets in the new state.
+     *
+     * @param switcher The runnable that switches the inbound state.
+     */
+    default void switchInboundState(Runnable switcher) {
+        switcher.run();
+
+        // We switched to the new inbound state
+        // we can start reading again
+        setAutoRead(true);
+    }
 }

protocol/src/main/java/org/geysermc/mcprotocollib/network/packet/Packet.java

@@ -1,6 +1,7 @@
 package org.geysermc.mcprotocollib.network.packet;
 
 import io.netty.buffer.ByteBuf;
+import org.geysermc.mcprotocollib.network.Session;
 
 /**
  * A network packet. Any given packet must have a constructor that takes in a {@link ByteBuf}.
@@ -17,4 +18,14 @@ public interface Packet {
     default boolean isPriority() {
         return false;
     }
+
+    /**
+     * Returns whether the packet is terminal. If true, this should be the last packet sent inside a protocol state.
+     * Subsequently, {@link Session#setAutoRead(boolean)} should be disabled when a terminal packet is received, until the session has switched into a new state and is ready to receive more packets.
+     *
+     * @return Whether the packet is terminal.
+     */
+    default boolean isTerminal() {
+        return false;
+    }
 }

protocol/src/main/java/org/geysermc/mcprotocollib/network/packet/PacketProtocol.java

@@ -49,9 +49,16 @@ public abstract class PacketProtocol {
     public abstract void newServerSession(Server server, Session session);
 
     /**
-     * Gets the packet registry for this protocol.
+     * Gets the inbound packet registry for this protocol.
      *
-     * @return The protocol's packet registry.
+     * @return The protocol's inbound packet registry.
      */
-    public abstract PacketRegistry getPacketRegistry();
+    public abstract PacketRegistry getInboundPacketRegistry();
+
+    /**
+     * Gets the outbound packet registry for this protocol.
+     *
+     * @return The protocol's outbound packet registry.
+     */
+    public abstract PacketRegistry getOutboundPacketRegistry();
 }

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpClientSession.java

@@ -117,6 +117,7 @@ public void initChannel(Channel channel) {
                                 pipeline.addLast("sizer", new TcpPacketSizer(TcpClientSession.this, size));
                             }
 
+                            pipeline.addLast("flow-control", new TcpFlowControlHandler());
                             pipeline.addLast("codec", new TcpPacketCodec(TcpClientSession.this, true));
                             pipeline.addLast("manager", TcpClientSession.this);
 

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpFlowControlHandler.java

@@ -0,0 +1,20 @@
+package org.geysermc.mcprotocollib.network.tcp;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.flow.FlowControlHandler;
+
+/**
+ * A flow control handler for TCP connections.
+ * When auto-read is disabled, this will halt decoding of packets until auto-read is re-enabled.
+ * This is needed because auto-read still allows packets to be decoded, even if the channel is not reading anymore from the network.
+ * This can happen when the channel already read a packet, but the packet is not yet decoded.
+ * This will halt all decoding until the channel is ready to process more packets.
+ */
+public class TcpFlowControlHandler extends FlowControlHandler {
+    @Override
+    public void read(ChannelHandlerContext ctx) throws Exception {
+        if (ctx.channel().config().isAutoRead()) {
+            super.read(ctx);
+        }
+    }
+}

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpPacketCodec.java

@@ -9,10 +9,15 @@
 import org.geysermc.mcprotocollib.network.event.session.PacketErrorEvent;
 import org.geysermc.mcprotocollib.network.packet.Packet;
 import org.geysermc.mcprotocollib.network.packet.PacketProtocol;
+import org.geysermc.mcprotocollib.network.packet.PacketRegistry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.List;
 
 public class TcpPacketCodec extends ByteToMessageCodec<Packet> {
+    private static final Logger log = LoggerFactory.getLogger(TcpPacketCodec.class);
+
     private final Session session;
     private final boolean client;
 
@@ -24,16 +29,23 @@ public TcpPacketCodec(Session session, boolean client) {
     @SuppressWarnings({"rawtypes", "unchecked"})
     @Override
     public void encode(ChannelHandlerContext ctx, Packet packet, ByteBuf buf) {
+        if (log.isDebugEnabled()) {
+            log.debug("Encoding packet: {}", packet.getClass().getSimpleName());
+        }
+
         int initial = buf.writerIndex();
 
         PacketProtocol packetProtocol = this.session.getPacketProtocol();
+        PacketRegistry packetRegistry = packetProtocol.getOutboundPacketRegistry();
         PacketCodecHelper codecHelper = this.session.getCodecHelper();
         try {
-            int packetId = this.client ? packetProtocol.getPacketRegistry().getServerboundId(packet) : packetProtocol.getPacketRegistry().getClientboundId(packet);
-            PacketDefinition definition = this.client ? packetProtocol.getPacketRegistry().getServerboundDefinition(packetId) : packetProtocol.getPacketRegistry().getClientboundDefinition(packetId);
+            int packetId = this.client ? packetRegistry.getServerboundId(packet) : packetRegistry.getClientboundId(packet);
+            PacketDefinition definition = this.client ? packetRegistry.getServerboundDefinition(packetId) : packetRegistry.getClientboundDefinition(packetId);
 
             packetProtocol.getPacketHeader().writePacketId(buf, codecHelper, packetId);
             definition.getSerializer().serialize(buf, codecHelper, packet);
+
+            log.debug("Encoded packet with id: {}", packetId);
         } catch (Throwable t) {
             // Reset writer index to make sure incomplete data is not written out.
             buf.writerIndex(initial);
@@ -51,6 +63,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out)
         int initial = buf.readerIndex();
 
         PacketProtocol packetProtocol = this.session.getPacketProtocol();
+        PacketRegistry packetRegistry = packetProtocol.getInboundPacketRegistry();
         PacketCodecHelper codecHelper = this.session.getCodecHelper();
         try {
             int id = packetProtocol.getPacketHeader().readPacketId(buf, codecHelper);
@@ -59,13 +72,19 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out)
                 return;
             }
 
-            Packet packet = this.client ? packetProtocol.getPacketRegistry().createClientboundPacket(id, buf, codecHelper) : packetProtocol.getPacketRegistry().createServerboundPacket(id, buf, codecHelper);
+            log.debug("Decoding packet with id: {}", id);
+
+            Packet packet = this.client ? packetRegistry.createClientboundPacket(id, buf, codecHelper) : packetRegistry.createServerboundPacket(id, buf, codecHelper);
 
             if (buf.readableBytes() > 0) {
                 throw new IllegalStateException("Packet \"" + packet.getClass().getSimpleName() + "\" not fully read.");
             }
 
             out.add(packet);
+
+            if (log.isDebugEnabled()) {
+                log.debug("Decoded packet: {}", packet.getClass().getSimpleName());
+            }
         } catch (Throwable t) {
             // Advance buffer to end to make sure remaining data in this packet is skipped.
             buf.readerIndex(buf.readerIndex() + buf.readableBytes());

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpServer.java

@@ -68,6 +68,7 @@ public void initChannel(Channel channel) {
                     pipeline.addLast("sizer", new TcpPacketSizer(session, size));
                 }
 
+                pipeline.addLast("flow-control", new TcpFlowControlHandler());
                 pipeline.addLast("codec", new TcpPacketCodec(session, false));
                 pipeline.addLast("manager", session);
             }

protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpSession.java

@@ -256,7 +256,7 @@ public boolean isConnected() {
     }
 
     @Override
-    public void send(Packet packet) {
+    public void send(Packet packet, @Nullable Runnable onSent) {
         if (this.channel == null) {
             return;
         }
@@ -268,6 +268,10 @@ public void send(Packet packet) {
             final Packet toSend = sendingEvent.getPacket();
             this.channel.writeAndFlush(toSend).addListener((ChannelFutureListener) future -> {
                 if (future.isSuccess()) {
+                    if (onSent != null) {
+                        onSent.run();
+                    }
+
                     callPacketSent(toSend);
                 } else {
                     exceptionCaught(null, future.cause());
@@ -292,6 +296,13 @@ public void disconnect(@NonNull Component reason, @Nullable Throwable cause) {
         }
     }
 
+    @Override
+    public void setAutoRead(boolean autoRead) {
+        if (this.channel != null) {
+            this.channel.config().setAutoRead(autoRead);
+        }
+    }
+
     private @Nullable EventLoop createEventLoop() {
         if (!USE_EVENT_LOOP_FOR_PACKETS) {
             return null;
@@ -384,6 +395,12 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
 
     @Override
     protected void channelRead0(ChannelHandlerContext ctx, Packet packet) {
+        if (packet.isTerminal()) {
+            // Next packets are in a different protocol state, so we must
+            // disable auto-read to prevent reading wrong packets.
+            setAutoRead(false);
+        }
+
         if (!packet.isPriority() && eventLoop != null) {
             eventLoop.execute(() -> this.callPacketReceived(packet));
         } else {

protocol/src/main/java/org/geysermc/mcprotocollib/protocol/ClientListener.java

@@ -58,7 +58,7 @@ public class ClientListener extends SessionAdapter {
     @Override
     public void packetReceived(Session session, Packet packet) {
         MinecraftProtocol protocol = (MinecraftProtocol) session.getPacketProtocol();
-        if (protocol.getState() == ProtocolState.LOGIN) {
+        if (protocol.getInboundState() == ProtocolState.LOGIN) {
             if (packet instanceof ClientboundHelloPacket helloPacket) {
                 GameProfile profile = session.getFlag(MinecraftConstants.PROFILE_KEY);
                 String accessToken = session.getFlag(MinecraftConstants.ACCESS_TOKEN_KEY);
@@ -87,16 +87,17 @@ public void packetReceived(Session session, Packet packet) {
                     return;
                 }
 
-                session.send(new ServerboundKeyPacket(helloPacket.getPublicKey(), key, helloPacket.getChallenge()));
-                session.enableEncryption(protocol.enableEncryption(key));
+                session.send(new ServerboundKeyPacket(helloPacket.getPublicKey(), key, helloPacket.getChallenge()),
+                    () -> session.enableEncryption(protocol.enableEncryption(key)));
             } else if (packet instanceof ClientboundGameProfilePacket) {
-                session.send(new ServerboundLoginAcknowledgedPacket());
+                session.switchInboundState(() -> protocol.setInboundState(ProtocolState.CONFIGURATION));
+                session.send(new ServerboundLoginAcknowledgedPacket(), () -> protocol.setOutboundState(ProtocolState.CONFIGURATION));
             } else if (packet instanceof ClientboundLoginDisconnectPacket loginDisconnectPacket) {
                 session.disconnect(loginDisconnectPacket.getReason());
             } else if (packet instanceof ClientboundLoginCompressionPacket loginCompressionPacket) {
                 session.setCompressionThreshold(loginCompressionPacket.getThreshold(), false);
             }
-        } else if (protocol.getState() == ProtocolState.STATUS) {
+        } else if (protocol.getInboundState() == ProtocolState.STATUS) {
             if (packet instanceof ClientboundStatusResponsePacket statusResponsePacket) {
                 ServerStatusInfo info = statusResponsePacket.parseInfo();
                 ServerInfoHandler handler = session.getFlag(MinecraftConstants.SERVER_INFO_HANDLER_KEY);
@@ -114,13 +115,14 @@ public void packetReceived(Session session, Packet packet) {
 
                 session.disconnect(Component.translatable("multiplayer.status.finished"));
             }
-        } else if (protocol.getState() == ProtocolState.GAME) {
+        } else if (protocol.getInboundState() == ProtocolState.GAME) {
             if (packet instanceof ClientboundKeepAlivePacket keepAlivePacket && session.getFlag(MinecraftConstants.AUTOMATIC_KEEP_ALIVE_MANAGEMENT, true)) {
                 session.send(new ServerboundKeepAlivePacket(keepAlivePacket.getPingId()));
             } else if (packet instanceof ClientboundDisconnectPacket disconnectPacket) {
                 session.disconnect(disconnectPacket.getReason());
             } else if (packet instanceof ClientboundStartConfigurationPacket) {
-                session.send(new ServerboundConfigurationAcknowledgedPacket());
+                session.switchInboundState(() -> protocol.setInboundState(ProtocolState.CONFIGURATION));
+                session.send(new ServerboundConfigurationAcknowledgedPacket(), () -> protocol.setOutboundState(ProtocolState.CONFIGURATION));
             } else if (packet instanceof ClientboundTransferPacket transferPacket) {
                 if (session.getFlag(MinecraftConstants.FOLLOW_TRANSFERS, true)) {
                     TcpClientSession newSession = new TcpClientSession(transferPacket.getHost(), transferPacket.getPort(), session.getPacketProtocol());
@@ -129,9 +131,12 @@ public void packetReceived(Session session, Packet packet) {
                     newSession.connect(true, true);
                 }
             }
-        } else if (protocol.getState() == ProtocolState.CONFIGURATION) {
-            if (packet instanceof ClientboundFinishConfigurationPacket) {
-                session.send(new ServerboundFinishConfigurationPacket());
+        } else if (protocol.getInboundState() == ProtocolState.CONFIGURATION) {
+            if (packet instanceof ClientboundKeepAlivePacket keepAlivePacket && session.getFlag(MinecraftConstants.AUTOMATIC_KEEP_ALIVE_MANAGEMENT, true)) {
+                session.send(new ServerboundKeepAlivePacket(keepAlivePacket.getPingId()));
+            } else if (packet instanceof ClientboundFinishConfigurationPacket) {
+                session.switchInboundState(() -> protocol.setInboundState(ProtocolState.GAME));
+                session.send(new ServerboundFinishConfigurationPacket(), () -> protocol.setOutboundState(ProtocolState.GAME));
             } else if (packet instanceof ClientboundSelectKnownPacks) {
                 if (session.getFlag(MinecraftConstants.SEND_BLANK_KNOWN_PACKS_RESPONSE, true)) {
                     session.send(new ServerboundSelectKnownPacks(Collections.emptyList()));
@@ -148,38 +153,24 @@ public void packetReceived(Session session, Packet packet) {
     }
 
     @Override
-    public void packetSent(Session session, Packet packet) {
+    public void connected(ConnectedEvent event) {
+        Session session = event.getSession();
         MinecraftProtocol protocol = (MinecraftProtocol) session.getPacketProtocol();
-        if (packet instanceof ClientIntentionPacket) {
-            // Once the HandshakePacket has been sent, switch to the next protocol mode.
-            protocol.setState(this.targetState);
+        ClientIntentionPacket intention = new ClientIntentionPacket(protocol.getCodec().getProtocolVersion(), session.getHost(), session.getPort(), switch (targetState) {
+            case LOGIN -> transferring ? HandshakeIntent.TRANSFER : HandshakeIntent.LOGIN;
+            case STATUS -> HandshakeIntent.STATUS;
+            default -> throw new IllegalStateException("Unexpected value: " + targetState);
+        });
 
-            if (this.targetState == ProtocolState.LOGIN) {
+        session.switchInboundState(() -> protocol.setInboundState(this.targetState));
+        session.send(intention, () -> protocol.setOutboundState(this.targetState));
+        switch (this.targetState) {
+            case LOGIN -> {
                 GameProfile profile = session.getFlag(MinecraftConstants.PROFILE_KEY);
                 session.send(new ServerboundHelloPacket(profile.getName(), profile.getId()));
-            } else {
-                session.send(new ServerboundStatusRequestPacket());
-            }
-        } else if (packet instanceof ServerboundLoginAcknowledgedPacket) {
-            protocol.setState(ProtocolState.CONFIGURATION); // LOGIN -> CONFIGURATION
-        } else if (packet instanceof ServerboundFinishConfigurationPacket) {
-            protocol.setState(ProtocolState.GAME); // CONFIGURATION -> GAME
-        } else if (packet instanceof ServerboundConfigurationAcknowledgedPacket) {
-            protocol.setState(ProtocolState.CONFIGURATION); // GAME -> CONFIGURATION
-        }
-    }
-
-    @Override
-    public void connected(ConnectedEvent event) {
-        MinecraftProtocol protocol = (MinecraftProtocol) event.getSession().getPacketProtocol();
-        if (this.targetState == ProtocolState.LOGIN) {
-            if (this.transferring) {
-                event.getSession().send(new ClientIntentionPacket(protocol.getCodec().getProtocolVersion(), event.getSession().getHost(), event.getSession().getPort(), HandshakeIntent.TRANSFER));
-            } else {
-                event.getSession().send(new ClientIntentionPacket(protocol.getCodec().getProtocolVersion(), event.getSession().getHost(), event.getSession().getPort(), HandshakeIntent.LOGIN));
             }
-        } else if (this.targetState == ProtocolState.STATUS) {
-            event.getSession().send(new ClientIntentionPacket(protocol.getCodec().getProtocolVersion(), event.getSession().getHost(), event.getSession().getPort(), HandshakeIntent.STATUS));
+            case STATUS -> session.send(new ServerboundStatusRequestPacket());
+            default -> throw new IllegalStateException("Unexpected value: " + targetState);
         }
     }
 }