diff --git a/api/src/main/java/org/geysermc/geyser/api/block/custom/nonvanilla/JavaBlockState.java b/api/src/main/java/org/geysermc/geyser/api/block/custom/nonvanilla/JavaBlockState.java
index 3096752bb..0dd0d3b33 100644
--- a/api/src/main/java/org/geysermc/geyser/api/block/custom/nonvanilla/JavaBlockState.java
+++ b/api/src/main/java/org/geysermc/geyser/api/block/custom/nonvanilla/JavaBlockState.java
@@ -111,7 +111,7 @@ public interface JavaBlockState {
          * @deprecated Does not have an effect. If you were using this to
          *      * set piston behavior, use {@link #pistonBehavior(String)} instead.
          */
-        @Deprecated
+        @Deprecated(forRemoval = true)
         Builder hasBlockEntity(boolean hasBlockEntity);
 
         JavaBlockState build();
diff --git a/core/src/main/java/org/geysermc/geyser/level/block/GeyserJavaBlockState.java b/core/src/main/java/org/geysermc/geyser/level/block/GeyserJavaBlockState.java
index 8028a4355..782f664fc 100644
--- a/core/src/main/java/org/geysermc/geyser/level/block/GeyserJavaBlockState.java
+++ b/core/src/main/java/org/geysermc/geyser/level/block/GeyserJavaBlockState.java
@@ -16,7 +16,6 @@ public class GeyserJavaBlockState implements JavaBlockState {
     boolean canBreakWithHand;
     String pickItem;
     String pistonBehavior;
-    boolean hasBlockEntity;
 
     private GeyserJavaBlockState(Builder builder) {
         this.identifier = builder.identifier;
@@ -28,7 +27,6 @@ public class GeyserJavaBlockState implements JavaBlockState {
         this.canBreakWithHand = builder.canBreakWithHand;
         this.pickItem = builder.pickItem;
         this.pistonBehavior = builder.pistonBehavior;
-        this.hasBlockEntity = builder.hasBlockEntity;
     }
 
     @Override
@@ -76,9 +74,10 @@ public class GeyserJavaBlockState implements JavaBlockState {
         return pistonBehavior;
     }
 
+    @SuppressWarnings("removal")
     @Override
     public boolean hasBlockEntity() {
-        return hasBlockEntity;
+        return false;
     }
 
     public static class Builder implements JavaBlockState.Builder {
@@ -91,7 +90,6 @@ public class GeyserJavaBlockState implements JavaBlockState {
         private boolean canBreakWithHand;
         private String pickItem;
         private String pistonBehavior;
-        private boolean hasBlockEntity;
 
         @Override
         public Builder identifier(@NonNull String identifier) {
@@ -147,9 +145,13 @@ public class GeyserJavaBlockState implements JavaBlockState {
             return this;
         }
 
+        @SuppressWarnings("removal")
         @Override
         public Builder hasBlockEntity(boolean hasBlockEntity) {
-            this.hasBlockEntity = hasBlockEntity;
+            // keep the current behavior
+            if (this.pistonBehavior == null && hasBlockEntity) {
+                this.pistonBehavior = "BLOCK";
+            }
             return this;
         }
 
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
index 4210a8fe0..24d419c24 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
@@ -33,11 +33,11 @@ import org.cloudburstmc.protocol.bedrock.data.structure.StructureMirror;
 import org.cloudburstmc.protocol.bedrock.data.structure.StructureRotation;
 import org.cloudburstmc.protocol.bedrock.packet.ContainerOpenPacket;
 import org.cloudburstmc.protocol.bedrock.packet.UpdateBlockPacket;
+import org.geysermc.geyser.level.block.Blocks;
 import org.geysermc.geyser.level.block.type.BlockState;
 import org.geysermc.geyser.registry.BlockRegistries;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.level.block.entity.BlockEntityTranslator;
-import org.geysermc.geyser.translator.level.block.entity.RequiresBlockState;
 import org.geysermc.geyser.translator.level.block.entity.SkullBlockEntityTranslator;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
 import org.geysermc.geyser.translator.protocol.Translator;
@@ -59,11 +59,11 @@ public class JavaBlockEntityDataTranslator extends PacketTranslator<ClientboundB
         BlockEntityTranslator translator = BlockEntityUtils.getBlockEntityTranslator(type);
         // The Java block state is used in BlockEntityTranslator.translateTag() to make up for some inconsistencies
         // between Java block states and Bedrock block entity data
-        BlockState blockState;
-        if (translator instanceof RequiresBlockState) {
-            blockState = BlockRegistries.BLOCK_STATES.get(session.getGeyser().getWorldManager().getBlockAt(session, packet.getPosition()));
-        } else {
-            blockState = BlockRegistries.BLOCK_STATES.get(0); //TODO
+        BlockState blockState = BlockRegistries.BLOCK_STATES.getOrDefault(session.getGeyser().getWorldManager().getBlockAt(session, packet.getPosition()),
+                Blocks.AIR.defaultBlockState());
+
+        if (blockState.block().blockEntityType() != type) {
+            return;
         }
 
         Vector3i position = packet.getPosition();
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
index f11d27b8c..0f7d5ced4 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
@@ -39,6 +39,7 @@ import org.cloudburstmc.protocol.bedrock.data.definitions.BlockDefinition;
 import org.cloudburstmc.protocol.bedrock.packet.LevelChunkPacket;
 import org.geysermc.geyser.entity.type.ItemFrameEntity;
 import org.geysermc.geyser.level.BedrockDimension;
+import org.geysermc.geyser.level.block.Blocks;
 import org.geysermc.geyser.level.block.type.Block;
 import org.geysermc.geyser.level.block.type.BlockState;
 import org.geysermc.geyser.level.chunk.BlockStorage;
@@ -398,7 +399,7 @@ public class JavaLevelChunkWithLightTranslator extends PacketTranslator<Clientbo
 
                 // Get the Java block state ID from block entity position
                 DataPalette section = javaChunks[(y >> 4) - yOffset];
-                BlockState blockState = BlockRegistries.BLOCK_STATES.get(section.get(x, y & 0xF, z));
+                BlockState blockState = BlockRegistries.BLOCK_STATES.getOrDefault(section.get(x, y & 0xF, z), Blocks.AIR.defaultBlockState());
 
                 // Note that, since 1.20.5, tags can be null, but Bedrock still needs a default tag to render the item
                 // Also, some properties - like banner base colors - are part of the tag and is processed here.
@@ -406,7 +407,7 @@ public class JavaLevelChunkWithLightTranslator extends PacketTranslator<Clientbo
 
                 // The Java server can send block entity data for blocks that aren't actually those blocks.
                 // A Java client ignores these
-                if (blockState != null && blockState.block().hasBlockEntity() && type.equals(blockState.block().blockEntityType())) {
+                if (blockState.block().hasBlockEntity() && type == blockState.block().blockEntityType()) {
                     bedrockBlockEntities.add(blockEntityTranslator.getBlockEntityTag(session, type, x + chunkBlockX, y, z + chunkBlockZ, tag, blockState));
 
                     // Check for custom skulls