From 6104354fa1a87a8a09d852ef8ec388611a9b51f1 Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Thu, 22 Dec 2016 22:29:00 +1100
Subject: [PATCH] Implement basic entity metadata mapping to remap fireworks
 for boosting.

---
 proxy/pom.xml                                 |  6 ++
 .../net/md_5/bungee/entitymap/EntityMap.java  | 89 +++++++++++++++++++
 .../md_5/bungee/entitymap/EntityMap_1_11.java |  4 +
 3 files changed, 99 insertions(+)

diff --git a/proxy/pom.xml b/proxy/pom.xml
index 703fbda3fe..2bb124e04c 100644
--- a/proxy/pom.xml
+++ b/proxy/pom.xml
@@ -19,6 +19,12 @@
     <description>Proxy component of the Elastic Portal Suite</description>
 
     <dependencies>
+        <dependency>
+            <groupId>com.flowpowered</groupId>
+            <artifactId>flow-nbt</artifactId>
+            <version>1.0.0</version>
+            <scope>compile</scope>
+        </dependency>
         <dependency>
             <groupId>io.netty</groupId>
             <artifactId>netty-codec-http</artifactId>
diff --git a/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap.java b/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap.java
index a8728bb3ec..e103cc0a69 100644
--- a/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap.java
+++ b/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap.java
@@ -1,7 +1,11 @@
 package net.md_5.bungee.entitymap;
 
+import com.flowpowered.nbt.stream.NBTInputStream;
+import com.google.common.base.Throwables;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufInputStream;
+import java.io.IOException;
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
 import net.md_5.bungee.protocol.DefinedPacket;
@@ -101,6 +105,91 @@ protected static void rewriteVarInt(ByteBuf packet, int oldId, int newId, int of
         }
     }
 
+    protected static void rewriteMetaVarInt(ByteBuf packet, int oldId, int newId, int metaIndex)
+    {
+        int readerIndex = packet.readerIndex();
+
+        short index;
+        while ( ( index = packet.readUnsignedByte() ) != 0xFF )
+        {
+            int type = DefinedPacket.readVarInt( packet );
+
+            switch ( type )
+            {
+                case 0:
+                    packet.skipBytes( 1 ); // byte
+                    break;
+                case 1:
+                    if ( index == metaIndex )
+                    {
+                        int position = packet.readerIndex();
+                        rewriteVarInt( packet, oldId, newId, position );
+                        packet.readerIndex( position );
+                    }
+                    DefinedPacket.readVarInt( packet );
+                    break;
+                case 2:
+                    packet.skipBytes( 4 ); // float
+                    break;
+                case 3:
+                case 4:
+                    DefinedPacket.readString( packet );
+                    break;
+                case 5:
+                    if ( packet.readShort() != -1 )
+                    {
+                        packet.skipBytes( 3 ); // byte, short
+
+                        int position = packet.readerIndex();
+                        if ( packet.readByte() != 0 )
+                        {
+                            packet.readerIndex( position );
+
+                            try
+                            {
+                                new NBTInputStream( new ByteBufInputStream( packet ), false ).readTag();
+                            } catch ( IOException ex )
+                            {
+                                throw Throwables.propagate( ex );
+                            }
+                        }
+                    }
+                    break;
+                case 6:
+                    packet.skipBytes( 1 ); // boolean
+                    break;
+                case 7:
+                    packet.skipBytes( 12 ); // float, float, float
+                    break;
+                case 8:
+                    packet.readLong();
+                    break;
+                case 9:
+                    if ( packet.readBoolean() )
+                    {
+                        packet.skipBytes( 8 ); // long
+                    }
+                    break;
+                case 10:
+                    DefinedPacket.readVarInt( packet );
+                    break;
+                case 11:
+                    if ( packet.readBoolean() )
+                    {
+                        packet.skipBytes( 16 ); // long, long
+                    }
+                    break;
+                case 12:
+                    DefinedPacket.readVarInt( packet );
+                    break;
+                default:
+                    throw new IllegalArgumentException( "Unknown meta type " + type );
+            }
+        }
+
+        packet.readerIndex( readerIndex );
+    }
+
     // Handles simple packets
     private static void rewrite(ByteBuf packet, int oldId, int newId, boolean[] ints, boolean[] varints)
     {
diff --git a/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap_1_11.java b/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap_1_11.java
index aa0fb7306e..260da17ef4 100644
--- a/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap_1_11.java
+++ b/proxy/src/main/java/net/md_5/bungee/entitymap/EntityMap_1_11.java
@@ -147,6 +147,10 @@ public void rewriteClientbound(ByteBuf packet, int oldId, int newId)
                     rewriteInt( packet, oldId, newId, packet.readerIndex() );
                 }
                 break;
+            case 0x39 /* EntityMetadata : PacketPlayOutEntityMetadata */:
+                DefinedPacket.readVarInt( packet ); // Entity ID
+                rewriteMetaVarInt( packet, oldId, newId, 7 ); // fireworks (et al)
+                break;
         }
         packet.readerIndex( readerIndex );
     }