From 7b409fd55b08e0d92bc768061b1c338bd05eb9ec Mon Sep 17 00:00:00 2001
From: Konicai <71294714+Konicai@users.noreply.github.com>
Date: Tue, 1 Aug 2023 10:58:59 -0400
Subject: [PATCH] Cache the Keep Alive timestamp for forwarding ping (#4002)

* Cache the (clientbound) Keep Alive timestamp and use that for forwarding ping
* Use a Queue of keep alive IDs to handle KeepAlive packets sent in succession
* Don't force NetworkStackLatencyTranslator on the session's event loop
* Send clientbound NetworkStackLatencyPacket immediately
* Avoid sending negative NetworkStackLatencyPackets that are not from the form-image-hack in FormCache
* Downsize timestamps that would lead to overflow on the client
---
 .../geyser/session/GeyserSession.java         |  8 +++
 .../geyser/session/cache/FormCache.java       |  9 ++-
 .../BedrockNetworkStackLatencyTranslator.java | 59 +++++++++----------
 .../java/JavaKeepAliveTranslator.java         | 24 +++++++-
 4 files changed, 67 insertions(+), 33 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
index dba4bd112..8960fb626 100644
--- a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
+++ b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
@@ -114,6 +114,7 @@ import org.geysermc.geyser.api.network.AuthType;
 import org.geysermc.geyser.api.network.RemoteServer;
 import org.geysermc.geyser.command.GeyserCommandSource;
 import org.geysermc.geyser.configuration.EmoteOffhandWorkaroundOption;
+import org.geysermc.geyser.configuration.GeyserConfiguration;
 import org.geysermc.geyser.entity.EntityDefinitions;
 import org.geysermc.geyser.entity.attribute.GeyserAttributeType;
 import org.geysermc.geyser.entity.type.Entity;
@@ -155,6 +156,7 @@ import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -564,6 +566,12 @@ public class GeyserSession implements GeyserConnection, GeyserCommandSource {
     @Setter
     private ScheduledFuture<?> mountVehicleScheduledFuture = null;
 
+    /**
+     * A cache of IDs from ClientboundKeepAlivePackets that have been sent to the Bedrock client, but haven't been returned to the server.
+     * Only used if {@link GeyserConfiguration#isForwardPlayerPing()} is enabled.
+     */
+    private final Queue<Long> keepAliveCache = new ConcurrentLinkedQueue<>();
+
     private MinecraftProtocol protocol;
 
     public GeyserSession(GeyserImpl geyser, BedrockServerSession bedrockServerSession, EventLoop eventLoop) {
diff --git a/core/src/main/java/org/geysermc/geyser/session/cache/FormCache.java b/core/src/main/java/org/geysermc/geyser/session/cache/FormCache.java
index 6ba6a1b7e..3f7df97c1 100644
--- a/core/src/main/java/org/geysermc/geyser/session/cache/FormCache.java
+++ b/core/src/main/java/org/geysermc/geyser/session/cache/FormCache.java
@@ -42,6 +42,13 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 @RequiredArgsConstructor
 public class FormCache {
+
+    /**
+     * The magnitude of this doesn't actually matter, but it must be negative so that
+     * BedrockNetworkStackLatencyTranslator can detect the hack.
+     */
+    private static final long MAGIC_FORM_IMAGE_HACK_TIMESTAMP = -1234567890L;
+
     private final FormDefinitions formDefinitions = FormDefinitions.instance();
     private final AtomicInteger formIdCounter = new AtomicInteger(0);
     private final Int2ObjectMap<Form> forms = new Int2ObjectOpenHashMap<>();
@@ -73,7 +80,7 @@ public class FormCache {
         if (form instanceof SimpleForm) {
             NetworkStackLatencyPacket latencyPacket = new NetworkStackLatencyPacket();
             latencyPacket.setFromServer(true);
-            latencyPacket.setTimestamp(-System.currentTimeMillis());
+            latencyPacket.setTimestamp(MAGIC_FORM_IMAGE_HACK_TIMESTAMP);
             session.scheduleInEventLoop(
                     () -> session.sendUpstreamPacket(latencyPacket),
                     500, TimeUnit.MILLISECONDS
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockNetworkStackLatencyTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockNetworkStackLatencyTranslator.java
index 0fe272ea9..1ccd5ced9 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockNetworkStackLatencyTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockNetworkStackLatencyTranslator.java
@@ -29,9 +29,7 @@ import com.github.steveice10.mc.protocol.packet.ingame.serverbound.ServerboundKe
 import org.cloudburstmc.protocol.bedrock.data.AttributeData;
 import org.cloudburstmc.protocol.bedrock.packet.NetworkStackLatencyPacket;
 import org.cloudburstmc.protocol.bedrock.packet.UpdateAttributesPacket;
-import org.geysermc.floodgate.util.DeviceOs;
 import org.geysermc.geyser.entity.attribute.GeyserAttributeType;
-import org.geysermc.geyser.network.GameProtocol;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
 import org.geysermc.geyser.translator.protocol.Translator;
@@ -47,42 +45,43 @@ public class BedrockNetworkStackLatencyTranslator extends PacketTranslator<Netwo
 
     @Override
     public void translate(GeyserSession session, NetworkStackLatencyPacket packet) {
-        long pingId;
-        // so apparently, as of 1.16.200
-        // PS4 divides the network stack latency timestamp FOR US!!!
-        // WTF
-        if (GameProtocol.isPre1_20_10(session)) {
-            if (session.getClientData().getDeviceOs().equals(DeviceOs.PS4)) {
-                pingId = packet.getTimestamp();
-            } else {
-                pingId = packet.getTimestamp() / 1000;
-            }
-        } else {
-            // changed in 1.20.10 todo: is ps4 still different?
-            pingId = packet.getTimestamp() / (1000 * 1000 * 1000);
-        }
-
         // negative timestamps are used as hack to fix the url image loading bug
-        if (packet.getTimestamp() > 0) {
+        if (packet.getTimestamp() >= 0) {
             if (session.getGeyser().getConfig().isForwardPlayerPing()) {
-                ServerboundKeepAlivePacket keepAlivePacket = new ServerboundKeepAlivePacket(pingId);
+                // use our cached value because
+                // a) bedrock can be inaccurate with the value returned
+                // b) playstation replies with a different magnitude than other platforms
+                // c) 1.20.10 and later reply with a different magnitude
+                Long keepAliveId = session.getKeepAliveCache().poll();
+                if (keepAliveId == null) {
+                    session.getGeyser().getLogger().debug("Received a latency packet that we don't have a KeepAlive for: " + packet);
+                    return;
+                }
+
+                ServerboundKeepAlivePacket keepAlivePacket = new ServerboundKeepAlivePacket(keepAliveId);
                 session.sendDownstreamPacket(keepAlivePacket);
             }
             return;
         }
 
-        // Hack to fix the url image loading bug
-        UpdateAttributesPacket attributesPacket = new UpdateAttributesPacket();
-        attributesPacket.setRuntimeEntityId(session.getPlayerEntity().getGeyserId());
+        session.scheduleInEventLoop(() -> {
+            // Hack to fix the url image loading bug
+            UpdateAttributesPacket attributesPacket = new UpdateAttributesPacket();
+            attributesPacket.setRuntimeEntityId(session.getPlayerEntity().getGeyserId());
 
-        AttributeData attribute = session.getPlayerEntity().getAttributes().get(GeyserAttributeType.EXPERIENCE_LEVEL);
-        if (attribute != null) {
-            attributesPacket.setAttributes(Collections.singletonList(attribute));
-        } else {
-            attributesPacket.setAttributes(Collections.singletonList(GeyserAttributeType.EXPERIENCE_LEVEL.getAttribute(0)));
-        }
+            AttributeData attribute = session.getPlayerEntity().getAttributes().get(GeyserAttributeType.EXPERIENCE_LEVEL);
+            if (attribute != null) {
+                attributesPacket.setAttributes(Collections.singletonList(attribute));
+            } else {
+                attributesPacket.setAttributes(Collections.singletonList(GeyserAttributeType.EXPERIENCE_LEVEL.getAttribute(0)));
+            }
 
-        session.scheduleInEventLoop(() -> session.sendUpstreamPacket(attributesPacket),
-                500, TimeUnit.MILLISECONDS);
+            session.sendUpstreamPacket(attributesPacket);
+        }, 500, TimeUnit.MILLISECONDS);
+    }
+
+    @Override
+    public boolean shouldExecuteInEventLoop() {
+        return false;
     }
 }
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/JavaKeepAliveTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/JavaKeepAliveTranslator.java
index 41eb5062a..dc7b7f316 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/JavaKeepAliveTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/JavaKeepAliveTranslator.java
@@ -42,10 +42,30 @@ public class JavaKeepAliveTranslator extends PacketTranslator<ClientboundKeepAli
         if (!session.getGeyser().getConfig().isForwardPlayerPing()) {
             return;
         }
+        // We use this once the client replies (see BedrockNetworkStackLatencyTranslator)
+        session.getKeepAliveCache().add(packet.getPingId());
+
+        long timestamp = packet.getPingId();
+
+        // We take the abs because we rely on the client responding with a negative value ONLY when we send
+        // a negative timestamp in the form-image-hack performed in FormCache.
+        // Apart from that case, we don't actually use the value the client responds with, instead using our keep alive cache.
+        if (timestamp == Long.MIN_VALUE) {
+            timestamp = Long.MAX_VALUE; // There is not an abs representation of MIN_VALUE (2's complement)
+        } else {
+            timestamp = Math.absExact(timestamp);
+        }
+
+        // Bedrock will overflow on timestamps that are too large, and respond with a mangled negative value.
+        // Keeping leftmost digits allows for easier debugging
+        while (timestamp > 1e10) {
+            timestamp /= 10;
+        }
+
         NetworkStackLatencyPacket latencyPacket = new NetworkStackLatencyPacket();
         latencyPacket.setFromServer(true);
-        latencyPacket.setTimestamp(packet.getPingId() * 1000);
-        session.sendUpstreamPacket(latencyPacket);
+        latencyPacket.setTimestamp(timestamp);
+        session.sendUpstreamPacketImmediately(latencyPacket);
     }
 
     @Override