feature: Cancellable connection establish event

Author: kezz

api/src/main/java/com/velocitypowered/api/event/connection/ConnectionEstablishEvent.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2018-2021 Velocity Contributors
+ *
+ * The Velocity API is licensed under the terms of the MIT License. For more details,
+ * reference the LICENSE file in the api top-level directory.
+ */
+
+package com.velocitypowered.api.event.connection;
+
+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 org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Event called when a connection is initially established with the proxy.
+ * Velocity will wait for this event to fire before continuing with the connection.
+ */
+@AwaitingEvent
+public class ConnectionEstablishEvent implements ResultedEvent<ResultedEvent.GenericResult> {
+  private final InboundConnection connection;
+  private final Intention intention;
+  private GenericResult result = GenericResult.allowed();
+
+  public ConnectionEstablishEvent(
+          final @NonNull InboundConnection connection,
+          final @Nullable Intention intention
+  ) {
+    this.connection = Preconditions.checkNotNull(connection, "connection");
+    this.intention = intention;
+  }
+
+  /**
+   * The intention of the connection.
+   */
+  public enum Intention {
+    /**
+     * The user intends to ping the server to fetch the status.
+     */
+    STATUS,
+    /**
+     * The user intends to log in to the server.
+     */
+    LOGIN,
+  }
+
+  /**
+   * Returns the inbound connection that is being established.
+   *
+   * @return the connection
+   */
+  public @NonNull InboundConnection getConnection() {
+    return this.connection;
+  }
+
+  /**
+   * Returns the intention for which the connection is being established, if known.
+   *
+   * @return the intention
+   */
+  public @Nullable Intention getIntention() {
+    return this.intention;
+  }
+
+  @Override
+  public GenericResult getResult() {
+    return this.result;
+  }
+
+  @Override
+  public void setResult(final @NonNull GenericResult result) {
+    this.result = Preconditions.checkNotNull(result, "result");
+  }
+}

proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java

@@ -19,6 +19,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.velocitypowered.api.event.connection.ConnectionEstablishEvent;
 import com.velocitypowered.api.event.connection.ConnectionHandshakeEvent;
 import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.proxy.VelocityServer;
@@ -85,31 +86,57 @@ public boolean handle(LegacyHandshake packet) {
   public boolean handle(Handshake handshake) {
     InitialInboundConnection ic = new InitialInboundConnection(connection,
         cleanVhost(handshake.getServerAddress()), handshake);
-    StateRegistry nextState = getStateForProtocol(handshake.getNextStatus());
-    if (nextState == null) {
-      LOGGER.error("{} provided invalid protocol {}", ic, handshake.getNextStatus());
-      connection.close(true);
-    } else {
-      connection.setProtocolVersion(handshake.getProtocolVersion());
-      connection.setAssociation(ic);
-
-      switch (nextState) {
-        case STATUS:
-          connection.setActiveSessionHandler(StateRegistry.STATUS,
-              new StatusSessionHandler(server, ic));
-          break;
-        case LOGIN:
-          this.handleLogin(handshake, ic);
-          break;
-        default:
-          // If you get this, it's a bug in Velocity.
-          throw new AssertionError("getStateForProtocol provided invalid state!");
-      }
-    }
+
+    // Handle connection establish event.
+    connection.setAutoReading(false);
+    server.getEventManager()
+            .fire(new ConnectionEstablishEvent(
+                    ic, getIntentionForStatus(handshake.getNextStatus())))
+            .thenAcceptAsync(result -> {
+              // Clean up the disabling of auto-read.
+              connection.setAutoReading(true);
+
+              if (!result.getResult().isAllowed()) {
+                connection.close(true);
+              } else {
+                StateRegistry nextState = getStateForProtocol(handshake.getNextStatus());
+                if (nextState == null) {
+                  LOGGER.error("{} provided invalid protocol {}", ic, handshake.getNextStatus());
+                  connection.close(true);
+                } else {
+                  connection.setProtocolVersion(handshake.getProtocolVersion());
+                  connection.setAssociation(ic);
+
+                  switch (nextState) {
+                    case STATUS:
+                      connection.setActiveSessionHandler(StateRegistry.STATUS,
+                              new StatusSessionHandler(server, ic));
+                      break;
+                    case LOGIN:
+                      this.handleLogin(handshake, ic);
+                      break;
+                    default:
+                      // If you get this, it's a bug in Velocity.
+                      throw new AssertionError("getStateForProtocol provided invalid state!");
+                  }
+                }
+              }
+            });
 
     return true;
   }
 
+  private static ConnectionEstablishEvent.@Nullable Intention getIntentionForStatus(int status) {
+    switch (status) {
+      case StateRegistry.STATUS_ID:
+        return ConnectionEstablishEvent.Intention.STATUS;
+      case StateRegistry.LOGIN_ID:
+        return ConnectionEstablishEvent.Intention.LOGIN;
+      default:
+        return null;
+    }
+  }
+
   private static @Nullable StateRegistry getStateForProtocol(int status) {
     switch (status) {
       case StateRegistry.STATUS_ID: