From 2d015cdf5cafd55c89b4fc9b5cd4373a2ef6f04c Mon Sep 17 00:00:00 2001
From: Adrian <68704415+4drian3d@users.noreply.github.com>
Date: Sat, 8 Jun 2024 14:28:33 -0500
Subject: [PATCH] Add support for canceling ProxyPingEvent (#1264)

* Added support for canceling ProxyPingEvent
---
 .../api/event/proxy/ProxyPingEvent.java       | 55 +++++++++++++++++--
 .../client/StatusSessionHandler.java          | 22 +++++---
 2 files changed, 65 insertions(+), 12 deletions(-)

diff --git a/api/src/main/java/com/velocitypowered/api/event/proxy/ProxyPingEvent.java b/api/src/main/java/com/velocitypowered/api/event/proxy/ProxyPingEvent.java
index 0e04acfee..635554366 100644
--- a/api/src/main/java/com/velocitypowered/api/event/proxy/ProxyPingEvent.java
+++ b/api/src/main/java/com/velocitypowered/api/event/proxy/ProxyPingEvent.java
@@ -8,9 +8,11 @@
 package com.velocitypowered.api.event.proxy;
 
 import com.google.common.base.Preconditions;
+import com.velocitypowered.api.event.ResultedEvent;
 import com.velocitypowered.api.event.annotation.AwaitingEvent;
 import com.velocitypowered.api.proxy.InboundConnection;
 import com.velocitypowered.api.proxy.server.ServerPing;
+import org.jetbrains.annotations.NotNull;
 
 /**
  * This event is fired when a request for server information is sent by a remote client, or when the
@@ -20,28 +22,71 @@ import com.velocitypowered.api.proxy.server.ServerPing;
  * amount of ping packets a client can send.
  */
 @AwaitingEvent
-public final class ProxyPingEvent {
+public final class ProxyPingEvent implements ResultedEvent<ResultedEvent.GenericResult> {
 
   private final InboundConnection connection;
   private ServerPing ping;
+  private GenericResult result = GenericResult.allowed();
 
-  public ProxyPingEvent(InboundConnection connection, ServerPing ping) {
+  public ProxyPingEvent(final InboundConnection connection, final ServerPing ping) {
     this.connection = Preconditions.checkNotNull(connection, "connection");
     this.ping = Preconditions.checkNotNull(ping, "ping");
   }
 
+  /**
+   * Obtain the connection to which the corresponding ServerPing will be sent.
+   *
+   * @return the connection that has sent the ServerPing request
+   */
   public InboundConnection getConnection() {
-    return connection;
+    return this.connection;
   }
 
+  /**
+   * Get the ServerPing to send to the connection.
+   *
+   * @return the ServerPing to send
+   */
   public ServerPing getPing() {
-    return ping;
+    return this.ping;
   }
 
-  public void setPing(ServerPing ping) {
+  /**
+   * Sets the ServerPing to send to the connection.
+   *
+   * @param ping sets the ServerPing to send
+   */
+  public void setPing(final @NotNull ServerPing ping) {
     this.ping = Preconditions.checkNotNull(ping, "ping");
   }
 
+  /**
+   * Gets whether to avoid sending a ping response to the connection.
+   *
+   * @return if a ping response to the connection will be avoided
+   * @apiNote For the ProxyPingEvent executed to obtain the MOTD for the ServerData
+   *     sent to players of versions higher than 1.19.1,
+   *     the cancellation of this event will have no effect.
+   */
+  @Override
+  public GenericResult getResult() {
+    return this.result;
+  }
+
+  /**
+   * Sets whether to avoid sending a ping response to the connection.
+   * This will automatically close the connection.
+   *
+   * @param result if a ping response to the connection will be avoided
+   * @apiNote For the ProxyPingEvent executed to obtain the MOTD for the ServerData
+   *     sent to players of versions higher than 1.19.1,
+   *     the cancellation of this event will have no effect.
+   */
+  @Override
+  public void setResult(final @NotNull GenericResult result) {
+    this.result = Preconditions.checkNotNull(result, "result");
+  }
+
   @Override
   public String toString() {
     return "ProxyPingEvent{"
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 0f7c43b84..2c2d7984f 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
@@ -68,9 +68,13 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
     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())),
-            connection.eventLoop())
+        .thenAcceptAsync(event -> {
+          if (event.getResult().isAllowed()) {
+            connection.closeWith(LegacyDisconnect.fromServerPing(event.getPing(), packet.getVersion()));
+          } else {
+            connection.close();
+          }
+        }, connection.eventLoop())
         .exceptionally((ex) -> {
           logger.error("Exception while handling legacy ping {}", packet, ex);
           return null;
@@ -95,10 +99,14 @@ public class StatusSessionHandler implements MinecraftSessionHandler {
         .thenCompose(ping -> server.getEventManager().fire(new ProxyPingEvent(inbound, ping)))
         .thenAcceptAsync(
             (event) -> {
-              final StringBuilder json = new StringBuilder();
-              VelocityServer.getPingGsonInstance(connection.getProtocolVersion())
-                  .toJson(event.getPing(), json);
-              connection.write(new StatusResponsePacket(json));
+              if (event.getResult().isAllowed()) {
+                final StringBuilder json = new StringBuilder();
+                VelocityServer.getPingGsonInstance(connection.getProtocolVersion())
+                        .toJson(event.getPing(), json);
+                connection.write(new StatusResponsePacket(json));
+              } else {
+                connection.close();
+              }
             },
             connection.eventLoop())
         .exceptionally((ex) -> {