From 74c932e5792135bdef5b2048b75bcf137cc4aba8 Mon Sep 17 00:00:00 2001
From: Adrian <68704415+4drian3d@users.noreply.github.com>
Date: Sat, 24 Feb 2024 14:24:06 -0500
Subject: [PATCH] Fix sending keep alive response to wrong backend server
 (#1253)

---
 .../client/ClientConfigSessionHandler.java    | 34 +++++++++++++------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientConfigSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientConfigSessionHandler.java
index 48ea03c41..667fe8d12 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientConfigSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientConfigSessionHandler.java
@@ -42,6 +42,8 @@ import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Handles the client config stage.
@@ -72,17 +74,11 @@ public class ClientConfigSessionHandler implements MinecraftSessionHandler {
   }
 
   @Override
-  public boolean handle(KeepAlivePacket packet) {
+  public boolean handle(final KeepAlivePacket packet) {
     final VelocityServerConnection serverConnection = player.getConnectedServer();
-    if (serverConnection != null) {
-      final Long sentTime = serverConnection.getPendingPings().remove(packet.getRandomId());
-      if (sentTime != null) {
-        final MinecraftConnection smc = serverConnection.getConnection();
-        if (smc != null) {
-          player.setPing(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - sentTime));
-          smc.write(packet);
-        }
-      }
+    if (!this.sendKeepAliveToBackend(serverConnection, packet)) {
+      final VelocityServerConnection connectionInFlight = player.getConnectionInFlight();
+      this.sendKeepAliveToBackend(connectionInFlight, packet);
     }
     return true;
   }
@@ -177,6 +173,24 @@ public class ClientConfigSessionHandler implements MinecraftSessionHandler {
         Component.translatable("velocity.error.player-connection-error", NamedTextColor.RED));
   }
 
+  private boolean sendKeepAliveToBackend(
+          final @Nullable VelocityServerConnection serverConnection,
+          final @NotNull KeepAlivePacket packet
+  ) {
+    if (serverConnection != null) {
+      final Long sentTime = serverConnection.getPendingPings().remove(packet.getRandomId());
+      if (sentTime != null) {
+        final MinecraftConnection smc = serverConnection.getConnection();
+        if (smc != null) {
+          player.setPing(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - sentTime));
+          smc.write(packet);
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
   /**
    * Handles the backend finishing the config stage.
    *