From 86b6a6036834fda25aed40235138ad91955bac7d Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Fri, 27 Oct 2023 17:40:56 -0400
Subject: [PATCH] Refactor `StatusSessionHandler` so we can introduce cached
 server pings more easily

---
 .../client/StatusSessionHandler.java          | 25 +++-----------
 .../util/ServerListPingHandler.java           | 34 +++++++++++++++++++
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/StatusSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/StatusSessionHandler.java
index ada570079..c290c9450 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/StatusSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/StatusSessionHandler.java
@@ -17,7 +17,6 @@
 
 package com.velocitypowered.proxy.connection.client;
 
-import com.velocitypowered.api.event.proxy.ProxyPingEvent;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
@@ -66,10 +65,9 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
       throw EXPECTED_AWAITING_REQUEST;
     }
     this.pingReceived = true;
-    server.getServerListPingHandler().getInitialPing(this.inbound)
-        .thenCompose(ping -> server.getEventManager().fire(new ProxyPingEvent(inbound, ping)))
-        .thenAcceptAsync(event -> connection.closeWith(
-                LegacyDisconnect.fromServerPing(event.getPing(), packet.getVersion())),
+    server.getServerListPingHandler().getPing(this.inbound)
+        .thenAcceptAsync(ping -> connection.closeWith(
+                LegacyDisconnect.fromServerPing(ping, packet.getVersion())),
             connection.eventLoop())
         .exceptionally((ex) -> {
           logger.error("Exception while handling legacy ping {}", packet, ex);
@@ -91,16 +89,8 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
     }
     this.pingReceived = true;
 
-    this.server.getServerListPingHandler().getInitialPing(inbound)
-        .thenCompose(ping -> server.getEventManager().fire(new ProxyPingEvent(inbound, ping)))
-        .thenAcceptAsync(
-            (event) -> {
-              StringBuilder json = new StringBuilder();
-              VelocityServer.getPingGsonInstance(connection.getProtocolVersion())
-                  .toJson(event.getPing(), json);
-              connection.write(new StatusResponse(json));
-            },
-            connection.eventLoop())
+    server.getServerListPingHandler().getPacketResponse(this.inbound)
+        .thenAcceptAsync(connection::write, connection.eventLoop())
         .exceptionally((ex) -> {
           logger.error("Exception while handling status request {}", packet, ex);
           return null;
@@ -113,9 +103,4 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
     // what even is going on?
     connection.close(true);
   }
-
-  private enum State {
-    AWAITING_REQUEST,
-    RECEIVED_REQUEST
-  }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/util/ServerListPingHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/util/ServerListPingHandler.java
index d253d3594..d168d2ee2 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/util/ServerListPingHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/util/ServerListPingHandler.java
@@ -19,6 +19,7 @@ package com.velocitypowered.proxy.connection.util;
 
 import com.google.common.collect.ImmutableList;
 import com.spotify.futures.CompletableFutures;
+import com.velocitypowered.api.event.proxy.ProxyPingEvent;
 import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.api.proxy.server.PingOptions;
 import com.velocitypowered.api.proxy.server.RegisteredServer;
@@ -27,6 +28,7 @@ import com.velocitypowered.api.util.ModInfo;
 import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.config.PingPassthroughMode;
 import com.velocitypowered.proxy.config.VelocityConfiguration;
+import com.velocitypowered.proxy.protocol.packet.StatusResponse;
 import com.velocitypowered.proxy.server.VelocityRegisteredServer;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
@@ -158,4 +160,36 @@ public class ServerListPingHandler {
       return attemptPingPassthrough(connection, passthroughMode, serversToTry, shownVersion);
     }
   }
+
+  /**
+   * Gets the current server ping for this connection, firing {@code ProxyPingEvent} if the
+   * ping is not cached.
+   *
+   * @param connection the connection being pinged
+   *
+   * @return the server ping as a completable future
+   */
+  public CompletableFuture<ServerPing> getPing(VelocityInboundConnection connection) {
+    return this.getInitialPing(connection)
+        .thenCompose(ping -> server.getEventManager().fire(new ProxyPingEvent(connection, ping)))
+        .thenApply(ProxyPingEvent::getPing);
+  }
+
+  /**
+   * Gets the current server ping for this connection, firing {@code ProxyPingEvent} if the
+   * ping is not cached.
+   *
+   * @param connection the connection being pinged
+   *
+   * @return the server ping as a completable future
+   */
+  public CompletableFuture<StatusResponse> getPacketResponse(VelocityInboundConnection connection) {
+    return this.getInitialPing(connection)
+        .thenApply(ping -> {
+          StringBuilder json = new StringBuilder();
+          VelocityServer.getPingGsonInstance(connection.getProtocolVersion())
+              .toJson(ping, json);
+          return new StatusResponse(json);
+        });
+  }
 }