From 41b0d95d70449a9ed5b07d1dca08f220c0b402e8 Mon Sep 17 00:00:00 2001
From: EnZaXD <florian.michael07@gmail.com>
Date: Sat, 27 Jul 2024 20:26:53 +0200
Subject: [PATCH] Fix horse/llama's armor/carpet not visible in 1.20.3->1.20.5
 (#4052)

---
 .../BlockItemPacketRewriter1_20_5.java        |  4 +-
 .../rewriter/EntityPacketRewriter1_20_5.java  | 39 +++++++++++++++++--
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/BlockItemPacketRewriter1_20_5.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/BlockItemPacketRewriter1_20_5.java
index 6c6f6867e..cef85579f 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/BlockItemPacketRewriter1_20_5.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/BlockItemPacketRewriter1_20_5.java
@@ -174,10 +174,10 @@ public final class BlockItemPacketRewriter1_20_5 extends ItemRewriter<Clientboun
         registerCooldown(ClientboundPackets1_20_3.COOLDOWN);
         registerSetContent1_17_1(ClientboundPackets1_20_3.CONTAINER_SET_CONTENT);
         registerSetSlot1_17_1(ClientboundPackets1_20_3.CONTAINER_SET_SLOT);
-        registerSetEquipment(ClientboundPackets1_20_3.SET_EQUIPMENT);
         registerContainerClick1_17_1(ServerboundPackets1_20_5.CONTAINER_CLICK);
         registerContainerSetData(ClientboundPackets1_20_3.CONTAINER_SET_DATA);
         registerSetCreativeModeSlot(ServerboundPackets1_20_5.SET_CREATIVE_MODE_SLOT);
+
         protocol.registerServerbound(ServerboundPackets1_20_5.CONTAINER_BUTTON_CLICK, wrapper -> {
             final byte containerId = wrapper.read(Types.VAR_INT).byteValue();
             final byte buttonId = wrapper.read(Types.VAR_INT).byteValue();
@@ -1597,4 +1597,4 @@ public final class BlockItemPacketRewriter1_20_5 extends ItemRewriter<Clientboun
 
         return true;
     }
-}
\ No newline at end of file
+}
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/EntityPacketRewriter1_20_5.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/EntityPacketRewriter1_20_5.java
index 31b05684c..063e79aa3 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/EntityPacketRewriter1_20_5.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_3to1_20_5/rewriter/EntityPacketRewriter1_20_5.java
@@ -26,10 +26,12 @@ import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.data.entity.DimensionData;
 import com.viaversion.viaversion.api.minecraft.Particle;
 import com.viaversion.viaversion.api.minecraft.RegistryEntry;
+import com.viaversion.viaversion.api.minecraft.data.StructuredDataContainer;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_20_5;
 import com.viaversion.viaversion.api.minecraft.entitydata.EntityData;
 import com.viaversion.viaversion.api.minecraft.item.Item;
+import com.viaversion.viaversion.api.minecraft.item.StructuredItem;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
 import com.viaversion.viaversion.api.protocol.remapper.PacketHandlers;
 import com.viaversion.viaversion.api.type.Types;
@@ -45,7 +47,6 @@ import com.viaversion.viaversion.protocols.v1_20_3to1_20_5.packet.ClientboundCon
 import com.viaversion.viaversion.protocols.v1_20_3to1_20_5.packet.ClientboundPackets1_20_5;
 import com.viaversion.viaversion.protocols.v1_20_3to1_20_5.storage.AcknowledgedMessagesStorage;
 import com.viaversion.viaversion.rewriter.EntityRewriter;
-import com.viaversion.viaversion.rewriter.entitydata.EntityDataHandler;
 import com.viaversion.viaversion.util.Key;
 import com.viaversion.viaversion.util.TagUtil;
 import it.unimi.dsi.fastutil.ints.IntArraySet;
@@ -73,6 +74,22 @@ public final class EntityPacketRewriter1_20_5 extends EntityRewriter<Clientbound
         registerSetEntityData(ClientboundPackets1_20_3.SET_ENTITY_DATA, Types1_20_3.ENTITY_DATA_LIST, Types1_20_5.ENTITY_DATA_LIST);
         registerRemoveEntities(ClientboundPackets1_20_3.REMOVE_ENTITIES);
 
+        protocol.registerClientbound(ClientboundPackets1_20_3.SET_EQUIPMENT, wrapper -> {
+            final int entityId = wrapper.passthrough(Types.VAR_INT); // Entity id
+            final EntityType type = tracker(wrapper.user()).entityType(entityId);
+
+            byte slot;
+            do {
+                slot = wrapper.read(Types.BYTE);
+                if (type != null && type.isOrHasParent(EntityTypes1_20_5.ABSTRACT_HORSE) && slot == 4) {
+                    slot = 6; // Map chest slot index to body slot index for horses
+                }
+                wrapper.write(Types.BYTE, slot);
+                Item item = protocol.getItemRewriter().handleItemToClient(wrapper.user(), wrapper.read(Types.ITEM1_20_2));
+                wrapper.write(Types1_20_5.ITEM, item);
+            } while (slot < 0);
+        });
+
         protocol.registerClientbound(ClientboundConfigurationPackets1_20_3.REGISTRY_DATA, wrapper -> {
             final PacketWrapper knownPacksPacket = wrapper.create(ClientboundConfigurationPackets1_20_5.SELECT_KNOWN_PACKS);
             knownPacksPacket.write(Types.VAR_INT, 0); // No known packs, everything is sent here
@@ -430,7 +447,23 @@ public final class EntityPacketRewriter1_20_5 extends EntityRewriter<Clientbound
             data.setTypeAndValue(Types1_20_5.ENTITY_DATA_TYPES.particlesType, new Particle[]{particle});
         });
 
-        filter().type(EntityTypes1_20_5.LLAMA).removeIndex(20); // Carpet color
+        filter().type(EntityTypes1_20_5.LLAMA).handler((event, data) -> {
+            // Carpet color removed - now set via the set equipment packet
+            final int dataIndex = event.index();
+            if (dataIndex == 20) {
+                event.cancel();
+                final int color = data.value();
+
+                // Convert dyed color id to carpet item id
+                final PacketWrapper setEquipment = PacketWrapper.create(ClientboundPackets1_20_5.SET_EQUIPMENT, event.user());
+                setEquipment.write(Types.VAR_INT, event.entityId());
+                setEquipment.write(Types.BYTE, (byte) 6);
+                setEquipment.write(Types1_20_5.ITEM, new StructuredItem(color + 446, 1, new StructuredDataContainer()));
+                setEquipment.scheduleSend(Protocol1_20_3To1_20_5.class);
+            } else if (dataIndex > 20) {
+                event.setIndex(dataIndex - 1);
+            }
+        });
         filter().type(EntityTypes1_20_5.AREA_EFFECT_CLOUD).handler((event, data) -> {
             // Color removed - Now put into the actual particle
             final int dataIndex = event.index();
@@ -510,4 +543,4 @@ public final class EntityPacketRewriter1_20_5 extends EntityRewriter<Clientbound
     public EntityType typeFromId(final int type) {
         return EntityTypes1_20_5.getTypeFromId(type);
     }
-}
\ No newline at end of file
+}