Merge pull request #4 from kashike/protocol-storage

Don't search through protocol versions all the time

src/main/java/com/velocitypowered/network/ConnectionManager.java

@@ -96,8 +96,8 @@ public final class ConnectionManager {
                                 .addLast(FRAME_DECODER, new MinecraftVarintFrameDecoder())
                                 .addLast(LEGACY_PING_ENCODER, LegacyPingEncoder.INSTANCE)
                                 .addLast(FRAME_ENCODER, MinecraftVarintLengthEncoder.INSTANCE)
-                                .addLast(MINECRAFT_DECODER, new MinecraftDecoder(ProtocolConstants.Direction.TO_SERVER))
+                                .addLast(MINECRAFT_DECODER, new MinecraftDecoder(ProtocolConstants.Direction.SERVERBOUND))
-                                .addLast(MINECRAFT_ENCODER, new MinecraftEncoder(ProtocolConstants.Direction.TO_CLIENT));
+                                .addLast(MINECRAFT_ENCODER, new MinecraftEncoder(ProtocolConstants.Direction.CLIENTBOUND));
 
                         final MinecraftConnection connection = new MinecraftConnection(ch);
                         connection.setState(StateRegistry.HANDSHAKE);

src/main/java/com/velocitypowered/proxy/connection/backend/ServerConnection.java

@@ -48,8 +48,8 @@ public class ServerConnection implements MinecraftConnectionAssociation {
                                 .addLast(READ_TIMEOUT, new ReadTimeoutHandler(SERVER_READ_TIMEOUT_SECONDS, TimeUnit.SECONDS))
                                 .addLast(FRAME_DECODER, new MinecraftVarintFrameDecoder())
                                 .addLast(FRAME_ENCODER, MinecraftVarintLengthEncoder.INSTANCE)
-                                .addLast(MINECRAFT_DECODER, new MinecraftDecoder(ProtocolConstants.Direction.TO_CLIENT))
+                                .addLast(MINECRAFT_DECODER, new MinecraftDecoder(ProtocolConstants.Direction.CLIENTBOUND))
-                                .addLast(MINECRAFT_ENCODER, new MinecraftEncoder(ProtocolConstants.Direction.TO_SERVER));
+                                .addLast(MINECRAFT_ENCODER, new MinecraftEncoder(ProtocolConstants.Direction.SERVERBOUND));
 
                         MinecraftConnection connection = new MinecraftConnection(ch);
                         connection.setState(StateRegistry.HANDSHAKE);

src/main/java/com/velocitypowered/proxy/protocol/ProtocolConstants.java

@@ -14,7 +14,7 @@ public enum ProtocolConstants { ;
     }
 
     public enum Direction {
-        TO_SERVER,
+        SERVERBOUND,
-        TO_CLIENT
+        CLIENTBOUND
     }
 }

src/main/java/com/velocitypowered/proxy/protocol/StateRegistry.java

@@ -1,13 +1,10 @@
 package com.velocitypowered.proxy.protocol;
 
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
 import com.velocitypowered.proxy.protocol.packets.*;
 import io.netty.util.collection.IntObjectHashMap;
 import io.netty.util.collection.IntObjectMap;
 
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Supplier;
@@ -17,124 +14,127 @@ import static com.velocitypowered.proxy.protocol.ProtocolConstants.MINECRAFT_1_1
 public enum StateRegistry {
     HANDSHAKE {
         {
-            TO_SERVER.register(Handshake.class, Handshake::new,
+            SERVERBOUND.register(Handshake.class, Handshake::new,
                     generic(0x00));
         }
     },
     STATUS {
         {
-            TO_SERVER.register(StatusRequest.class, StatusRequest::new,
+            SERVERBOUND.register(StatusRequest.class, StatusRequest::new,
                     generic(0x00));
-            TO_SERVER.register(Ping.class, Ping::new,
+            SERVERBOUND.register(Ping.class, Ping::new,
                     generic(0x01));
 
-            TO_CLIENT.register(StatusResponse.class, StatusResponse::new,
+            CLIENTBOUND.register(StatusResponse.class, StatusResponse::new,
                     generic(0x00));
-            TO_CLIENT.register(Ping.class, Ping::new,
+            CLIENTBOUND.register(Ping.class, Ping::new,
                     generic(0x01));
         }
     },
     PLAY {
         {
-            TO_SERVER.register(Chat.class, Chat::new,
+            SERVERBOUND.register(Chat.class, Chat::new,
                     map(0x02, MINECRAFT_1_12));
-            TO_SERVER.register(Ping.class, Ping::new,
+            SERVERBOUND.register(Ping.class, Ping::new,
                     map(0x0b, MINECRAFT_1_12));
 
-            TO_CLIENT.register(Chat.class, Chat::new,
+            CLIENTBOUND.register(Chat.class, Chat::new,
                     map(0x0F, MINECRAFT_1_12));
-            TO_CLIENT.register(Disconnect.class, Disconnect::new,
+            CLIENTBOUND.register(Disconnect.class, Disconnect::new,
                     map(0x1A, MINECRAFT_1_12));
-            TO_CLIENT.register(Ping.class, Ping::new,
+            CLIENTBOUND.register(Ping.class, Ping::new,
                     map(0x1F, MINECRAFT_1_12));
-            TO_CLIENT.register(JoinGame.class, JoinGame::new,
+            CLIENTBOUND.register(JoinGame.class, JoinGame::new,
                     map(0x23, MINECRAFT_1_12));
-            TO_CLIENT.register(Respawn.class, Respawn::new,
+            CLIENTBOUND.register(Respawn.class, Respawn::new,
                     map(0x35, MINECRAFT_1_12));
         }
     },
     LOGIN {
         {
-            TO_SERVER.register(ServerLogin.class, ServerLogin::new,
+            SERVERBOUND.register(ServerLogin.class, ServerLogin::new,
                     generic(0x00));
-            TO_SERVER.register(EncryptionResponse.class, EncryptionResponse::new,
+            SERVERBOUND.register(EncryptionResponse.class, EncryptionResponse::new,
                     generic(0x01));
 
-            TO_CLIENT.register(Disconnect.class, Disconnect::new,
+            CLIENTBOUND.register(Disconnect.class, Disconnect::new,
                     generic(0x00));
-            TO_CLIENT.register(EncryptionRequest.class, EncryptionRequest::new,
+            CLIENTBOUND.register(EncryptionRequest.class, EncryptionRequest::new,
                     generic(0x01));
-            TO_CLIENT.register(ServerLoginSuccess.class, ServerLoginSuccess::new,
+            CLIENTBOUND.register(ServerLoginSuccess.class, ServerLoginSuccess::new,
                     generic(0x02));
-            TO_CLIENT.register(SetCompression.class, SetCompression::new,
+            CLIENTBOUND.register(SetCompression.class, SetCompression::new,
                     generic(0x03));
         }
     };
 
-    public final PacketRegistry TO_CLIENT = new PacketRegistry(ProtocolConstants.Direction.TO_CLIENT, this);
+    public final PacketRegistry CLIENTBOUND = new PacketRegistry(ProtocolConstants.Direction.CLIENTBOUND);
-    public final PacketRegistry TO_SERVER = new PacketRegistry(ProtocolConstants.Direction.TO_SERVER, this);
+    public final PacketRegistry SERVERBOUND = new PacketRegistry(ProtocolConstants.Direction.SERVERBOUND);
 
     public static class PacketRegistry {
         private final ProtocolConstants.Direction direction;
-        private final StateRegistry state;
+        private final IntObjectMap<ProtocolVersion> versions = new IntObjectHashMap<>();
-        private final IntObjectMap<IntObjectMap<Supplier<? extends MinecraftPacket>>> byProtocolVersionToProtocolIds = new IntObjectHashMap<>();
-        private final Map<Class<? extends MinecraftPacket>, List<PacketMapping>> idMappers = new HashMap<>();
 
-        public PacketRegistry(ProtocolConstants.Direction direction, StateRegistry state) {
+        public PacketRegistry(ProtocolConstants.Direction direction) {
             this.direction = direction;
-            this.state = state;
+        }
+
+        public ProtocolVersion getVersion(final int version) {
+            ProtocolVersion result = null;
+            for (final IntObjectMap.PrimitiveEntry<ProtocolVersion> entry : this.versions.entries()) {
+                if (entry.key() <= version) {
+                    result = entry.value();
+                }
+            }
+            if (result == null) {
+                throw new IllegalArgumentException("Could not find data for protocol version " + version);
+            }
+            return result;
         }
 
         public <P extends MinecraftPacket> void register(Class<P> clazz, Supplier<P> packetSupplier, PacketMapping... mappings) {
             if (mappings.length == 0) {
                 throw new IllegalArgumentException("At least one mapping must be provided.");
             }
-            for (PacketMapping mapping : mappings) {
+
-                IntObjectMap<Supplier<? extends MinecraftPacket>> ids = byProtocolVersionToProtocolIds.get(mapping.protocolVersion);
+            for (final PacketMapping mapping : mappings) {
-                if (ids == null) {
+                ProtocolVersion version = this.versions.get(mapping.protocolVersion);
-                    byProtocolVersionToProtocolIds.put(mapping.protocolVersion, ids = new IntObjectHashMap<>());
+                if (version == null) {
+                    version = new ProtocolVersion(mapping.protocolVersion);
+                    this.versions.put(mapping.protocolVersion, version);
                 }
-                ids.put(mapping.id, packetSupplier);
+                version.packetIdToSupplier.put(mapping.id, packetSupplier);
+                version.packetClassToId.put(clazz, mapping.id);
             }
-            idMappers.put(clazz, ImmutableList.copyOf(mappings));
         }
 
-        public MinecraftPacket createPacket(int id, int protocolVersion) {
+        public class ProtocolVersion {
-            IntObjectMap<Supplier<? extends MinecraftPacket>> bestLookup = null;
+            public final int id;
-            for (IntObjectMap.PrimitiveEntry<IntObjectMap<Supplier<? extends MinecraftPacket>>> entry : byProtocolVersionToProtocolIds.entries()) {
+            final IntObjectMap<Supplier<? extends MinecraftPacket>> packetIdToSupplier = new IntObjectHashMap<>();
-                if (entry.key() <= protocolVersion) {
+            final Map<Class<? extends MinecraftPacket>, Integer> packetClassToId = new HashMap<>();
-                    bestLookup = entry.value();
-                }
-            }
-            if (bestLookup == null) {
-                return null;
-            }
-            Supplier<? extends MinecraftPacket> supplier = bestLookup.get(id);
-            if (supplier == null) {
-                return null;
-            }
-            return supplier.get();
-        }
 
-        public int getId(MinecraftPacket packet, int protocolVersion) {
+            ProtocolVersion(final int id) {
-            Preconditions.checkNotNull(packet, "packet");
+                this.id = id;
+            }
 
-            List<PacketMapping> mappings = idMappers.get(packet.getClass());
+            public MinecraftPacket createPacket(final int id) {
-            if (mappings == null || mappings.isEmpty()) {
+                final Supplier<? extends MinecraftPacket> supplier = this.packetIdToSupplier.get(id);
-                throw new IllegalArgumentException("Supplied packet " + packet.getClass().getName() +
+                if (supplier == null) {
-                        " doesn't have any mappings. Direction " + direction + " State " + state);
+                    return null;
-            }
-            int useId = -1;
-            for (PacketMapping mapping : mappings) {
-                if (mapping.protocolVersion <= protocolVersion) {
-                    useId = mapping.id;
                 }
+                return supplier.get();
             }
-            if (useId == -1) {
+
-                throw new IllegalArgumentException("Unable to find a mapping for " + packet.getClass().getName()
+            public int getPacketId(final MinecraftPacket packet) {
-                        + " Version " + protocolVersion + " Direction " + direction + " State " + state);
+                final Integer id = this.packetClassToId.get(packet.getClass());
+                if (id == null) {
+                    throw new IllegalArgumentException(String.format(
+                            "Unable to find id for packet of type %s in %s protocol %s",
+                            packet.getClass().getName(), PacketRegistry.this.direction, this.id
+                    ));
+                }
+                return id;
             }
-            return useId;
         }
     }
 

src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftDecoder.java

@@ -12,7 +12,7 @@ import java.util.List;
 public class MinecraftDecoder extends MessageToMessageDecoder<ByteBuf> {
     private StateRegistry state;
     private final ProtocolConstants.Direction direction;
-    private int protocolVersion;
+    private StateRegistry.PacketRegistry.ProtocolVersion protocolVersion;
 
     public MinecraftDecoder(ProtocolConstants.Direction direction) {
         this.state = StateRegistry.HANDSHAKE;
@@ -28,14 +28,13 @@ public class MinecraftDecoder extends MessageToMessageDecoder<ByteBuf> {
         ByteBuf slice = msg.slice().retain();
 
         int packetId = ProtocolUtils.readVarInt(msg);
-        StateRegistry.PacketRegistry mappings = direction == ProtocolConstants.Direction.TO_CLIENT ? state.TO_CLIENT : state.TO_SERVER;
+        MinecraftPacket packet = this.protocolVersion.createPacket(packetId);
-        MinecraftPacket packet = mappings.createPacket(packetId, protocolVersion);
         if (packet == null) {
             msg.skipBytes(msg.readableBytes());
             out.add(new PacketWrapper(null, slice));
         } else {
             try {
-                packet.decode(msg, direction, protocolVersion);
+                packet.decode(msg, direction, protocolVersion.id);
             } catch (Exception e) {
                 throw new CorruptedFrameException("Error decoding " + packet.getClass() + " Direction " + direction
                         + " Protocol " + protocolVersion + " State " + state + " ID " + Integer.toHexString(packetId), e);
@@ -44,12 +43,12 @@ public class MinecraftDecoder extends MessageToMessageDecoder<ByteBuf> {
         }
     }
 
-    public int getProtocolVersion() {
+    public StateRegistry.PacketRegistry.ProtocolVersion getProtocolVersion() {
         return protocolVersion;
     }
 
     public void setProtocolVersion(int protocolVersion) {
-        this.protocolVersion = protocolVersion;
+        this.protocolVersion = (this.direction == ProtocolConstants.Direction.CLIENTBOUND ? this.state.CLIENTBOUND : this.state.SERVERBOUND).getVersion(protocolVersion);
     }
 
     public StateRegistry getState() {

src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftEncoder.java

@@ -12,7 +12,7 @@ import io.netty.handler.codec.MessageToByteEncoder;
 public class MinecraftEncoder extends MessageToByteEncoder<MinecraftPacket> {
     private StateRegistry state;
     private final ProtocolConstants.Direction direction;
-    private int protocolVersion;
+    private StateRegistry.PacketRegistry.ProtocolVersion protocolVersion;
 
     public MinecraftEncoder(ProtocolConstants.Direction direction) {
         this.state = StateRegistry.HANDSHAKE;
@@ -20,19 +20,18 @@ public class MinecraftEncoder extends MessageToByteEncoder<MinecraftPacket> {
     }
 
     @Override
-    protected void encode(ChannelHandlerContext ctx, MinecraftPacket msg, ByteBuf out) throws Exception {
+    protected void encode(ChannelHandlerContext ctx, MinecraftPacket msg, ByteBuf out) {
-        StateRegistry.PacketRegistry mappings = direction == ProtocolConstants.Direction.TO_CLIENT ? state.TO_CLIENT : state.TO_SERVER;
+        int packetId = this.protocolVersion.getPacketId(msg);
-        int packetId = mappings.getId(msg, protocolVersion);
         ProtocolUtils.writeVarInt(out, packetId);
-        msg.encode(out, direction, protocolVersion);
+        msg.encode(out, direction, protocolVersion.id);
     }
 
-    public int getProtocolVersion() {
+    public StateRegistry.PacketRegistry.ProtocolVersion getProtocolVersion() {
         return protocolVersion;
     }
 
-    public void setProtocolVersion(int protocolVersion) {
+    public void setProtocolVersion(final int protocolVersion) {
-        this.protocolVersion = protocolVersion;
+        this.protocolVersion = (this.direction == ProtocolConstants.Direction.CLIENTBOUND ? this.state.CLIENTBOUND : this.state.SERVERBOUND).getVersion(protocolVersion);
     }
 
     public StateRegistry getState() {

src/main/java/com/velocitypowered/proxy/protocol/packets/Chat.java

@@ -47,7 +47,7 @@ public class Chat implements MinecraftPacket {
     @Override
     public void decode(ByteBuf buf, ProtocolConstants.Direction direction, int protocolVersion) {
         message = ProtocolUtils.readString(buf);
-        if (direction == ProtocolConstants.Direction.TO_CLIENT) {
+        if (direction == ProtocolConstants.Direction.CLIENTBOUND) {
             position = buf.readByte();
         }
     }
@@ -55,7 +55,7 @@ public class Chat implements MinecraftPacket {
     @Override
     public void encode(ByteBuf buf, ProtocolConstants.Direction direction, int protocolVersion) {
         ProtocolUtils.writeString(buf, message);
-        if (direction == ProtocolConstants.Direction.TO_CLIENT) {
+        if (direction == ProtocolConstants.Direction.CLIENTBOUND) {
             buf.writeByte(position);
         }
     }

src/test/java/com/velocitypowered/proxy/protocol/PacketRegistryTest.java

@@ -8,7 +8,7 @@ import static org.junit.jupiter.api.Assertions.*;
 
 class PacketRegistryTest {
     private StateRegistry.PacketRegistry setupRegistry() {
-        StateRegistry.PacketRegistry registry = new StateRegistry.PacketRegistry(ProtocolConstants.Direction.TO_CLIENT, StateRegistry.HANDSHAKE);
+        StateRegistry.PacketRegistry registry = new StateRegistry.PacketRegistry(ProtocolConstants.Direction.CLIENTBOUND);
         registry.register(Handshake.class, Handshake::new, new StateRegistry.PacketMapping(0x00, 1));
         registry.register(Ping.class, Ping::new, new StateRegistry.PacketMapping(0x01, 1),
                 new StateRegistry.PacketMapping(0x02, 5));
@@ -18,47 +18,47 @@ class PacketRegistryTest {
     @Test
     void packetRegistryWorks() {
         StateRegistry.PacketRegistry registry = setupRegistry();
-        MinecraftPacket packet = registry.createPacket(0, 1);
+        MinecraftPacket packet = registry.getVersion(1).createPacket(0);
         assertNotNull(packet, "Packet was not found in registry");
         assertEquals(Handshake.class, packet.getClass(), "Registry returned wrong class");
 
-        assertEquals(0, registry.getId(packet, 1), "Registry did not return the correct packet ID");
+        assertEquals(0, registry.getVersion(1).getPacketId(packet), "Registry did not return the correct packet ID");
     }
 
     @Test
     void packetRegistryRevertsToBestOldVersion() {
         StateRegistry.PacketRegistry registry = setupRegistry();
-        MinecraftPacket packet = registry.createPacket(0, 2);
+        MinecraftPacket packet = registry.getVersion(2).createPacket(0);
         assertNotNull(packet, "Packet was not found in registry");
         assertEquals(Handshake.class, packet.getClass(), "Registry returned wrong class");
 
-        assertEquals(0, registry.getId(packet, 2), "Registry did not return the correct packet ID");
+        assertEquals(0, registry.getVersion(2).getPacketId(packet), "Registry did not return the correct packet ID");
     }
 
     @Test
     void packetRegistryDoesntProvideNewPacketsForOld() {
         StateRegistry.PacketRegistry registry = setupRegistry();
-        assertNull(registry.createPacket(0, 0), "Packet was found in registry despite being too new");
+        assertNull(registry.getVersion(0).createPacket(0), "Packet was found in registry despite being too new");
 
-        assertThrows(IllegalArgumentException.class, () -> registry.getId(new Handshake(), 0), "Registry provided new packets for an old protocol version");
+        assertThrows(IllegalArgumentException.class, () -> registry.getVersion(0).getPacketId(new Handshake()), "Registry provided new packets for an old protocol version");
     }
 
     @Test
     void failOnNoMappings() {
-        StateRegistry.PacketRegistry registry = new StateRegistry.PacketRegistry(ProtocolConstants.Direction.TO_CLIENT, StateRegistry.HANDSHAKE);
+        StateRegistry.PacketRegistry registry = new StateRegistry.PacketRegistry(ProtocolConstants.Direction.CLIENTBOUND);
         assertThrows(IllegalArgumentException.class, () -> registry.register(Handshake.class, Handshake::new));
-        assertThrows(IllegalArgumentException.class, () -> registry.getId(new Handshake(), 0));
+        assertThrows(IllegalArgumentException.class, () -> registry.getVersion(0).getPacketId(new Handshake()));
     }
 
     @Test
     void packetRegistryProvidesCorrectVersionsForMultipleMappings() {
         StateRegistry.PacketRegistry registry = setupRegistry();
-        assertNotNull(registry.createPacket(1, 1), "Packet was not found in registry despite being being registered with ID 1 and version 1");
+        assertNotNull(registry.getVersion(1).createPacket(1), "Packet was not found in registry despite being being registered with ID 1 and version 1");
-        assertNotNull(registry.createPacket(1, 2), "Packet was not found in registry despite being being registered with ID 1 and version 1 (we are looking up version 2)");
+        assertNotNull(registry.getVersion(2).createPacket(1), "Packet was not found in registry despite being being registered with ID 1 and version 1 (we are looking up version 2)");
-        assertNotNull(registry.createPacket(2, 5), "Packet was not found in registry despite being being registered with ID 2 and version 5");
+        assertNotNull(registry.getVersion(5).createPacket(2), "Packet was not found in registry despite being being registered with ID 2 and version 5");
-        assertNotNull(registry.createPacket(2, 6), "Packet was not found in registry despite being being registered with ID 2 and version 5 (we are looking up version 6)");
+        assertNotNull(registry.getVersion(6).createPacket(2), "Packet was not found in registry despite being being registered with ID 2 and version 5 (we are looking up version 6)");
 
-        assertEquals(1, registry.getId(new Ping(), 1), "Wrong ID provided from registry");
+        assertEquals(1, registry.getVersion(1).getPacketId(new Ping()), "Wrong ID provided from registry");
-        assertEquals(2, registry.getId(new Ping(), 5), "Wrong ID provided from registry");
+        assertEquals(2, registry.getVersion(5).getPacketId(new Ping()), "Wrong ID provided from registry");
     }
 }