From 0ea01bfa4876fd9722698f096dbbfefbf5ccd7d1 Mon Sep 17 00:00:00 2001
From: Camotoy <20743703+Camotoy@users.noreply.github.com>
Date: Sat, 25 May 2024 18:59:37 -0400
Subject: [PATCH] Allow skull pick item NBT

---
 .../geyser/inventory/GeyserItemStack.java     |  2 +-
 .../geyser/level/block/type/SkullBlock.java   | 36 +++++++++++++++++++
 .../BedrockBlockPickRequestTranslator.java    |  7 +++-
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/inventory/GeyserItemStack.java b/core/src/main/java/org/geysermc/geyser/inventory/GeyserItemStack.java
index fa62769fe..744ad70b6 100644
--- a/core/src/main/java/org/geysermc/geyser/inventory/GeyserItemStack.java
+++ b/core/src/main/java/org/geysermc/geyser/inventory/GeyserItemStack.java
@@ -69,7 +69,7 @@ public class GeyserItemStack {
         return of(javaId, amount, null);
     }
 
-    public static @NonNull GeyserItemStack of(int javaId, int amount, DataComponents components) {
+    public static @NonNull GeyserItemStack of(int javaId, int amount, @Nullable DataComponents components) {
         return new GeyserItemStack(javaId, amount, components);
     }
 
diff --git a/core/src/main/java/org/geysermc/geyser/level/block/type/SkullBlock.java b/core/src/main/java/org/geysermc/geyser/level/block/type/SkullBlock.java
index 76b532919..c4aae46a2 100644
--- a/core/src/main/java/org/geysermc/geyser/level/block/type/SkullBlock.java
+++ b/core/src/main/java/org/geysermc/geyser/level/block/type/SkullBlock.java
@@ -25,10 +25,20 @@
 
 package org.geysermc.geyser.level.block.type;
 
+import com.github.steveice10.mc.auth.data.GameProfile;
 import org.cloudburstmc.math.vector.Vector3i;
+import org.cloudburstmc.nbt.NbtMap;
+import org.cloudburstmc.nbt.NbtMapBuilder;
 import org.cloudburstmc.protocol.bedrock.data.definitions.BlockDefinition;
+import org.geysermc.geyser.inventory.GeyserItemStack;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.SkullCache;
+import org.geysermc.mcprotocollib.protocol.data.game.item.ItemStack;
+import org.geysermc.mcprotocollib.protocol.data.game.item.component.DataComponentType;
+import org.geysermc.mcprotocollib.protocol.data.game.item.component.DataComponents;
+
+import java.util.Collections;
+import java.util.UUID;
 
 public class SkullBlock extends Block {
     private final Type type;
@@ -55,6 +65,32 @@ public class SkullBlock extends Block {
         // It's not an empty skull.
     }
 
+    public ItemStack pickItem(GeyserSession session, BlockState state, Vector3i position) {
+        SkullCache.Skull skull = session.getSkullCache().getSkulls().get(position);
+        if (skull == null) {
+            return new ItemStack(pickItem(state).getId());
+        }
+
+        GeyserItemStack itemStack = GeyserItemStack.of(pickItem(state).getId(), 1);
+        // This is a universal block entity behavior, but hardcode how it works for now.
+        NbtMapBuilder builder = NbtMap.builder()
+                .putString("id", "minecraft:skull")
+                .putInt("x", position.getX())
+                .putInt("y", position.getY())
+                .putInt("z", position.getZ());
+        DataComponents components = itemStack.getOrCreateComponents();
+        components.put(DataComponentType.BLOCK_ENTITY_DATA, builder.build());
+
+        UUID uuid = skull.getUuid();
+        String texturesProperty = skull.getTexturesProperty();
+        GameProfile profile = new GameProfile(uuid, null);
+        if (texturesProperty != null) {
+            profile.setProperties(Collections.singletonList(new GameProfile.Property("textures", texturesProperty)));
+        }
+        components.put(DataComponentType.PROFILE, profile);
+        return itemStack.getItemStack();
+    }
+
     public Type skullType() {
         return type;
     }
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockBlockPickRequestTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockBlockPickRequestTranslator.java
index 0a721e4b0..94368a6d4 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockBlockPickRequestTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockBlockPickRequestTranslator.java
@@ -33,6 +33,7 @@ import org.geysermc.geyser.item.Items;
 import org.geysermc.geyser.level.block.Blocks;
 import org.geysermc.geyser.level.block.type.BannerBlock;
 import org.geysermc.geyser.level.block.type.BlockState;
+import org.geysermc.geyser.level.block.type.SkullBlock;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
 import org.geysermc.geyser.translator.protocol.Translator;
@@ -65,7 +66,11 @@ public class BedrockBlockPickRequestTranslator extends PacketTranslator<BlockPic
         }
 
         boolean addExtraData = packet.isAddUserData() && blockToPick.block().hasBlockEntity(); // Holding down CTRL
-        if (blockToPick.block() instanceof BannerBlock || addExtraData) {
+        if (session.isInstabuild() && addExtraData && blockToPick.block() instanceof SkullBlock skull) {
+            InventoryUtils.findOrCreateItem(session, skull.pickItem(session, blockToPick, vector));
+            return;
+        }
+        if (blockToPick.block() instanceof BannerBlock) {
             session.getGeyser().getWorldManager().getPickItemComponents(session, vector.getX(), vector.getY(), vector.getZ(), addExtraData)
                     .whenComplete((components, ex) -> session.ensureInEventLoop(() -> {
                         if (components == null) {