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

@@ -1,12 +1,14 @@
 package org.geysermc.mcprotocollib.network;
 
+import io.netty.channel.Channel;
 import net.kyori.adventure.text.Component;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.geysermc.mcprotocollib.network.codec.PacketCodecHelper;
 import org.geysermc.mcprotocollib.network.crypt.PacketEncryption;
 import org.geysermc.mcprotocollib.network.event.session.SessionEvent;
 import org.geysermc.mcprotocollib.network.event.session.SessionListener;
+import org.geysermc.mcprotocollib.network.packet.FakeFlushPacket;
 import org.geysermc.mcprotocollib.network.packet.Packet;
 import org.geysermc.mcprotocollib.network.packet.PacketProtocol;
 
@@ -216,7 +218,17 @@ public interface Session {
      *
      * @param packet Packet to send.
      */
-    void send(Packet packet);
+    default void send(@NonNull 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(@NonNull Packet packet, @Nullable Runnable onSent);
 
     /**
      * Disconnects the session.
@@ -259,4 +271,48 @@ 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.
+     *                 Default is true.
+     */
+    void setAutoRead(boolean autoRead);
+
+    /**
+     * Returns the underlying netty channel of this session.
+     *
+     * @return The netty channel
+     */
+    Channel getChannel();
+
+    /**
+     * 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);
+    }
+
+    /**
+     * Flushes all packets due to be sent and changes the outbound state of the session.
+     * This makes sure no other threads scheduled packets to be sent.
+     *
+     * @param switcher The runnable that switches the outbound state.
+     */
+    default void switchOutboundState(Runnable switcher) {
+        getChannel().writeAndFlush(FakeFlushPacket.INSTANCE).syncUninterruptibly();
+
+        switcher.run();
+    }
 }

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

@@ -0,0 +1,11 @@
+package org.geysermc.mcprotocollib.network.packet;
+
+/**
+ * Packet only exists to be writeAndFlushed to ensure all packets that came before are encoded and written out.
+ */
+public class FakeFlushPacket implements Packet {
+    public static final FakeFlushPacket INSTANCE = new FakeFlushPacket();
+
+    private FakeFlushPacket() {
+    }
+}

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

@@ -115,9 +115,10 @@ public void initChannel(Channel channel) {
 
                     pipeline.addLast("sizer", new TcpPacketSizer(protocol.getPacketHeader(), getCodecHelper()));
 
-                    pipeline.addLast("codec", new TcpPacketCodec(TcpClientSession.this, true));
-                    pipeline.addLast("manager", TcpClientSession.this);
-                }
+                    pipeline.addLast("flow-control", new TcpFlowControlHandler());
+                            pipeline.addLast("codec", new TcpPacketCodec(TcpClientSession.this, true));
+                            pipeline.addLast("manager", TcpClientSession.this);
+}
             });
 
         if (getFlag(BuiltinFlags.TCP_FAST_OPEN, false) && TRANSPORT_TYPE.supportsTcpFastOpenClient()) {

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

@@ -7,12 +7,18 @@
 import org.geysermc.mcprotocollib.network.codec.PacketCodecHelper;
 import org.geysermc.mcprotocollib.network.codec.PacketDefinition;
 import org.geysermc.mcprotocollib.network.event.session.PacketErrorEvent;
+import org.geysermc.mcprotocollib.network.packet.FakeFlushPacket;
 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,17 +30,33 @@ public TcpPacketCodec(Session session, boolean client) {
     @SuppressWarnings({"rawtypes", "unchecked"})
     @Override
     public void encode(ChannelHandlerContext ctx, Packet packet, ByteBuf buf) {
+        if (packet == FakeFlushPacket.INSTANCE) {
+            log.debug("Fake flush packet reached");
+            return;
+        }
+
+        if (log.isTraceEnabled()) {
+            log.trace("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);
+
+            if (log.isDebugEnabled()) {
+                log.debug("Encoded packet {} ({})", packet.getClass().getSimpleName(), packetId);
+            }
         } catch (Throwable t) {
+            log.debug("Error encoding packet", t);
+
             // Reset writer index to make sure incomplete data is not written out.
             buf.writerIndex(initial);
 
@@ -51,6 +73,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,14 +82,22 @@ 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.trace("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(), id);
+            }
         } catch (Throwable t) {
+            log.debug("Error decoding packet", 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

@@ -66,6 +66,7 @@ public void initChannel(Channel channel) {
 
                 pipeline.addLast("sizer", new TcpPacketSizer(protocol.getPacketHeader(), session.getCodecHelper()));
 
+                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

@@ -221,18 +221,27 @@ public boolean isConnected() {
     }
 
     @Override
-    public void send(Packet packet) {
+    public void send(@NonNull Packet packet, @Nullable Runnable onSent) {
         if (this.channel == null) {
             return;
         }
 
+        if (!this.channel.eventLoop().inEventLoop()) {
+            this.channel.eventLoop().execute(() -> this.send(packet, onSent));
+            return;
+        }
+
         PacketSendingEvent sendingEvent = new PacketSendingEvent(this, packet);
         this.callEvent(sendingEvent);
 
         if (!sendingEvent.isCancelled()) {
             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());
@@ -257,6 +266,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;
@@ -272,6 +288,7 @@ public void disconnect(@NonNull Component reason, @Nullable Throwable cause) {
         return PACKET_EVENT_LOOP.next();
     }
 
+    @Override
     public Channel getChannel() {
         return this.channel;
     }
@@ -309,6 +326,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 {