From 7aed76ee3deab3b42051e11606173f1f7791174c Mon Sep 17 00:00:00 2001
From: 0x22 <0x22@futureclient.net>
Date: Sat, 23 Jan 2021 18:35:14 -0500
Subject: [PATCH] Use keepalive queue instead of just the last keepalive.

---
 .../backend/BackendPlaySessionHandler.java    |  2 +-
 .../backend/VelocityServerConnection.java     | 22 +++++--------------
 .../client/ClientPlaySessionHandler.java      | 14 +++++++-----
 3 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
index 6dab7d8a5..667b1867e 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
@@ -83,7 +83,7 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
 
   @Override
   public boolean handle(KeepAlive packet) {
-    serverConn.setLastPingId(packet.getRandomId());
+    serverConn.getPendingPings().put(packet.getRandomId(), System.currentTimeMillis());
     return false; // forwards on
   }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
index fe1144157..93958663e 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
@@ -29,7 +29,9 @@ import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelFutureListener;
 import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.UnaryOperator;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -44,8 +46,7 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
   private boolean hasCompletedJoin = false;
   private boolean gracefulDisconnect = false;
   private BackendConnectionPhase connectionPhase = BackendConnectionPhases.UNKNOWN;
-  private long lastPingId;
-  private long lastPingSent;
+  private final Map<Long, Long> pendingPings = new HashMap<>();
   private @MonotonicNonNull DimensionRegistry activeDimensionRegistry;
 
   /**
@@ -244,21 +245,8 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
     return gracefulDisconnect;
   }
 
-  public long getLastPingId() {
-    return lastPingId;
-  }
-
-  public long getLastPingSent() {
-    return lastPingSent;
-  }
-
-  void setLastPingId(long lastPingId) {
-    this.lastPingId = lastPingId;
-    this.lastPingSent = System.currentTimeMillis();
-  }
-
-  public void resetLastPingId() {
-    this.lastPingId = -1;
+  public Map<Long, Long> getPendingPings() {
+    return pendingPings;
   }
 
   /**
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
index b41a5ae32..414e2520b 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
@@ -99,12 +99,14 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
   @Override
   public boolean handle(KeepAlive packet) {
     VelocityServerConnection serverConnection = player.getConnectedServer();
-    if (serverConnection != null && packet.getRandomId() == serverConnection.getLastPingId()) {
-      MinecraftConnection smc = serverConnection.getConnection();
-      if (smc != null) {
-        player.setPing(System.currentTimeMillis() - serverConnection.getLastPingSent());
-        smc.write(packet);
-        serverConnection.resetLastPingId();
+    if (serverConnection != null) {
+      Long sentTime = serverConnection.getPendingPings().remove(packet.getRandomId());
+      if (sentTime != null) {
+        MinecraftConnection smc = serverConnection.getConnection();
+        if (smc != null) {
+          player.setPing(System.currentTimeMillis() - sentTime);
+          smc.write(packet);
+        }
       }
     }
     return true;