Merge pull request #498 from VelocityPowered/feature/extended-connection-handshake-event

Reimplement #434 but aligned to the Velocity 2.0.0 API

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

@@ -7,12 +7,27 @@
 
 package com.velocitypowered.api.event.connection;
 
+import com.velocitypowered.api.event.ResultedEvent;
+import com.velocitypowered.api.event.ResultedEvent.ComponentResult;
 import com.velocitypowered.api.proxy.connection.InboundConnection;
+import java.net.InetAddress;
+import java.net.SocketAddress;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * This event is fired when a handshake is established between a client and the proxy.
  */
-public interface ConnectionHandshakeEvent {
+public interface ConnectionHandshakeEvent extends ResultedEvent<ComponentResult> {
 
   InboundConnection connection();
+
+  String currentHostname();
+
+  String originalHostname();
+
+  void setCurrentHostname(String hostname);
+
+  @Nullable SocketAddress currentRemoteHostAddress();
+
+  void setCurrentRemoteHostAddress(@Nullable SocketAddress address);
 }

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

@@ -9,6 +9,9 @@ package com.velocitypowered.api.event.connection;
 
 import com.google.common.base.Preconditions;
 import com.velocitypowered.api.proxy.connection.InboundConnection;
+import java.net.InetAddress;
+import java.net.SocketAddress;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * This event is fired when a handshake is established between a client and the proxy.
@@ -16,9 +19,18 @@ import com.velocitypowered.api.proxy.connection.InboundConnection;
 public final class ConnectionHandshakeEventImpl implements ConnectionHandshakeEvent {
 
   private final InboundConnection connection;
+  private final String originalHostname;
+  private String currentHostname;
+  private ComponentResult result;
+  private SocketAddress currentRemoteAddress;
 
-  public ConnectionHandshakeEventImpl(InboundConnection connection) {
+  public ConnectionHandshakeEventImpl(InboundConnection connection,
+      String originalHostname) {
     this.connection = Preconditions.checkNotNull(connection, "connection");
+    this.originalHostname = Preconditions.checkNotNull(originalHostname, "originalHostname");
+    this.currentHostname = originalHostname;
+    this.result = ComponentResult.allowed();
+    this.currentRemoteAddress = connection.remoteAddress();
   }
 
   @Override
@@ -26,10 +38,49 @@ public final class ConnectionHandshakeEventImpl implements ConnectionHandshakeEv
     return connection;
   }
 
+  @Override
+  public String currentHostname() {
+    return currentHostname;
+  }
+
+  @Override
+  public String originalHostname() {
+    return originalHostname;
+  }
+
+  @Override
+  public void setCurrentHostname(String hostname) {
+    currentHostname = Preconditions.checkNotNull(hostname, "hostname");
+  }
+
+  @Override
+  public @Nullable SocketAddress currentRemoteHostAddress() {
+    return currentRemoteAddress;
+  }
+
+  @Override
+  public void setCurrentRemoteHostAddress(SocketAddress address) {
+    currentRemoteAddress = address;
+  }
+
+  @Override
+  public ComponentResult result() {
+    return result;
+  }
+
+  @Override
+  public void setResult(ComponentResult result) {
+    this.result = Preconditions.checkNotNull(result, "result");
+  }
+
   @Override
   public String toString() {
-    return "ConnectionHandshakeEvent{"
+    return "ConnectionHandshakeEventImpl{"
         + "connection=" + connection
+        + ", originalHostname='" + originalHostname + '\''
+        + ", currentHostname='" + currentHostname + '\''
+        + ", result=" + result
+        + ", currentRemoteAddress=" + currentRemoteAddress
         + '}';
   }
 }

proxy/build.gradle

@@ -75,7 +75,7 @@ dependencies {
     runtimeOnly 'com.lmax:disruptor:3.4.2' // Async loggers
 
     implementation 'it.unimi.dsi:fastutil:8.4.1'
-    implementation 'net.kyori:adventure-nbt:4.0.0-SNAPSHOT'
+    implementation "net.kyori:adventure-nbt:${adventureVersion}"
 
     implementation 'org.asynchttpclient:async-http-client:2.12.1'
 

proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java

@@ -75,7 +75,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
   private static final Logger logger = LogManager.getLogger(MinecraftConnection.class);
 
   private final Channel channel;
-  private SocketAddress remoteAddress;
+  private @Nullable SocketAddress remoteAddress;
   private ProtocolRegistry state;
   private @Nullable MinecraftSessionHandler sessionHandler;
   private ProtocolVersion protocolVersion;
@@ -304,10 +304,16 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
     return !channel.isActive();
   }
 
-  public SocketAddress getRemoteAddress() {
+  public @Nullable SocketAddress getRemoteAddress() {
     return remoteAddress;
   }
 
+  public void setRemoteAddress(@Nullable SocketAddress address) {
+    ensureOpen();
+    ensureInEventLoop();
+    this.remoteAddress = address;
+  }
+
   public ProtocolRegistry getState() {
     return state;
   }

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

@@ -39,6 +39,7 @@ import com.velocitypowered.proxy.network.registry.state.ProtocolStates;
 import io.netty.buffer.ByteBuf;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.util.Objects;
 import java.util.Optional;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
@@ -121,12 +122,6 @@ public class HandshakeSessionHandler implements MinecraftSessionHandler {
       return;
     }
 
-    InetAddress address = ((InetSocketAddress) connection.getRemoteAddress()).getAddress();
-    if (!server.getIpAttemptLimiter().attempt(address)) {
-      ic.disconnectQuietly(Component.translatable("velocity.error.logging-in-too-fast"));
-      return;
-    }
-
     connection.setType(getHandshakeConnectionType(handshake));
 
     // If the proxy is configured for modern forwarding, we must deny connections from 1.12.2
@@ -138,8 +133,36 @@ public class HandshakeSessionHandler implements MinecraftSessionHandler {
       return;
     }
 
-    server.eventManager().fireAndForget(new ConnectionHandshakeEventImpl(ic));
-    connection.setSessionHandler(new LoginSessionHandler(server, connection, ic));
+    connection.setAutoReading(false);
+    server.eventManager().fire(new ConnectionHandshakeEventImpl(ic, handshake.getServerAddress()))
+        .thenAcceptAsync(event -> {
+          connection.setAutoReading(true);
+
+          if (!event.result().isAllowed()) {
+            ic.disconnectQuietly(event.result().reason().get());
+          } else {
+            // if the handshake is changed, propagate the change
+            if (!event.currentHostname().equals(event.originalHostname())) {
+              ic.setCleanedHostname(cleanVhost(event.currentHostname()));
+            }
+
+            if (!Objects.equals(event.currentRemoteHostAddress(), ic.remoteAddress())) {
+              ic.setRemoteAddress(event.currentRemoteHostAddress());
+            }
+
+            if (connection.getRemoteAddress() instanceof InetSocketAddress) {
+              InetAddress address = ((InetSocketAddress) connection.getRemoteAddress())
+                  .getAddress();
+              if (!server.getIpAttemptLimiter().attempt(address)) {
+                ic.disconnectQuietly(
+                    Component.translatable("velocity.error.logging-in-too-fast"));
+                return;
+              }
+            }
+
+            connection.setSessionHandler(new LoginSessionHandler(server, connection, ic));
+          }
+        }, connection.eventLoop());
   }
 
   private ConnectionType getHandshakeConnectionType(ServerboundHandshakePacket handshake) {

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

@@ -17,6 +17,7 @@
 
 package com.velocitypowered.proxy.connection.client;
 
+import com.google.common.base.Preconditions;
 import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.api.proxy.connection.InboundConnection;
 import com.velocitypowered.proxy.connection.MinecraftConnection;
@@ -24,13 +25,16 @@ import com.velocitypowered.proxy.connection.MinecraftConnectionAssociation;
 import com.velocitypowered.proxy.network.packet.clientbound.ClientboundDisconnectPacket;
 import com.velocitypowered.proxy.network.packet.serverbound.ServerboundHandshakePacket;
 import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.util.Locale;
+import java.util.Objects;
 import java.util.Optional;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.serializer.legacy.LegacyComponentSerializer;
 import net.kyori.adventure.translation.GlobalTranslator;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public final class InitialInboundConnection implements InboundConnection,
     MinecraftConnectionAssociation {
@@ -38,24 +42,24 @@ public final class InitialInboundConnection implements InboundConnection,
   private static final Logger logger = LogManager.getLogger(InitialInboundConnection.class);
 
   private final MinecraftConnection connection;
-  private final String cleanedAddress;
+  private String cleanedHostname;
   private final ServerboundHandshakePacket handshake;
 
-  InitialInboundConnection(MinecraftConnection connection, String cleanedAddress,
+  InitialInboundConnection(MinecraftConnection connection, String cleanedHostname,
       ServerboundHandshakePacket handshake) {
     this.connection = connection;
-    this.cleanedAddress = cleanedAddress;
+    this.cleanedHostname = cleanedHostname;
     this.handshake = handshake;
   }
 
   @Override
-  public InetSocketAddress remoteAddress() {
-    return (InetSocketAddress) connection.getRemoteAddress();
+  public SocketAddress remoteAddress() {
+    return connection.getRemoteAddress();
   }
 
   @Override
   public Optional<InetSocketAddress> connectedHostname() {
-    return Optional.of(InetSocketAddress.createUnresolved(cleanedAddress, handshake.getPort()));
+    return Optional.of(InetSocketAddress.createUnresolved(cleanedHostname, handshake.getPort()));
   }
 
   @Override
@@ -70,7 +74,11 @@ public final class InitialInboundConnection implements InboundConnection,
 
   @Override
   public String toString() {
-    return "[initial connection] " + connection.getRemoteAddress().toString();
+    return "[initial connection] " + connection.getRemoteAddress();
+  }
+
+  public void setCleanedHostname(String hostname) {
+    this.cleanedHostname = Preconditions.checkNotNull(hostname, "hostname");
   }
 
   /**
@@ -93,4 +101,8 @@ public final class InitialInboundConnection implements InboundConnection,
     Component translated = GlobalTranslator.render(reason, Locale.getDefault());
     connection.closeWith(ClientboundDisconnectPacket.create(translated, protocolVersion()));
   }
+
+  public void setRemoteAddress(@Nullable SocketAddress currentRemoteHostAddress) {
+    connection.setRemoteAddress(currentRemoteHostAddress);
+  }
 }

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

@@ -54,6 +54,7 @@ import com.velocitypowered.proxy.network.packet.serverbound.ServerboundServerLog
 import com.velocitypowered.proxy.network.registry.state.ProtocolStates;
 import io.netty.buffer.ByteBuf;
 import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.security.GeneralSecurityException;
 import java.security.KeyPair;
 import java.security.MessageDigest;
@@ -119,12 +120,15 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
       byte[] decryptedSharedSecret = decryptRsa(serverKeyPair, packet.getSharedSecret());
       String serverId = generateServerId(decryptedSharedSecret, serverKeyPair.getPublic());
 
-      String playerIp = ((InetSocketAddress) mcConnection.getRemoteAddress()).getHostString();
       String url = String.format(MOJANG_HASJOINED_URL,
           urlFormParameterEscaper().escape(login.getUsername()), serverId);
 
       if (server.configuration().shouldPreventClientProxyConnections()) {
-        url += "&ip=" + urlFormParameterEscaper().escape(playerIp);
+        SocketAddress playerRemoteAddress = mcConnection.getRemoteAddress();
+        if (playerRemoteAddress instanceof InetSocketAddress) {
+          url += "&ip=" + urlFormParameterEscaper().escape(
+              ((InetSocketAddress) playerRemoteAddress).getHostString());
+        }
       }
 
       ListenableFuture<Response> hasJoinedResponse = server.getAsyncHttpClient().prepareGet(url)
@@ -156,8 +160,8 @@ public class LoginSessionHandler implements MinecraftSessionHandler {
           } else {
             // Something else went wrong
             logger.error(
-                "Got an unexpected error code {} whilst contacting Mojang to log in {} ({})",
-                profileResponse.getStatusCode(), login.getUsername(), playerIp);
+                "Got an unexpected error code {} whilst contacting Mojang to log in {}",
+                profileResponse.getStatusCode(), login.getUsername());
             mcConnection.close(true);
           }
         } catch (ExecutionException e) {