From 7c770269836cfcda771207e9e40fd3d23ef85467 Mon Sep 17 00:00:00 2001 From: Kieran Wallbanks Date: Tue, 19 Sep 2023 17:53:44 +0100 Subject: [PATCH] feature: Cancellable connection establish event --- .../connection/ConnectionEstablishEvent.java | 76 +++++++++++++++++++ .../client/HandshakeSessionHandler.java | 69 ++++++++++++----- 2 files changed, 124 insertions(+), 21 deletions(-) create mode 100644 api/src/main/java/com/velocitypowered/api/event/connection/ConnectionEstablishEvent.java diff --git a/api/src/main/java/com/velocitypowered/api/event/connection/ConnectionEstablishEvent.java b/api/src/main/java/com/velocitypowered/api/event/connection/ConnectionEstablishEvent.java new file mode 100644 index 0000000000..dd5c12f690 --- /dev/null +++ b/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 { + 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"); + } +} diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java index 079ce035f8..347af888df 100644 --- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java +++ b/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: