Waterfall/Waterfall-Proxy-Patches/0028-InitialHandler-Processing-State.patch

From cdf449f0be0a29592f76f2b53afc2da0949d58f1 Mon Sep 17 00:00:00 2001
From: LinsaFTW <25271111+linsaftw@users.noreply.github.com>
Date: Fri, 25 Feb 2022 12:28:31 -0300
Subject: [PATCH] InitialHandler Processing State


diff --git a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
index 0c608e0f..4d9e9b89 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
@@ -128,12 +128,12 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     private enum State
     {
 
-        HANDSHAKE, STATUS, PING, USERNAME, ENCRYPT, FINISHING;
+        PROCESSING, PROCESSING_USERNAME, HANDSHAKE, STATUS, PING, USERNAME, ENCRYPT, FINISHING;
     }
 
     private boolean canSendKickMessage()
     {
-        return thisState == State.USERNAME || thisState == State.ENCRYPT || thisState == State.FINISHING;
+        return thisState == State.PROCESSING_USERNAME || thisState == State.USERNAME || thisState == State.ENCRYPT || thisState == State.FINISHING;
     }
 
     @Override
@@ -248,6 +248,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     public void handle(StatusRequest statusRequest) throws Exception
     {
         Preconditions.checkState( thisState == State.STATUS, "Not expecting STATUS" );
+        thisState = State.PROCESSING;
 
         ServerInfo forced = AbstractReconnectHandler.getForcedHost( this );
         // FlameCord - Custom motd
@@ -328,6 +329,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     {
         // FlameCord - Never accept invalid packets
         Preconditions.checkState( thisState == State.PING, "Not expecting PING" );
+        thisState = State.PROCESSING;
 
         unsafe.sendPacket( ping );
 
@@ -339,6 +341,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     public void handle(Handshake handshake) throws Exception
     {
         Preconditions.checkState( thisState == State.HANDSHAKE, "Not expecting HANDSHAKE" );
+        thisState = State.PROCESSING;
         this.handshake = handshake;
         ch.setVersion( handshake.getProtocolVersion() );
 
@@ -411,6 +414,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     public void handle(LoginRequest loginRequest) throws Exception
     {
         Preconditions.checkState( thisState == State.USERNAME, "Not expecting USERNAME" );
+        thisState = State.PROCESSING_USERNAME;
 
         if ( !AllowedCharacters.isValidName( loginRequest.getData(), onlineMode ) )
         {
@@ -469,6 +473,7 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     public void handle(final EncryptionResponse encryptResponse) throws Exception
     {
         Preconditions.checkState( thisState == State.ENCRYPT, "Not expecting ENCRYPT" );
+        thisState = State.PROCESSING;
 
         SecretKey sharedKey = EncryptionUtil.getSecret( encryptResponse, request );
         // Waterfall start
@@ -696,14 +701,14 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Override
     public void setOnlineMode(boolean onlineMode)
     {
-        Preconditions.checkState( thisState == State.USERNAME, "Can only set online mode status whilst state is username" );
+        Preconditions.checkState( thisState == State.USERNAME || thisState == State.PROCESSING_USERNAME, "Can only set online mode status whilst state is username" );
         this.onlineMode = onlineMode;
     }
 
     @Override
     public void setUniqueId(UUID uuid)
     {
-        Preconditions.checkState( thisState == State.USERNAME, "Can only set uuid while state is username" );
+        Preconditions.checkState( thisState == State.USERNAME || thisState == State.PROCESSING_USERNAME, "Can only set uuid while state is username" );
         // FlameCord - Allow custom uuids even if onlineMode is true
         this.uniqueId = uuid;
     }
-- 
2.32.0