Add packet api to change out the packet during listening

Signed-off-by: Irmo van den Berge <irmo.vandenberge@ziggo.nl>

src/main/java/com/bergerkiller/bukkit/common/events/PacketEvent.java

@@ -3,6 +3,7 @@ package com.bergerkiller.bukkit.common.events;
 import com.bergerkiller.bukkit.common.protocol.CommonPacket;
 import com.bergerkiller.bukkit.common.protocol.PacketType;
 
+import com.bergerkiller.generated.net.minecraft.network.protocol.PacketHandle;
 import org.bukkit.entity.Player;
 import org.bukkit.event.Cancellable;
 
@@ -56,4 +57,30 @@ public abstract class PacketEvent implements Cancellable {
     public CommonPacket getPacket() {
         return packet;
     }
+
+    /**
+     * Switches out the current packet with a new one. Can be used when immutable packets
+     * need to be changed before being sent to the player / received by the server.
+     *
+     * @param packet New packet to swap it out with. Not null.
+     */
+    public void setPacket(CommonPacket packet) {
+        if (packet == null) {
+            throw new IllegalArgumentException("Cannot set null packet");
+        }
+        this.packet = packet;
+    }
+
+    /**
+     * Switches out the current packet with a new one. Can be used when immutable packets
+     * need to be changed before being sent to the player / received by the server.
+     *
+     * @param packet New packet to swap it out with. Not null.
+     */
+    public void setPacket(PacketHandle packet) {
+        if (packet == null) {
+            throw new IllegalArgumentException("Cannot set null packet");
+        }
+        this.packet = packet.toCommonPacket();
+    }
 }

src/main/java/com/bergerkiller/bukkit/common/internal/network/CommonPacketHandler.java

@@ -128,15 +128,17 @@ public class CommonPacketHandler extends PacketHandlerHooked {
 
         @Override
         public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
-            if (handler.handlePacketSend(player, msg, false)) {
-                super.write(ctx, msg, promise);
+            PacketHandlerRegistration.HandlerResult result = handler.handlePacketSend(player, msg, false);
+            if (!result.isCancelled) {
+                super.write(ctx, result.packet, promise);
             }
         }
 
         @Override
         public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-            if (handler.handlePacketReceive(player, msg, false)) {
-                super.channelRead(ctx, msg);
+            PacketHandlerRegistration.HandlerResult result = handler.handlePacketReceive(player, msg, false);
+            if (!result.isCancelled) {
+                super.channelRead(ctx, result.packet);
             }
         }
     }

src/main/java/com/bergerkiller/bukkit/common/internal/network/PacketHandlerHooked.java

@@ -88,16 +88,32 @@ public abstract class PacketHandlerHooked implements PacketHandler {
         type.preprocess(packet);
         SafeMethod<?> method = this.receiverMethods.get(packet);
         if (method == null) {
-        	Logging.LOGGER_NETWORK.log(Level.WARNING, "Could not find suitable packet handler for " + packet.getClass().getSimpleName());
-        } else {
-            PlayerConnectionHandle connection = PlayerConnectionHandle.forPlayer(player);
-            if (connection != null) {
-                // We are bypassing the hook - make sure to handle receive packet listener
-                if (this.handlePacketReceive(player, packet, false)) {
-                    method.invoke(connection.getRaw(), packet);
-                }
+            Logging.LOGGER_NETWORK.log(Level.WARNING, "Could not find suitable packet handler for " + packet.getClass().getSimpleName());
+            return;
+        }
+
+        PlayerConnectionHandle connection = PlayerConnectionHandle.forPlayer(player);
+        if (connection == null) {
+            return;
+        }
+
+        // We are bypassing the hook - make sure to handle receive packet listener
+        PacketHandlerRegistration.HandlerResult result = this.handlePacketReceive(player, packet, false);
+        if (result.isCancelled) {
+            return;
+        }
+
+        // Listeners could have changed the type of packet being received, then we
+        // need to find different receiver methods for it
+        if (result.packetType != type) {
+            method = this.receiverMethods.get(result.packet);
+            if (method == null) {
+                Logging.LOGGER_NETWORK.log(Level.WARNING, "Could not find suitable packet handler for listener-changed " + result.packet.getClass().getSimpleName());
+                return;
             }
         }
+
+        method.invoke(connection.getRaw(), result.packet);
     }
 
     @Override
@@ -167,11 +183,11 @@ public abstract class PacketHandlerHooked implements PacketHandler {
      * @param player for which the packet was meant
      * @param packet that is handled
      * @param wasCancelled - True if it was originally cancelled, False if not
-     * @return True if the packet is allowed to be sent, False if not
+     * @return Result of handling the packet
      */
-    public boolean handlePacketSend(Player player, Object packet, boolean wasCancelled) {
+    public PacketHandlerRegistration.HandlerResult handlePacketSend(Player player, Object packet, boolean wasCancelled) {
         if (player == null || !PacketHandle.T.isAssignableFrom(packet)) {
-            return true;
+            return new PacketHandlerRegistration.HandlerResult(packet, null, wasCancelled);
         }
 
         // Check if silent
@@ -190,9 +206,9 @@ public abstract class PacketHandlerHooked implements PacketHandler {
      * @param packet that is handled
      * @param wasCancelled - True if the packet is allowed to be received, False
      * if not
-     * @return True if the packet is allowed to be received, False if not
+     * @return Result of handling the packet
      */
-    public boolean handlePacketReceive(Player player, Object packet, boolean wasCancelled) {
+    public PacketHandlerRegistration.HandlerResult handlePacketReceive(Player player, Object packet, boolean wasCancelled) {
         return handlers.handlePacketReceive(player, packet, wasCancelled);
     }
 }

src/main/java/com/bergerkiller/bukkit/common/internal/network/PacketHandlerRegistration.java

@@ -176,7 +176,17 @@ class PacketHandlerRegistration {
         }
     }
 
-    public boolean handlePacketSend(Player player, PacketType packetType, Object packet, boolean is_silent, boolean wasCancelled) {
+    /**
+     * Handles when a packet is sent to a player
+     *
+     * @param player Player
+     * @param packetType Packet Type
+     * @param packet The packet
+     * @param is_silent Whether listeners should be skipped for this packet
+     * @param wasCancelled Whether it was originally cancelled
+     * @return Result of handling this packet
+     */
+    public HandlerResult handlePacketSend(Player player, PacketType packetType, Object packet, boolean is_silent, boolean wasCancelled) {
         if (packetType == PacketType.OUT_BUNDLE) {
             ClientboundBundlePacketHandle bundle = ClientboundBundlePacketHandle.createHandle(packet);
             if (is_silent) {
@@ -184,25 +194,48 @@ class PacketHandlerRegistration {
                 for (Object subPacket : bundle.subPackets()) {
                     handlePacketSendMonitor(player, packetType, subPacket);
                 }
-                return true;
+                return new HandlerResult(packet, packetType, false);
             } else {
                 // Filter the packets contained within the bundle packet. If all packets end up filtered,
                 // cancel the entire Bundle packet. We skip the monitors as every packet already passes
                 // by the monitor.
-                return bundle.filterSubPackets(p -> handlePacketSend(player, PacketType.getType(p), p, false, wasCancelled));
+                // TODO: Actually swap out the packets if changed by listener callbacks...
+                if (bundle.filterSubPackets(p -> handlePacketSend(player, PacketType.getType(p), p, false, wasCancelled) != null)) {
+                    return new HandlerResult(packet, packetType, false);
+                } else {
+                    return new HandlerResult(packet, packetType, true);
+                }
             }
         }
 
-        if (!is_silent && !handlePacketSendListener(player, packetType, packet, wasCancelled)) {
-            return false;
+        HandlerResult result;
+        if (is_silent) {
+            result = new HandlerResult(packet, packetType, false);
+        } else {
+            result = handlePacketSendListener(player, packetType, packet, wasCancelled);
+            if (result.isCancelled) {
+                return result;
+            } else {
+                packet = result.packet;
+                packetType = result.packetType;
+            }
         }
 
         // Handle monitors
         handlePacketSendMonitor(player, packetType, packet);
-        return true;
+        return result;
     }
 
-    public boolean handlePacketSendListener(Player player, PacketType packetType, Object packet, boolean wasCancelled) {
+    /**
+     * Handles when a packet is sent to a player
+     *
+     * @param player Player
+     * @param packetType Packet Type
+     * @param packet The packet
+     * @param wasCancelled Whether it was originally cancelled
+     * @return Result of handling this packet
+     */
+    public HandlerResult handlePacketSendListener(Player player, PacketType packetType, Object packet, boolean wasCancelled) {
         List<PacketListener> listenerList = listeners.get(packetType);
         if (listenerList != null) {
             CommonPacket cp = new CommonPacket(packet, packetType);
@@ -211,15 +244,28 @@ class PacketHandlerRegistration {
             for (PacketListener listener : listenerList) {
                 try {
                     listener.onPacketSend(ev);
+
+                    if (ev.getPacket() != cp) {
+                        cp = ev.getPacket();
+                        packet = cp.getHandle();
+
+                        // If packet type is changed, abort any further listener callbacks
+                        // Not sure how to proceed with this honestly...
+                        PacketType newType = ev.getType();
+                        if (newType != packetType) {
+                            packetType = newType;
+                            break;
+                        }
+                    }
                 } catch (Throwable t) {
                     Logging.LOGGER_NETWORK.log(Level.SEVERE, "Error occurred in onPacketSend handling " + packetType + ":", t);
                 }
             }
             if (ev.isCancelled()) {
-                return false;
+                return new HandlerResult(packet, packetType, true); // Cancelled
             }
         }
-        return true;
+        return new HandlerResult(packet, packetType, false);
     }
 
     public void handlePacketSendMonitor(Player player, PacketType packetType, Object packet) {
@@ -243,12 +289,13 @@ class PacketHandlerRegistration {
      * @param packet that is handled
      * @param wasCancelled - True if the packet is allowed to be received, False
      * if not
-     * @return True if the packet is allowed to be received, False if not
+     * @return Result of handling this packet
      */
-    public boolean handlePacketReceive(Player player, Object packet, boolean wasCancelled) {
+    public HandlerResult handlePacketReceive(Player player, Object packet, boolean wasCancelled) {
         if (player == null || !PacketHandle.T.isAssignableFrom(packet)) {
-            return true;
+            return new HandlerResult(packet, null, wasCancelled); // not handled
         }
+
         // Handle listeners
         PacketType type = PacketType.getType(packet);
         List<PacketListener> listenerList = listeners.get(type);
@@ -256,15 +303,30 @@ class PacketHandlerRegistration {
             CommonPacket cp = new CommonPacket(packet, type);
             PacketReceiveEvent ev = new PacketReceiveEvent(player, cp);
             ev.setCancelled(wasCancelled);
+
             for (PacketListener listener : listenerList) {
                 try {
                     listener.onPacketReceive(ev);
+
+                    if (ev.getPacket() != cp) {
+                        cp = ev.getPacket();
+                        packet = cp.getHandle();
+
+                        // If packet type is changed, abort any further listener callbacks
+                        // Not sure how to proceed with this honestly...
+                        PacketType newType = ev.getType();
+                        if (newType != type) {
+                            type = newType;
+                            break;
+                        }
+                    }
                 } catch (Throwable t) {
                     Logging.LOGGER_NETWORK.log(Level.SEVERE, "Error occurred in onPacketReceive handling " + type + ":", t);
                 }
             }
+
             if (ev.isCancelled()) {
-                return false;
+                return new HandlerResult(packet, type, true); // Cancelled
             }
         }
         // Handle monitors
@@ -279,6 +341,18 @@ class PacketHandlerRegistration {
                 }
             }
         }
-        return true;
+        return new HandlerResult(packet, type, false); // Not cancelled, but packet instance may have changed
+    }
+
+    public static class HandlerResult {
+        public final Object packet;
+        public final PacketType packetType;
+        public final boolean isCancelled;
+
+        public HandlerResult(Object packet, PacketType packetType, boolean isCancelled) {
+            this.packet = packet;
+            this.packetType = packetType;
+            this.isCancelled = isCancelled;
+        }
     }
 }

src/main/java/com/bergerkiller/bukkit/common/internal/network/ProtocolLibPacketHandler.java

@@ -176,6 +176,10 @@ public class ProtocolLibPacketHandler implements PacketHandler {
         return plugins;
     }
 
+    private PacketContainer toContainer(Object rawPacket) {
+        return new PacketContainer(getPacketType(rawPacket.getClass()), rawPacket);
+    }
+
     @Override
     public void receivePacket(Player player, PacketType type, Object packet) {
         if (PlayerConnectionHandle.forPlayer(player) == null) {
@@ -183,7 +187,7 @@ public class ProtocolLibPacketHandler implements PacketHandler {
         }
 
         type.preprocess(packet);
-        PacketContainer toReceive = new PacketContainer(getPacketType(packet.getClass()), packet);
+        PacketContainer toReceive = toContainer(packet);
         try {
             this.receivePacketMethod.invoke(getProtocolManager(), player, toReceive);
         } catch (RuntimeException ex) {
@@ -225,7 +229,7 @@ public class ProtocolLibPacketHandler implements PacketHandler {
         } else {
             // Silent - do not send it through listeners, only through monitors
             try {
-                final PacketContainer toSend = new PacketContainer(getPacketType(packet.getClass()), packet);
+                final PacketContainer toSend = toContainer(packet);
                 getProtocolManager().sendServerPacket(player, toSend, null, false);
             } catch (RuntimeException ex) {
                 if (this.loggedOutPlayerExceptionType.isInstance(ex)) {
@@ -517,6 +521,17 @@ public class ProtocolLibPacketHandler implements PacketHandler {
             this.listener = listener;
         }
 
+        private void swapPacket(PacketEvent event, CommonPacket commonpacket) {
+            PacketContainer container = toContainer(commonpacket.getHandle());
+
+            // TODO: Cancel the packet and schedule a new one instead to avoid warning
+            // Im not sure this actually works for receiving packets...
+            //if (event.getPacketType() != container.getType()) {
+            //    event.setCancelled(true);
+            //}
+            event.setPacket(container);
+        }
+
         @Override
         public void onPacketReceiving(PacketEvent event) {
             // Check not temporary
@@ -529,6 +544,9 @@ public class ProtocolLibPacketHandler implements PacketHandler {
             receiveEvent.setCancelled(event.isCancelled());
             listener.onPacketReceive(receiveEvent);
             event.setCancelled(receiveEvent.isCancelled());
+            if (receiveEvent.getPacket() != packet) {
+                swapPacket(event, receiveEvent.getPacket());
+            }
         }
 
         @Override
@@ -555,6 +573,9 @@ public class ProtocolLibPacketHandler implements PacketHandler {
                 listener.onPacketSend(sendEvent);
                 return sendEvent.isCancelled();
             }));
+            if (sendEvent.getPacket() != packet) {
+                swapPacket(event, sendEvent.getPacket());
+            }
         }
     }