fix missing server info in some events when kicked from current server (#703)

api/src/main/java/com/velocitypowered/api/event/player/ServerPreConnectEvent.java

@@ -13,6 +13,7 @@ import com.velocitypowered.api.event.annotation.AwaitingEvent;
 import com.velocitypowered.api.proxy.ConnectionRequestBuilder;
 import com.velocitypowered.api.proxy.ConnectionRequestBuilder.Status;
 import com.velocitypowered.api.proxy.Player;
+import com.velocitypowered.api.proxy.ServerConnection;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import java.util.Optional;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -27,6 +28,7 @@ public final class ServerPreConnectEvent implements
 
   private final Player player;
   private final RegisteredServer originalServer;
+  private final RegisteredServer previousServer;
   private ServerResult result;
 
   /**
@@ -35,8 +37,21 @@ public final class ServerPreConnectEvent implements
    * @param originalServer the server the player was trying to connect to
    */
   public ServerPreConnectEvent(Player player, RegisteredServer originalServer) {
+    this(player, originalServer,
+        player.getCurrentServer().map(ServerConnection::getServer).orElse(null));
+  }
+
+  /**
+   * Creates the ServerPreConnectEvent.
+   * @param player the player who is connecting to a server
+   * @param originalServer the server the player was trying to connect to
+   * @param previousServer the server the player ís connected to
+   */
+  public ServerPreConnectEvent(Player player, RegisteredServer originalServer,
+      @Nullable RegisteredServer previousServer) {
     this.player = Preconditions.checkNotNull(player, "player");
     this.originalServer = Preconditions.checkNotNull(originalServer, "originalServer");
+    this.previousServer = previousServer;
     this.result = ServerResult.allowed(originalServer);
   }
 
@@ -61,13 +76,24 @@ public final class ServerPreConnectEvent implements
   /**
    * Returns the server that the player originally tried to connect to. To get the server the
    * player will connect to, see the {@link ServerResult} of this event. To get the server the
-   * player is currently on when this event is fired, use {@link Player#getCurrentServer()}.
+   * player is currently on when this event is fired, use {@link #getPreviousServer()}.
    * @return the server that the player originally tried to connect to
    */
   public RegisteredServer getOriginalServer() {
     return originalServer;
   }
 
+  /**
+   * Returns the server that the player is currently connected to. Prefer this method over using
+   * {@link Player#getCurrentServer()} as the current server might get reset after server kicks to
+   * prevent connection issues. This is {@code null} if they were not connected to another server
+   * beforehand (for instance, if the player has just joined the proxy).
+   * @return the server the player is currently connected to.
+   */
+  public @Nullable RegisteredServer getPreviousServer() {
+    return previousServer;
+  }
+
   @Override
   public String toString() {
     return "ServerPreConnectEvent{"

api/src/main/java/com/velocitypowered/api/proxy/ServerConnection.java

@@ -11,6 +11,7 @@ import com.velocitypowered.api.proxy.messages.ChannelMessageSink;
 import com.velocitypowered.api.proxy.messages.ChannelMessageSource;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.api.proxy.server.ServerInfo;
+import java.util.Optional;
 
 /**
  * Represents a connection to a backend server from the proxy for a client.
@@ -24,6 +25,14 @@ public interface ServerConnection extends ChannelMessageSource, ChannelMessageSi
    */
   RegisteredServer getServer();
 
+  /**
+   * Returns the server that the player associated with this connection was connected to before
+   * switching to this connection.
+   *
+   * @return the server the player was connected to.
+   */
+  Optional<RegisteredServer> getPreviousServer();
+
   /**
    * Returns the server info for this connection.
    *

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

@@ -22,6 +22,7 @@ import static com.velocitypowered.proxy.connection.forge.legacy.LegacyForgeHands
 
 import com.velocitypowered.api.event.player.ServerConnectedEvent;
 import com.velocitypowered.api.event.player.ServerPostConnectEvent;
+import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.ConnectionTypes;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
@@ -88,6 +89,7 @@ public class TransitionSessionHandler implements MinecraftSessionHandler {
   @Override
   public boolean handle(JoinGame packet) {
     MinecraftConnection smc = serverConn.ensureConnected();
+    RegisteredServer previousServer = serverConn.getPreviousServer().orElse(null);
     VelocityServerConnection existingConnection = serverConn.getPlayer().getConnectedServer();
 
     final ConnectedPlayer player = serverConn.getPlayer();
@@ -104,8 +106,7 @@ public class TransitionSessionHandler implements MinecraftSessionHandler {
     // The goods are in hand! We got JoinGame. Let's transition completely to the new state.
     smc.setAutoReading(false);
     server.getEventManager()
-        .fire(new ServerConnectedEvent(player, serverConn.getServer(),
-            existingConnection != null ? existingConnection.getServer() : null))
+        .fire(new ServerConnectedEvent(player, serverConn.getServer(), previousServer))
         .thenRunAsync(() -> {
           // Make sure we can still transition (player might have disconnected here).
           if (!serverConn.isActive()) {
@@ -136,7 +137,7 @@ public class TransitionSessionHandler implements MinecraftSessionHandler {
 
           // We're done! :)
           server.getEventManager().fireAndForget(new ServerPostConnectEvent(player,
-              existingConnection == null ? null : existingConnection.getServer()));
+              previousServer));
           resultFuture.complete(ConnectionRequestResults.successful(serverConn.getServer()));
         }, smc.eventLoop())
         .exceptionally(exc -> {

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

@@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableList;
 import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.api.proxy.ServerConnection;
 import com.velocitypowered.api.proxy.messages.ChannelIdentifier;
+import com.velocitypowered.api.proxy.server.RegisteredServer;
 import com.velocitypowered.api.proxy.server.ServerInfo;
 import com.velocitypowered.api.util.GameProfile.Property;
 import com.velocitypowered.proxy.VelocityServer;
@@ -49,6 +50,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.UnaryOperator;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -57,6 +59,7 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 public class VelocityServerConnection implements MinecraftConnectionAssociation, ServerConnection {
 
   private final VelocityRegisteredServer registeredServer;
+  private final @Nullable VelocityRegisteredServer previousServer;
   private final ConnectedPlayer proxyPlayer;
   private final VelocityServer server;
   private @Nullable MinecraftConnection connection;
@@ -69,12 +72,15 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
   /**
    * Initializes a new server connection.
    * @param registeredServer the server to connect to
+   * @param previousServer the server the player is coming from
    * @param proxyPlayer the player connecting to the server
    * @param server the Velocity proxy instance
    */
   public VelocityServerConnection(VelocityRegisteredServer registeredServer,
+      @Nullable VelocityRegisteredServer previousServer,
       ConnectedPlayer proxyPlayer, VelocityServer server) {
     this.registeredServer = registeredServer;
+    this.previousServer = previousServer;
     this.proxyPlayer = proxyPlayer;
     this.server = server;
   }
@@ -209,6 +215,11 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
     return registeredServer;
   }
 
+  @Override
+  public Optional<RegisteredServer> getPreviousServer() {
+    return Optional.ofNullable(this.previousServer);
+  }
+
   @Override
   public ServerInfo getServerInfo() {
     return registeredServer.getServerInfo();

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

@@ -482,7 +482,12 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
 
   @Override
   public ConnectionRequestBuilder createConnectionRequest(RegisteredServer server) {
-    return new ConnectionRequestBuilderImpl(server);
+    return new ConnectionRequestBuilderImpl(server, this.connectedServer);
+  }
+
+  private ConnectionRequestBuilder createConnectionRequest(RegisteredServer server,
+      @Nullable VelocityServerConnection previousConnection) {
+    return new ConnectionRequestBuilderImpl(server, previousConnection);
   }
 
   @Override
@@ -652,7 +657,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
           connectionInFlight = null;
 
           // Make sure we clear the current connected server as the connection is invalid.
-          boolean previouslyConnected = connectedServer != null;
+          VelocityServerConnection previousConnection = connectedServer;
           if (kickedFromCurrent) {
             connectedServer = null;
           }
@@ -667,7 +672,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
             disconnect(res.getReasonComponent());
           } else if (event.getResult() instanceof RedirectPlayer) {
             RedirectPlayer res = (RedirectPlayer) event.getResult();
-            createConnectionRequest(res.getServer())
+            createConnectionRequest(res.getServer(), previousConnection)
                 .connect()
                 .whenCompleteAsync((status, throwable) -> {
                   if (throwable != null) {
@@ -710,7 +715,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
                 }, connection.eventLoop());
           } else if (event.getResult() instanceof Notify) {
             Notify res = (Notify) event.getResult();
-            if (event.kickedDuringServerConnect() && previouslyConnected) {
+            if (event.kickedDuringServerConnect() && previousConnection != null) {
               sendMessage(Identity.nil(), res.getMessageComponent());
             } else {
               disconnect(res.getMessageComponent());
@@ -1060,9 +1065,12 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
   private class ConnectionRequestBuilderImpl implements ConnectionRequestBuilder {
 
     private final RegisteredServer toConnect;
+    private final @Nullable VelocityRegisteredServer previousServer;
 
-    ConnectionRequestBuilderImpl(RegisteredServer toConnect) {
+    ConnectionRequestBuilderImpl(RegisteredServer toConnect,
+        @Nullable VelocityServerConnection previousConnection) {
       this.toConnect = Preconditions.checkNotNull(toConnect, "info");
+      this.previousServer = previousConnection == null ? null : previousConnection.getServer();
     }
 
     @Override
@@ -1096,7 +1104,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
             }
 
             ServerPreConnectEvent event = new ServerPreConnectEvent(ConnectedPlayer.this,
-                toConnect);
+                toConnect, previousServer);
             return server.getEventManager().fire(event)
                 .thenComposeAsync(newEvent -> {
                   Optional<RegisteredServer> newDest = newEvent.getResult().getServer();
@@ -1114,7 +1122,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player {
 
                   VelocityRegisteredServer vrs = (VelocityRegisteredServer) realDestination;
                   VelocityServerConnection con = new VelocityServerConnection(vrs,
-                      ConnectedPlayer.this, server);
+                      previousServer, ConnectedPlayer.this, server);
                   connectionInFlight = con;
                   return con.connect().whenCompleteAsync(
                       (result, exception) -> this.resetIfInFlightIs(con), connection.eventLoop());