diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java index 74b1d6730..030de405d 100644 --- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java +++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java @@ -50,7 +50,7 @@ import org.checkerframework.checker.nullness.qual.Nullable; public class ClientPlaySessionHandler implements MinecraftSessionHandler { private static final Logger logger = LogManager.getLogger(ClientPlaySessionHandler.class); - private static final int MAX_PLUGIN_CHANNELS = 1024; + static final int MAX_PLUGIN_CHANNELS = 1024; private final ConnectedPlayer player; private boolean spawned = false; @@ -63,6 +63,12 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler { public ClientPlaySessionHandler(VelocityServer server, ConnectedPlayer player) { this.player = player; this.server = server; + + if (player.getMinecraftConnection().getSessionHandler() + instanceof InitialConnectSessionHandler) { + this.knownChannels.addAll(((InitialConnectSessionHandler) player.getMinecraftConnection() + .getSessionHandler()).getKnownChannels()); + } } @Override diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java index 998e01cc5..33d78cd30 100644 --- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java +++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java @@ -1,22 +1,56 @@ package com.velocitypowered.proxy.connection.client; +import static com.velocitypowered.proxy.connection.client.ClientPlaySessionHandler.MAX_PLUGIN_CHANNELS; + import com.velocitypowered.proxy.connection.MinecraftSessionHandler; import com.velocitypowered.proxy.connection.backend.VelocityServerConnection; import com.velocitypowered.proxy.protocol.packet.PluginMessage; +import com.velocitypowered.proxy.protocol.util.PluginMessageUtil; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; public class InitialConnectSessionHandler implements MinecraftSessionHandler { private final ConnectedPlayer player; + private final Set knownChannels; InitialConnectSessionHandler(ConnectedPlayer player) { this.player = player; + this.knownChannels = new HashSet<>(); } @Override public boolean handle(PluginMessage packet) { VelocityServerConnection serverConn = player.getConnectionInFlight(); if (serverConn != null) { - if (!player.getPhase().handle(player, packet, serverConn)) { + if (player.getPhase().handle(player, packet, serverConn)) { + return true; + } + + if (PluginMessageUtil.isRegister(packet)) { + List actuallyRegistered = new ArrayList<>(); + List channels = PluginMessageUtil.getChannels(packet); + for (String channel : channels) { + if (knownChannels.size() >= MAX_PLUGIN_CHANNELS && !knownChannels.contains(channel)) { + throw new IllegalStateException("Too many plugin message channels registered"); + } + if (knownChannels.add(channel)) { + actuallyRegistered.add(channel); + } + } + + if (!actuallyRegistered.isEmpty()) { + PluginMessage newRegisterPacket = PluginMessageUtil.constructChannelsPacket(serverConn + .ensureConnected().getProtocolVersion(), actuallyRegistered); + serverConn.ensureConnected().write(newRegisterPacket); + } + } else if (PluginMessageUtil.isUnregister(packet)) { + List channels = PluginMessageUtil.getChannels(packet); + knownChannels.removeAll(channels); + serverConn.ensureConnected().write(packet); + } else { serverConn.ensureConnected().write(packet); } } @@ -28,4 +62,8 @@ public class InitialConnectSessionHandler implements MinecraftSessionHandler { // the user cancelled the login process player.teardown(); } + + Set getKnownChannels() { + return knownChannels; + } }