From 87ad8948c142ec51f6d1ef5a610a91579b52fd4f Mon Sep 17 00:00:00 2001
From: Moulberry <james.jenour@protonmail.com>
Date: Sat, 24 Feb 2024 14:57:50 +0800
Subject: [PATCH] Check maxChunkLoadDistance for block entities

---
 .../RequestChunkDataPacketListener.java       | 46 ++++++++++++-------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/main/java/com/moulberry/axiom/packet/RequestChunkDataPacketListener.java b/src/main/java/com/moulberry/axiom/packet/RequestChunkDataPacketListener.java
index de6a6b3..db44993 100644
--- a/src/main/java/com/moulberry/axiom/packet/RequestChunkDataPacketListener.java
+++ b/src/main/java/com/moulberry/axiom/packet/RequestChunkDataPacketListener.java
@@ -20,6 +20,7 @@ import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.entity.BlockEntity;
 import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.chunk.ChunkStatus;
 import net.minecraft.world.level.chunk.LevelChunk;
 import net.minecraft.world.level.chunk.LevelChunkSection;
 import net.minecraft.world.level.chunk.PalettedContainer;
@@ -76,22 +77,6 @@ public class RequestChunkDataPacketListener implements PluginMessageListener {
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         BlockPos.MutableBlockPos mutableBlockPos = new BlockPos.MutableBlockPos();
 
-        // Save and compress block entities
-        int count = friendlyByteBuf.readVarInt();
-        for (int i = 0; i < count; i++) {
-            long pos = friendlyByteBuf.readLong();
-            BlockEntity blockEntity = level.getBlockEntity(mutableBlockPos.set(pos));
-            if (blockEntity != null) {
-                CompoundTag tag = blockEntity.saveWithoutMetadata();
-                blockEntityMap.put(pos, CompressedBlockEntity.compress(tag, baos));
-            }
-        }
-
-        int playerSectionX = player.getBlockX() >> 4;
-        int playerSectionZ = player.getBlockZ() >> 4;
-
-        Long2ObjectMap<PalettedContainer<BlockState>> sections = new Long2ObjectOpenHashMap<>();
-
         int maxChunkLoadDistance = this.plugin.configuration.getInt("max-chunk-load-distance");
 
         // Don't allow loading chunks outside render distance for plot worlds
@@ -99,6 +84,35 @@ public class RequestChunkDataPacketListener implements PluginMessageListener {
             maxChunkLoadDistance = 0;
         }
 
+        int playerSectionX = player.getBlockX() >> 4;
+        int playerSectionZ = player.getBlockZ() >> 4;
+
+        // Save and compress block entities
+        int count = friendlyByteBuf.readVarInt();
+        for (int i = 0; i < count; i++) {
+            long pos = friendlyByteBuf.readLong();
+            mutableBlockPos.set(pos);
+
+            if (level.isOutsideBuildHeight(mutableBlockPos)) continue;
+
+            int chunkX = mutableBlockPos.getX() >> 4;
+            int chunkZ = mutableBlockPos.getZ() >> 4;
+
+            int distance = Math.abs(playerSectionX - chunkX) + Math.abs(playerSectionZ - chunkZ);
+            boolean canLoad = distance <= maxChunkLoadDistance;
+
+            LevelChunk chunk = (LevelChunk) level.getChunk(chunkX, chunkZ, ChunkStatus.FULL, canLoad);
+            if (chunk == null) continue;
+
+            BlockEntity blockEntity = chunk.getBlockEntity(mutableBlockPos, LevelChunk.EntityCreationType.IMMEDIATE);
+            if (blockEntity != null) {
+                CompoundTag tag = blockEntity.saveWithoutMetadata();
+                blockEntityMap.put(pos, CompressedBlockEntity.compress(tag, baos));
+            }
+        }
+
+        Long2ObjectMap<PalettedContainer<BlockState>> sections = new Long2ObjectOpenHashMap<>();
+
         if (maxChunkLoadDistance > 0) {
             count = friendlyByteBuf.readVarInt();
             for (int i = 0; i < count; i++) {