Fix login disconnect from 1.20.4 (#1245)

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

@@ -399,7 +399,8 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
       if (server.getConfiguration().isFailoverOnUnexpectedServerDisconnect()) {
         serverConn.getPlayer().handleConnectionException(serverConn.getServer(),
             DisconnectPacket.create(ConnectionMessages.INTERNAL_SERVER_CONNECTION_ERROR,
-                serverConn.getPlayer().getProtocolVersion(), false), true);
+                serverConn.getPlayer().getProtocolVersion(),
+                    serverConn.getPlayer().getConnection().getState()), true);
       } else {
         serverConn.getPlayer().disconnect(ConnectionMessages.INTERNAL_SERVER_CONNECTION_ERROR);
       }

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

@@ -606,7 +606,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
       logger.info(Component.text(this + " has disconnected: ").append(translated));
     }
     connection.closeWith(DisconnectPacket.create(translated,
-            this.getProtocolVersion(), duringLogin));
+            this.getProtocolVersion(), connection.getState()));
   }
 
   public @Nullable VelocityServerConnection getConnectedServer() {
@@ -775,7 +775,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
                   Component reason = status.getReasonComponent()
                       .orElse(ConnectionMessages.INTERNAL_SERVER_CONNECTION_ERROR);
                   handleConnectionException(res.getServer(),
-                      DisconnectPacket.create(reason, getProtocolVersion(), false),
+                      DisconnectPacket.create(reason, getProtocolVersion(), connection.getState()),
                       ((Impl) status).isSafe());
                   break;
                 case SUCCESS:
@@ -1276,7 +1276,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
             Component reason = status.getReasonComponent()
                 .orElse(ConnectionMessages.INTERNAL_SERVER_CONNECTION_ERROR);
             handleConnectionException(toConnect,
-                DisconnectPacket.create(reason, getProtocolVersion(), false), status.isSafe());
+                DisconnectPacket.create(reason, getProtocolVersion(), connection.getState()), status.isSafe());
             break;
           default:
             // The only remaining value is successful (no need to do anything!)

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

@@ -98,7 +98,7 @@ public final class InitialInboundConnection implements VelocityInboundConnection
     if (connection.server.getConfiguration().isLogPlayerConnections()) {
       logger.info(Component.text(this + " has disconnected: ").append(translated));
     }
-    connection.closeWith(DisconnectPacket.create(translated, getProtocolVersion(), true));
+    connection.closeWith(DisconnectPacket.create(translated, getProtocolVersion(), connection.getState()));
   }
 
   /**
@@ -109,6 +109,6 @@ public final class InitialInboundConnection implements VelocityInboundConnection
   public void disconnectQuietly(Component reason) {
     Component translated = GlobalTranslator.render(reason, ClosestLocaleMatcher.INSTANCE
         .lookupClosest(Locale.getDefault()));
-    connection.closeWith(DisconnectPacket.create(translated, getProtocolVersion(), true));
+    connection.closeWith(DisconnectPacket.create(translated, getProtocolVersion(), connection.getState()));
   }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/StateRegistry.java

@@ -161,7 +161,7 @@ public enum StateRegistry {
           PluginMessagePacket.class, PluginMessagePacket::new,
           map(0x00, MINECRAFT_1_20_2, false));
       clientbound.register(
-          DisconnectPacket.class, () -> new DisconnectPacket(false),
+          DisconnectPacket.class, () -> new DisconnectPacket(this),
           map(0x01, MINECRAFT_1_20_2, false));
       clientbound.register(
           FinishedUpdatePacket.class, () -> FinishedUpdatePacket.INSTANCE,
@@ -354,7 +354,7 @@ public enum StateRegistry {
           map(0x18, MINECRAFT_1_20_2, false));
       clientbound.register(
           DisconnectPacket.class,
-          () -> new DisconnectPacket(false),
+          () -> new DisconnectPacket(this),
           map(0x40, MINECRAFT_1_7_2, false),
           map(0x1A, MINECRAFT_1_9, false),
           map(0x1B, MINECRAFT_1_13, false),
@@ -595,7 +595,7 @@ public enum StateRegistry {
           map(0x03, MINECRAFT_1_20_2, false));
 
       clientbound.register(
-          DisconnectPacket.class, () -> new DisconnectPacket(true),
+          DisconnectPacket.class, () -> new DisconnectPacket(this),
           map(0x00, MINECRAFT_1_7_2, false));
       clientbound.register(
           EncryptionRequestPacket.class, EncryptionRequestPacket::new,
@@ -795,7 +795,8 @@ public enum StateRegistry {
     private final @Nullable ProtocolVersion lastValidProtocolVersion;
 
     PacketMapping(int id, ProtocolVersion protocolVersion,
-                  ProtocolVersion lastValidProtocolVersion, boolean packetDecoding) {
+                  @Nullable ProtocolVersion lastValidProtocolVersion,
+                  boolean packetDecoding) {
       this.id = id;
       this.protocolVersion = protocolVersion;
       this.lastValidProtocolVersion = lastValidProtocolVersion;

proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/DisconnectPacket.java

@@ -22,6 +22,7 @@ import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
+import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.packet.chat.ComponentHolder;
 import io.netty.buffer.ByteBuf;
 import net.kyori.adventure.text.Component;
@@ -30,14 +31,14 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 public class DisconnectPacket implements MinecraftPacket {
 
   private @Nullable ComponentHolder reason;
-  private final boolean login;
+  private final StateRegistry state;
 
-  public DisconnectPacket(boolean login) {
+  public DisconnectPacket(StateRegistry state) {
-    this.login = login;
+    this.state = state;
   }
 
-  private DisconnectPacket(boolean login, ComponentHolder reason) {
+  private DisconnectPacket(StateRegistry state, ComponentHolder reason) {
-    this.login = login;
+    this.state = state;
     this.reason = Preconditions.checkNotNull(reason, "reason");
   }
 
@@ -61,7 +62,8 @@ public class DisconnectPacket implements MinecraftPacket {
 
   @Override
   public void decode(ByteBuf buf, ProtocolUtils.Direction direction, ProtocolVersion version) {
-	  reason = ComponentHolder.read(buf, login ? ProtocolVersion.MINECRAFT_1_20_2 : version);
+	  reason = ComponentHolder.read(buf, state == StateRegistry.LOGIN
+              ? ProtocolVersion.MINECRAFT_1_20_2 : version);
   }
 
   @Override
@@ -74,8 +76,9 @@ public class DisconnectPacket implements MinecraftPacket {
     return handler.handle(this);
   }
 
-  public static DisconnectPacket create(Component component, ProtocolVersion version, boolean login) {
+  public static DisconnectPacket create(Component component, ProtocolVersion version, StateRegistry state) {
     Preconditions.checkNotNull(component, "component");
-    return new DisconnectPacket(login, new ComponentHolder(login ? ProtocolVersion.MINECRAFT_1_20_2 : version, component));
+    return new DisconnectPacket(state, new ComponentHolder(state == StateRegistry.LOGIN
+            ? ProtocolVersion.MINECRAFT_1_20_2 : version, component));
   }
 }