Merge pull request #58 from KJP12/fix/heads-creative

Fixes for textured heads & creative inventories

Author: samolego

src/main/java/org/samo_lego/golfiv/casts/ItemStackChecker.java

@@ -5,9 +5,7 @@
 import net.minecraft.entity.effect.StatusEffectInstance;
 import net.minecraft.entity.effect.StatusEffects;
 import net.minecraft.item.*;
-import net.minecraft.nbt.NbtCompound;
-import net.minecraft.nbt.NbtElement;
-import net.minecraft.nbt.NbtList;
+import net.minecraft.nbt.*;
 import net.minecraft.potion.PotionUtil;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.registry.Registry;
@@ -34,11 +32,11 @@ public interface ItemStackChecker {
     /**
      * Creates a fake ItemStack based on the
      * provided stack.
-     *
+     * <p>
      * E.g. removes enchantment info, stack size, etc.
      * Used on ground ItemEntities / opponent's stacks etc.
      *
-     * @param original the original ItemStack
+     * @param original   the original ItemStack
      * @param spoofCount whether or not the item count should be faked
      * @return the faked ItemStack
      */
@@ -47,52 +45,78 @@ static ItemStack fakeStack(ItemStack original, boolean spoofCount) {
                 new ItemStack(original.getItem(), original.getMaxCount()) :
                 new ItemStack(original.getItem(), original.getCount());
 
-        if(original.hasEnchantments())
+        if (original.hasEnchantments())
             fakedStack.addEnchantment(null, 0);
 
         Item item = original.getItem();
-        if(item instanceof DyeableItem dyeable) {
-            if(dyeable.hasColor(original)) {
+        if (item instanceof DyeableItem dyeable) {
+            if (dyeable.hasColor(original)) {
                 dyeable.setColor(fakedStack, dyeable.getColor(original));
             }
         }
 
-        if(item instanceof PotionItem || item instanceof TippedArrowItem) {
+        if (item instanceof PotionItem || item instanceof TippedArrowItem) {
             // Lets dropping potions and arrows to be less of a 'surprising' change.
-            fakedStack.getOrCreateNbt().putInt("CustomPotionColor", PotionUtil.getColor(original));
+            fakedStack.setSubNbt(PotionUtil.CUSTOM_POTION_COLOR_KEY, NbtInt.of(PotionUtil.getColor(original)));
             if (item.hasGlint(original)) {
                 PotionUtil.setCustomPotionEffects(fakedStack, UNLUCK);
             }
         }
 
-        if(item instanceof WritableBookItem || item instanceof WrittenBookItem) {
+        if (item instanceof WritableBookItem || item instanceof WrittenBookItem) {
             // Prevents issues with other mods expecting pages to be present.
             fakedStack.setSubNbt("pages", new NbtList());
         }
 
-        if(item instanceof BannerItem) {
+        if (item instanceof BannerItem) {
             // Lets dropping banners not be a surprising change, and for illager patrol leaders to have a proper banner.
-            cleanBanner(original.getSubNbt("BlockEntityTag"), fakedStack);
+            cleanBanner(original.getSubNbt(BlockItem.BLOCK_STATE_TAG_KEY), fakedStack);
+        }
+
+        if (item instanceof SkullItem) {
+            // Lets dropping player heads not be a surprising change, and for wearing heads to show the actual skin.
+            cleanSkull(original.getSubNbt(SkullItem.SKULL_OWNER_KEY), fakedStack);
+        }
+
+        if (item instanceof CrossbowItem && CrossbowItem.isCharged(original)) {
+            cleanCrossbow(original.getNbt(), fakedStack, false);
         }
 
         return fakedStack;
     }
 
+    /**
+     * This is for keeping creative from getting inadvertently sanitised,
+     * as creative has the property of allowing one to summon any item at will,
+     * including freely mutating the inventory.
+     *
+     * @param stack The original ItemStack.
+     * @return The faked ItemStack, with a GolfIV pointer injected.
+     * @author KJP12
+     * @see #inventoryStack(ItemStack)
+     */
+    static ItemStack creativeInventoryStack(ItemStack stack) {
+        ItemStack fake = inventoryStack(stack);
+        if (stack.hasNbt()) //noinspection ConstantConditions
+            fake.setSubNbt("GolfIV", NbtInt.of(stack.getNbt().hashCode()));
+        return fake;
+    }
+
     /**
      * Creates a fake ItemStack based on the
      * provided stack.
-     *
+     * <p>
      * Changes the tag to be the minimal required
      * NBT to render in an inventory.
      *
      * @param stack The original ItemStack.
      * @return The faked ItemStack
-     * */
+     */
     static ItemStack inventoryStack(ItemStack stack) {
         // TODO: Perhaps take a more dynamic approach to this?
         //  This is not really flexible as is and may leave modded items broken.
         NbtCompound tag = stack.getNbt();
-        if(tag == null || tag.isEmpty()) {
+        if (tag == null || tag.isEmpty()) {
             // Sanitization isn't necessary when it's already empty.
             return stack;
         }
@@ -101,14 +125,15 @@ static ItemStack inventoryStack(ItemStack stack) {
         ItemStack fake = new ItemStack(item, stack.getCount());
 
         // Rewrite display.
-        if(tag.contains(ItemStack.DISPLAY_KEY)) {
+        if (tag.contains(ItemStack.DISPLAY_KEY)) {
             NbtCompound display = tag.getCompound(ItemStack.DISPLAY_KEY);
             NbtCompound fakeDisplay = new NbtCompound();
             NbtElement name = display.get(ItemStack.NAME_KEY);
-            if(name != null) fakeDisplay.put(ItemStack.NAME_KEY, name);
-            if(display.contains(ItemStack.COLOR_KEY)) fakeDisplay.put(ItemStack.COLOR_KEY, display.get(ItemStack.COLOR_KEY));
+            if (name != null) fakeDisplay.put(ItemStack.NAME_KEY, name);
+            if (display.contains(ItemStack.COLOR_KEY))
+                fakeDisplay.put(ItemStack.COLOR_KEY, display.get(ItemStack.COLOR_KEY));
             NbtElement lore = display.get(ItemStack.LORE_KEY);
-            if(lore != null) fakeDisplay.put(ItemStack.LORE_KEY, lore);
+            if (lore != null) fakeDisplay.put(ItemStack.LORE_KEY, lore);
             fake.setSubNbt(ItemStack.DISPLAY_KEY, fakeDisplay);
         }
 
@@ -140,31 +165,40 @@ static ItemStack inventoryStack(ItemStack stack) {
         // Check block items.
         if(item instanceof BlockItem) {
             boolean flag = true;
-            if(item instanceof BannerItem) {
+            if (item instanceof BannerItem) {
                 flag = false;
-                cleanBanner(stack.getSubNbt("BlockEntityTag"), fake);
+                cleanBanner(stack.getSubNbt(BlockItem.BLOCK_STATE_TAG_KEY), fake);
             }
 
-            if(flag) {
+            // Transfer SkullProperties.
+            if (item instanceof SkullItem) {
+                flag = false;
+                cleanSkull(stack.getSubNbt(SkullItem.SKULL_OWNER_KEY), fake);
+            }
+
+            if (flag) {
                 Block block = ((BlockItem) item).getBlock();
 
                 // Rewrite shulker items
                 if (block instanceof ShulkerBoxBlock) {
-                    NbtCompound blockEntity = stack.getSubNbt("BlockEntityTag");
-                    if(blockEntity != null) {
-                        NbtCompound fakeEntity = fake.getOrCreateSubNbt("BlockEntityTag");
-                        if(blockEntity.contains("LootTable", NbtElement.STRING_TYPE)) {
+                    NbtCompound blockEntity = stack.getSubNbt(BlockItem.BLOCK_STATE_TAG_KEY);
+                    if (blockEntity != null) {
+                        NbtCompound fakeEntity = fake.getOrCreateSubNbt(BlockItem.BLOCK_STATE_TAG_KEY);
+                        if (blockEntity.contains("LootTable", NbtElement.STRING_TYPE)) {
                             fakeEntity.put("LootTable", blockEntity.get("LootTable"));
                         }
-                        if(blockEntity.contains("Items", NbtElement.LIST_TYPE)) {
+                        if (blockEntity.contains("Items", NbtElement.LIST_TYPE)) {
                             NbtList fakeItems = new NbtList();
-                            for(NbtElement $item : blockEntity.getList("Items", NbtElement.COMPOUND_TYPE)) {
-                                if($item == null) continue;
+                            for (NbtElement $item : blockEntity.getList("Items", NbtElement.COMPOUND_TYPE)) {
+                                if ($item == null) continue;
                                 NbtCompound oldItem = (NbtCompound) $item;
                                 NbtCompound fakeItem = new NbtCompound();
-                                if(oldItem.contains("Slot", NbtElement.BYTE_TYPE)) fakeItem.put("Slot", oldItem.get("Slot"));
-                                if(oldItem.contains("id", NbtElement.STRING_TYPE)) fakeItem.put("id", oldItem.get("id"));
-                                if(oldItem.contains("Count", NbtElement.BYTE_TYPE)) fakeItem.put("Count", oldItem.get("Count"));
+                                if (oldItem.contains("Slot", NbtElement.BYTE_TYPE))
+                                    fakeItem.put("Slot", oldItem.get("Slot"));
+                                if (oldItem.contains("id", NbtElement.STRING_TYPE))
+                                    fakeItem.put("id", oldItem.get("id"));
+                                if (oldItem.contains("Count", NbtElement.BYTE_TYPE))
+                                    fakeItem.put("Count", oldItem.get("Count"));
                                 // TODO: Add in display name
                                 fakeItems.add(fakeItem);
                             }
@@ -204,27 +238,94 @@ static ItemStack inventoryStack(ItemStack stack) {
             fake.setSubNbt("pages", new NbtList());
         }
 
-        if(item instanceof WritableBookItem) {
+        if (item instanceof WritableBookItem) {
             // FIXME: Pages need to be present for the book to work. Force update on selection?
             // Prevents issues with other mods expecting pages to be present.
             fake.setSubNbt("pages", new NbtList());
         }
 
-        if(item instanceof FilledMapItem) {
+        if (item instanceof FilledMapItem) {
             fake.getOrCreateNbt().putInt("map", tag.getInt("map"));
         }
 
+        if (item instanceof CrossbowItem && CrossbowItem.isCharged(stack)) {
+            cleanCrossbow(stack.getNbt(), fake, true);
+        }
+
         return fake;
     }
 
     /**
-     * Minimally copies over the banner NBT based on the provided blockEntity data.
+     * Minimally copies over the crossbow data based on the provided Crossbow data.
+     * <p>
+     * Only the NBT required to render the crossbow, including rocket, is copied over.
+     *
+     * @param crossbow    The raw crossbow NBT.
+     * @param faked       The faked ItemStack to copy to.
+     * @param isInventory Whether to send the raw item or not.
+     */
+    static void cleanCrossbow(NbtCompound crossbow, ItemStack faked, boolean isInventory) {
+        if (crossbow == null) return;
+        NbtList originalProjectiles = crossbow.getList("ChargedProjectiles", NbtElement.COMPOUND_TYPE);
+        if (originalProjectiles.isEmpty()) return;
+        String originalProjectile = originalProjectiles.getCompound(0).getString("id");
+        String projectile;
+
+        if (isInventory || "minecraft:firework".equals(originalProjectile)) {
+            projectile = originalProjectile;
+        } else {
+            projectile = "minecraft:arrow";
+        }
+
+        NbtCompound projectileStack = new NbtCompound();
+        projectileStack.putByte("Count", (byte) 1);
+        projectileStack.putString("id", projectile);
+        NbtList projectiles = new NbtList();
+        projectiles.add(projectileStack);
+        faked.setSubNbt("ChargedProjectiles", projectiles);
+        faked.setSubNbt("Charged", NbtByte.of(true));
+    }
+
+    /**
+     * Minimally copies over the skull data based on the provided SkullOwner data.
+     * <p>
+     * Only the NBT required to render the skull is copied over.
      *
+     * @param skullOwner The SkullOwner NBT compound. May be null.
+     * @param faked      The faked ItemStack to copy to.
+     */
+    static void cleanSkull(NbtCompound skullOwner, ItemStack faked) {
+        if (skullOwner != null) {
+            NbtCompound fakeSkullOwner = faked.getOrCreateSubNbt(SkullItem.SKULL_OWNER_KEY);
+            if (skullOwner.containsUuid("Id")) fakeSkullOwner.putUuid("Id", skullOwner.getUuid("Id"));
+            if (skullOwner.contains("Properties", NbtElement.COMPOUND_TYPE)) {
+                NbtCompound skullProperties = skullOwner.getCompound("Properties");
+                if (skullProperties.contains("textures", NbtElement.LIST_TYPE)) {
+                    NbtList skullTextures = skullProperties.getList("textures", NbtElement.COMPOUND_TYPE);
+                    if (!skullTextures.isEmpty()) {
+                        NbtCompound skullValueContainer = skullTextures.getCompound(0);
+                        String skullTexture = skullValueContainer.getString("Value");
+                        NbtCompound fakeProperties = new NbtCompound();
+                        NbtList fakeTextures = new NbtList();
+                        NbtCompound fakeTexture = new NbtCompound();
+                        fakeTexture.putString("Value", skullTexture);
+                        fakeTextures.add(fakeTexture);
+                        fakeProperties.put("textures", fakeTextures);
+                        fakeSkullOwner.put("Properties", fakeProperties);
+                    }
+                }
+            }
+        }
+    }
+
+    /**
+     * Minimally copies over the banner NBT based on the provided blockEntity data.
+     * <p>
      * Only the NBT required to render a layer is copied over.
      *
      * @param blockEntity The original block entity tag. May be null.
-     * @param faked The faked ItemStack to copy to.
-     * */
+     * @param faked       The faked ItemStack to copy to.
+     */
     static void cleanBanner(NbtCompound blockEntity, ItemStack faked) {
         if (blockEntity != null && blockEntity.contains("Patterns", NbtElement.LIST_TYPE)) {
             NbtList fakePatterns = new NbtList();
@@ -247,7 +348,7 @@ static void cleanBanner(NbtCompound blockEntity, ItemStack faked) {
                     }
                 }
             }
-            faked.getOrCreateSubNbt("BlockEntityTag").put("Patterns", fakePatterns);
+            faked.getOrCreateSubNbt(BlockItem.BLOCK_STATE_TAG_KEY).put("Patterns", fakePatterns);
         }
     }
 }

src/main/java/org/samo_lego/golfiv/event/S2CPacket/ItemInventoryKickPatch.java

@@ -13,45 +13,75 @@
 import org.samo_lego.golfiv.mixin.accessors.ScreenHandlerSlotUpdateS2CPacketAccessor;
 
 import java.util.List;
-import java.util.stream.Collectors;
+import java.util.function.Consumer;
 
 import static org.samo_lego.golfiv.GolfIV.golfConfig;
-import static org.samo_lego.golfiv.casts.ItemStackChecker.fakeStack;
+import static org.samo_lego.golfiv.casts.ItemStackChecker.creativeInventoryStack;
+import static org.samo_lego.golfiv.casts.ItemStackChecker.inventoryStack;
 
 public class ItemInventoryKickPatch implements S2CPacketCallback {
 
     public ItemInventoryKickPatch() {
     }
 
     /**
-     * Changes ItemStacks in {@link InventoryS2CPacket}
-     * to not include tags, as they get sent additionally by {@link ScreenHandlerSlotUpdateS2CPacket}
+     * Removes non-critical tags from item stacks in the event of a
+     * creative player and packet overflows.
+     * <p>
+     * Creative players will additionally also get a {@code GolfIV} tag
+     * injected as a hash of the original NBT to prevent Creative inventory
+     * management from unexpectedly mangling the item either by just opening
+     * the inventory or by {@link org.samo_lego.golfiv.storage.GolfConfig.IllegalItems.Creative#removeCreativeNBTTags
+     * Remove Creative NBT Tags}.
      *
      * @param packet packet being sent
      * @param player player getting the packet
      * @param server Minecraft Server
+     * @see ItemStackChecker#fakeStack(ItemStack, boolean)
+     * @see ItemStackChecker#inventoryStack(ItemStack)
+     * @see org.samo_lego.golfiv.mixin.illegal_items.ServerPlayNetworkHandlerMixin_CreativeItemsCheck
      */
     @Override
     public void preSendPacket(Packet<?> packet, ServerPlayerEntity player, MinecraftServer server) {
-        if(golfConfig.packet.patchItemKickExploit && packet instanceof InventoryS2CPacket) {
-            PacketByteBuf testBuf = new PacketByteBuf(Unpooled.buffer());
-            packet.write(testBuf);
-            if(testBuf.readableBytes() > 2097140) {
-                List<ItemStack> contents = ((InventoryS2CPacketAccessor) packet).getContents();
-                List<ItemStack> fakedContents = contents.stream().map(ItemStackChecker::inventoryStack).collect(Collectors.toList());
+        if (!golfConfig.packet.patchItemKickExploit) return;
+        if (packet instanceof InventoryS2CPacket inventoryPacket) {
+            // The creative player does some weirdness in regards to sending packets.
+            // This will try to guarantee that the creative player will not accidentally
+            // wipe their inventory of various NBT required for the item.
+
+            // For the moment, this may produce some weirdness in regards to rendering as this doesn't quite
+            // send all possible NBT. However, the inventory will remain intact, which arguably is better
+            // than some missing data for the client.
+            if (player.isCreative() || isOversized(packet::write)) {
+                List<ItemStack> contents = inventoryPacket.getContents();
+                List<ItemStack> fakedContents = contents.stream().map(player.isCreative() ?
+                                ItemStackChecker::creativeInventoryStack :
+                                ItemStackChecker::inventoryStack)
+                        .toList();
 
                 ((InventoryS2CPacketAccessor) packet).setContents(fakedContents);
             }
-            testBuf.release();
-        } else if(golfConfig.packet.patchItemKickExploit && packet instanceof ScreenHandlerSlotUpdateS2CPacket) {
+        } else if (packet instanceof ScreenHandlerSlotUpdateS2CPacket) {
             ItemStack stack = ((ScreenHandlerSlotUpdateS2CPacketAccessor) packet).getStack();
-            if(stack.getNbt() != null) {
-                PacketByteBuf testBuf = new PacketByteBuf(Unpooled.buffer());
-                if(testBuf.writeItemStack(stack).readableBytes() > 2097140) {
-                    ((ScreenHandlerSlotUpdateS2CPacketAccessor) packet).setStack(fakeStack(stack, false));
-                }
-                testBuf.release();
+            if (stack.hasNbt() && (player.isCreative() || isOversized(buf -> buf.writeItemStack(stack)))) {
+                ((ScreenHandlerSlotUpdateS2CPacketAccessor) packet).setStack(player.isCreative() ? creativeInventoryStack(stack) : inventoryStack(stack));
             }
         }
     }
+
+    /**
+     * Tests if the packet will overflow the maximum allowed for the buffer.
+     *
+     * @param packet The packet in the form of a consumer. Allows for arbitrary packets.
+     * @return true if the packet is greater than 2,097,140, false otherwise.
+     */
+    private static boolean isOversized(Consumer<PacketByteBuf> packet) {
+        PacketByteBuf testBuf = new PacketByteBuf(Unpooled.buffer());
+        try {
+            packet.accept(testBuf);
+            return testBuf.readableBytes() > 2097140;
+        } finally {
+            testBuf.release();
+        }
+    }
 }