Implement ProtocolV5 (Block Entity Support)

src/main/java/com/moulberry/axiom/AxiomConstants.java

@@ -12,7 +12,7 @@ public class AxiomConstants {
         }
     }
 
-    public static final int API_VERSION = 4;
+    public static final int API_VERSION = 5;
     public static final NamespacedKey ACTIVE_HOTBAR_INDEX = new NamespacedKey("axiom", "active_hotbar_index");
     public static final NamespacedKey HOTBAR_DATA = new NamespacedKey("axiom", "hotbar_data");
 

src/main/java/com/moulberry/axiom/AxiomPaper.java

@@ -18,11 +18,11 @@ import org.bukkit.entity.Player;
 import org.bukkit.event.EventHandler;
 import org.bukkit.event.Listener;
 import org.bukkit.plugin.java.JavaPlugin;
+import org.bukkit.plugin.messaging.Messenger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 
-import java.util.HashSet;
-import java.util.Map;
-import java.util.UUID;
+import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 
 public class AxiomPaper extends JavaPlugin implements Listener {
 
@@ -30,20 +30,23 @@ public class AxiomPaper extends JavaPlugin implements Listener {
     public void onEnable() {
         Bukkit.getPluginManager().registerEvents(this, this);
 
-        Bukkit.getMessenger().registerOutgoingPluginChannel(this, "axiom:enable");
-        Bukkit.getMessenger().registerOutgoingPluginChannel(this, "axiom:initialize_hotbars");
-        Bukkit.getMessenger().registerOutgoingPluginChannel(this, "axiom:set_editor_views");
+        Messenger msg = Bukkit.getMessenger();
 
-        HashSet<UUID> activeAxiomPlayers = new HashSet<>();
+        msg.registerOutgoingPluginChannel(this, "axiom:enable");
+        msg.registerOutgoingPluginChannel(this, "axiom:initialize_hotbars");
+        msg.registerOutgoingPluginChannel(this, "axiom:set_editor_views");
 
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:hello", new HelloPacketListener(this, activeAxiomPlayers));
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:set_gamemode", new SetGamemodePacketListener());
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:set_fly_speed", new SetFlySpeedPacketListener());
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:set_block", new SetBlockPacketListener(this));
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:set_hotbar_slot", new SetHotbarSlotPacketListener());
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:switch_active_hotbar", new SwitchActiveHotbarPacketListener());
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:teleport", new TeleportPacketListener());
-        Bukkit.getMessenger().registerIncomingPluginChannel(this, "axiom:set_editor_views", new SetEditorViewsPacketListener());
+        final Set<UUID> activeAxiomPlayers = Collections.newSetFromMap(new ConcurrentHashMap<>());
+
+        msg.registerIncomingPluginChannel(this, "axiom:hello", new HelloPacketListener(this, activeAxiomPlayers));
+        msg.registerIncomingPluginChannel(this, "axiom:set_gamemode", new SetGamemodePacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:set_fly_speed", new SetFlySpeedPacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:set_block", new SetBlockPacketListener(this));
+        msg.registerIncomingPluginChannel(this, "axiom:set_hotbar_slot", new SetHotbarSlotPacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:switch_active_hotbar", new SwitchActiveHotbarPacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:teleport", new TeleportPacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:set_editor_views", new SetEditorViewsPacketListener());
+        msg.registerIncomingPluginChannel(this, "axiom:request_block_entity", new RequestBlockEntityPacketListener(this));
 
         SetBlockBufferPacketListener setBlockBufferPacketListener = new SetBlockBufferPacketListener(this);
 

src/main/java/com/moulberry/axiom/buffer/BlockBuffer.java

@@ -6,12 +6,14 @@ import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectSet;
 import it.unimi.dsi.fastutil.shorts.Short2ObjectMap;
+import it.unimi.dsi.fastutil.shorts.Short2ObjectOpenHashMap;
 import net.minecraft.core.BlockPos;
 import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.chunk.PalettedContainer;
+import org.jetbrains.annotations.Nullable;
 
 public class BlockBuffer {
 
@@ -21,6 +23,7 @@ public class BlockBuffer {
 
     private PalettedContainer<BlockState> last = null;
     private long lastId = AxiomConstants.MIN_POSITION_LONG;
+    private final Long2ObjectMap<Short2ObjectMap<CompressedBlockEntity>> blockEntities = new Long2ObjectOpenHashMap<>();
 
     public BlockBuffer() {
         this.values = new Long2ObjectOpenHashMap<>();
@@ -34,6 +37,17 @@ public class BlockBuffer {
         for (Long2ObjectMap.Entry<PalettedContainer<BlockState>> entry : this.entrySet()) {
             friendlyByteBuf.writeLong(entry.getLongKey());
             entry.getValue().write(friendlyByteBuf);
+
+            Short2ObjectMap<CompressedBlockEntity> blockEntities = this.blockEntities.get(entry.getLongKey());
+            if (blockEntities != null) {
+                friendlyByteBuf.writeVarInt(blockEntities.size());
+                for (Short2ObjectMap.Entry<CompressedBlockEntity> entry2 : blockEntities.short2ObjectEntrySet()) {
+                    friendlyByteBuf.writeShort(entry2.getShortKey());
+                    entry2.getValue().write(friendlyByteBuf);
+                }
+            } else {
+                friendlyByteBuf.writeVarInt(0);
+            }
         }
 
         friendlyByteBuf.writeLong(AxiomConstants.MIN_POSITION_LONG);
@@ -48,6 +62,17 @@ public class BlockBuffer {
 
             PalettedContainer<BlockState> palettedContainer = buffer.getOrCreateSection(index);
             palettedContainer.read(friendlyByteBuf);
+
+            int blockEntitySize = Math.min(4096, friendlyByteBuf.readVarInt());
+            if (blockEntitySize > 0) {
+                Short2ObjectMap<CompressedBlockEntity> map = new Short2ObjectOpenHashMap<>(blockEntitySize);
+                for (int i = 0; i < blockEntitySize; i++) {
+                    short offset = friendlyByteBuf.readShort();
+                    CompressedBlockEntity blockEntity = CompressedBlockEntity.read(friendlyByteBuf);
+                    map.put(offset, blockEntity);
+                }
+                buffer.blockEntities.put(index, map);
+            }
         }
 
         return buffer;
@@ -59,6 +84,30 @@ public class BlockBuffer {
         this.values.clear();
     }
 
+    public void putBlockEntity(int x, int y, int z, CompressedBlockEntity blockEntity) {
+        long cpos = BlockPos.asLong(x >> 4, y >> 4, z >> 4);
+        Short2ObjectMap<CompressedBlockEntity> chunkMap = this.blockEntities.computeIfAbsent(cpos, k -> new Short2ObjectOpenHashMap<>());
+
+        int key = (x & 0xF) | ((y & 0xF) << 4) | ((z & 0xF) << 8);
+        chunkMap.put((short)key, blockEntity);
+    }
+
+    @Nullable
+    public CompressedBlockEntity getBlockEntity(int x, int y, int z) {
+        long cpos = BlockPos.asLong(x >> 4, y >> 4, z >> 4);
+        Short2ObjectMap<CompressedBlockEntity> chunkMap = this.blockEntities.get(cpos);
+
+        if (chunkMap == null) return null;
+
+        int key = (x & 0xF) | ((y & 0xF) << 4) | ((z & 0xF) << 8);
+        return chunkMap.get((short)key);
+    }
+
+    @Nullable
+    public Short2ObjectMap<CompressedBlockEntity> getBlockEntityChunkMap(long cpos) {
+        return this.blockEntities.get(cpos);
+    }
+
     public BlockState get(int x, int y, int z) {
         var container = this.getSectionForCoord(x, y, z);
         if (container == null) {

src/main/java/com/moulberry/axiom/buffer/CompressedBlockEntity.java

@@ -0,0 +1,66 @@
+package com.moulberry.axiom.buffer;
+
+import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdDictCompress;
+import com.github.luben.zstd.ZstdDictDecompress;
+import com.moulberry.axiom.AxiomPaper;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.nbt.NbtIo;
+import net.minecraft.network.FriendlyByteBuf;
+
+import java.io.*;
+import java.util.Objects;
+
+public record CompressedBlockEntity(int originalSize, byte compressionDict, byte[] compressed) {
+
+    private static ZstdDictCompress zstdDictCompress = null;
+    private static ZstdDictDecompress zstdDictDecompress = null;
+
+    public static void initialize(AxiomPaper plugin) {
+        try (InputStream is = Objects.requireNonNull(plugin.getResource("zstd_dictionaries/block_entities_v1.dict"))) {
+            byte[] bytes = is.readAllBytes();
+            zstdDictCompress = new ZstdDictCompress(bytes, Zstd.defaultCompressionLevel());
+            zstdDictDecompress = new ZstdDictDecompress(bytes);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static CompressedBlockEntity compress(CompoundTag tag, ByteArrayOutputStream baos) {
+        try {
+            baos.reset();
+            DataOutputStream dos = new DataOutputStream(baos);
+            NbtIo.write(tag, dos);
+            byte[] uncompressed = baos.toByteArray();
+            byte[] compressed = Zstd.compress(uncompressed, zstdDictCompress);
+            return new CompressedBlockEntity(uncompressed.length, (byte) 0, compressed);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public CompoundTag decompress() {
+        if (this.compressionDict != 0) throw new UnsupportedOperationException("Unknown compression dict: " + this.compressionDict);
+
+        try {
+            byte[] nbt = Zstd.decompress(this.compressed, zstdDictDecompress, this.originalSize);
+            return NbtIo.read(new DataInputStream(new ByteArrayInputStream(nbt)));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static CompressedBlockEntity read(FriendlyByteBuf friendlyByteBuf) {
+        int originalSize = friendlyByteBuf.readVarInt();
+        byte compressionDict = friendlyByteBuf.readByte();
+        byte[] compressed = friendlyByteBuf.readByteArray();
+        return new CompressedBlockEntity(originalSize, compressionDict, compressed);
+    }
+
+    public void write(FriendlyByteBuf friendlyByteBuf) {
+        friendlyByteBuf.writeVarInt(this.originalSize);
+        friendlyByteBuf.writeByte(this.compressionDict);
+        friendlyByteBuf.writeByteArray(this.compressed);
+    }
+
+}

src/main/java/com/moulberry/axiom/packet/RequestBlockEntityPacketListener.java

@@ -0,0 +1,90 @@
+package com.moulberry.axiom.packet;
+
+import com.moulberry.axiom.AxiomPaper;
+import com.moulberry.axiom.buffer.CompressedBlockEntity;
+import io.netty.buffer.Unpooled;
+import it.unimi.dsi.fastutil.longs.*;
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.registries.Registries;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.network.FriendlyByteBuf;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.ServerPlayer;
+import net.minecraft.world.level.Level;
+import net.minecraft.world.level.block.entity.BlockEntity;
+import org.bukkit.craftbukkit.v1_20_R1.entity.CraftPlayer;
+import org.bukkit.entity.Player;
+import org.bukkit.plugin.messaging.PluginMessageListener;
+import org.jetbrains.annotations.NotNull;
+
+import java.io.ByteArrayOutputStream;
+
+public class RequestBlockEntityPacketListener implements PluginMessageListener {
+
+    private final AxiomPaper plugin;
+
+    public RequestBlockEntityPacketListener(AxiomPaper plugin) {
+        this.plugin = plugin;
+    }
+
+    @Override
+    public void onPluginMessageReceived(@NotNull String channel, @NotNull Player bukkitPlayer, @NotNull byte[] message) {
+        FriendlyByteBuf friendlyByteBuf = new FriendlyByteBuf(Unpooled.wrappedBuffer(message));
+        long id = friendlyByteBuf.readLong();
+
+        if (!bukkitPlayer.hasPermission("axiom.*")) {
+            // We always send an 'empty' response in order to make the client happy
+            sendEmptyResponse(bukkitPlayer, id);
+            return;
+        }
+
+        ServerPlayer player = ((CraftPlayer)bukkitPlayer).getHandle();
+        MinecraftServer server = player.getServer();
+        if (server == null) {
+            sendEmptyResponse(bukkitPlayer, id);
+            return;
+        }
+
+        ResourceKey<Level> worldKey = friendlyByteBuf.readResourceKey(Registries.DIMENSION);
+        ServerLevel level = server.getLevel(worldKey);
+        if (level == null) {
+            sendEmptyResponse(bukkitPlayer, id);
+            return;
+        }
+
+        Long2ObjectMap<CompressedBlockEntity> map = new Long2ObjectOpenHashMap<>();
+        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();
+                map.put(pos, CompressedBlockEntity.compress(tag, baos));
+            }
+        }
+
+        // Send response packet
+        FriendlyByteBuf buf = new FriendlyByteBuf(Unpooled.buffer(16));
+        buf.writeLong(id);
+        buf.writeVarInt(map.size());
+        for (Long2ObjectMap.Entry<CompressedBlockEntity> entry : map.long2ObjectEntrySet()) {
+            buf.writeLong(entry.getLongKey());
+            entry.getValue().write(buf);
+        }
+        bukkitPlayer.sendPluginMessage(this.plugin, "axiom:block_entities", buf.accessByteBufWithCorrectSize());
+    }
+
+    private void sendEmptyResponse(Player player, long id) {
+        FriendlyByteBuf buf = new FriendlyByteBuf(Unpooled.buffer(16));
+        buf.writeLong(id);
+        buf.writeByte(0); // no block entities
+        player.sendPluginMessage(this.plugin, "axiom:block_entities", buf.accessByteBufWithCorrectSize());
+    }
+
+}

src/main/java/com/moulberry/axiom/packet/SetBlockBufferPacketListener.java

@@ -3,8 +3,10 @@ package com.moulberry.axiom.packet;
 import com.moulberry.axiom.AxiomPaper;
 import com.moulberry.axiom.buffer.BiomeBuffer;
 import com.moulberry.axiom.buffer.BlockBuffer;
+import com.moulberry.axiom.buffer.CompressedBlockEntity;
 import io.netty.buffer.Unpooled;
 import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import it.unimi.dsi.fastutil.shorts.Short2ObjectMap;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
@@ -125,6 +127,8 @@ public class SetBlockBufferPacketListener {
                     }
                 }
 
+                Short2ObjectMap<CompressedBlockEntity> blockEntityChunkMap = buffer.getBlockEntityChunkMap(entry.getLongKey());
+
                 sectionStates.acquire();
                 try {
                     for (int x = 0; x < 16; x++) {
@@ -159,26 +163,41 @@ public class SetBlockBufferPacketListener {
                                         }
                                     }
 
-                                    boolean oldHasBlockEntity = old.hasBlockEntity();
-                                    if (old.is(block)) {
-                                        if (blockState.hasBlockEntity()) {
-                                            BlockEntity blockEntity = chunk.getBlockEntity(blockPos, LevelChunk.EntityCreationType.CHECK);
-                                            if (blockEntity == null) {
-                                                blockEntity = ((EntityBlock)block).newBlockEntity(blockPos, blockState);
-                                                if (blockEntity != null) {
-                                                    chunk.addAndRegisterBlockEntity(blockEntity);
-                                                }
-                                            } else {
-                                                blockEntity.setBlockState(blockState);
+                                    if (blockState.hasBlockEntity()) {
+                                        BlockEntity blockEntity = chunk.getBlockEntity(blockPos, LevelChunk.EntityCreationType.CHECK);
 
-                                                try {
-                                                    this.updateBlockEntityTicker.invoke(chunk, blockEntity);
-                                                } catch (IllegalAccessException | InvocationTargetException e) {
-                                                    throw new RuntimeException(e);
-                                                }
+                                        if (blockEntity == null) {
+                                            // There isn't a block entity here, create it!
+                                            blockEntity = ((EntityBlock)block).newBlockEntity(blockPos, blockState);
+                                            if (blockEntity != null) {
+                                                chunk.addAndRegisterBlockEntity(blockEntity);
+                                            }
+                                        } else if (blockEntity.getType().isValid(blockState)) {
+                                            // Block entity is here and the type is correct
+                                            blockEntity.setBlockState(blockState);
+
+                                            try {
+                                                this.updateBlockEntityTicker.invoke(chunk, blockEntity);
+                                            } catch (IllegalAccessException | InvocationTargetException e) {
+                                                throw new RuntimeException(e);
+                                            }
+                                        } else {
+                                            // Block entity type isn't correct, we need to recreate it
+                                            chunk.removeBlockEntity(blockPos);
+
+                                            blockEntity = ((EntityBlock)block).newBlockEntity(blockPos, blockState);
+                                            if (blockEntity != null) {
+                                                chunk.addAndRegisterBlockEntity(blockEntity);
                                             }
                                         }
-                                    } else if (oldHasBlockEntity) {
+                                        if (blockEntity != null && blockEntityChunkMap != null) {
+                                            int key = x | (y << 4) | (z << 8);
+                                            CompressedBlockEntity savedBlockEntity = blockEntityChunkMap.get((short) key);
+                                            if (savedBlockEntity != null) {
+                                                blockEntity.load(savedBlockEntity.decompress());
+                                            }
+                                        }
+                                    } else if (old.hasBlockEntity()) {
                                         chunk.removeBlockEntity(blockPos);
                                     }