Fixing 1.15 .schematic loading

SpigotCore_14/src/de/steamwar/sql/Schematic_14.java

@@ -1,19 +1,35 @@
 package de.steamwar.sql;
 
+import com.google.common.collect.ImmutableList;
+import com.sk89q.jnbt.*;
 import com.sk89q.worldedit.EmptyClipboardException;
+import com.sk89q.worldedit.WorldEditException;
 import com.sk89q.worldedit.bukkit.WorldEditPlugin;
 import com.sk89q.worldedit.extension.platform.Actor;
+import com.sk89q.worldedit.extent.clipboard.BlockArrayClipboard;
 import com.sk89q.worldedit.extent.clipboard.Clipboard;
 import com.sk89q.worldedit.extent.clipboard.io.BuiltInClipboardFormat;
 import com.sk89q.worldedit.extent.clipboard.io.ClipboardFormat;
 import com.sk89q.worldedit.extent.clipboard.io.ClipboardWriter;
+import com.sk89q.worldedit.extent.clipboard.io.NBTSchematicReader;
+import com.sk89q.worldedit.extent.clipboard.io.legacycompat.*;
+import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.regions.CuboidRegion;
+import com.sk89q.worldedit.regions.Region;
 import com.sk89q.worldedit.session.ClipboardHolder;
+import com.sk89q.worldedit.world.DataFixer;
+import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.registry.LegacyMapper;
 import org.bukkit.Bukkit;
 import org.bukkit.entity.Player;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.*;
+import java.util.zip.GZIPInputStream;
+
+import static com.google.common.base.Preconditions.checkNotNull;
 
 class Schematic_14 {
     private Schematic_14(){}
@@ -63,7 +79,7 @@ class Schematic_14 {
             if(schemFormat){
                 return SCHEM.getReader(is).read();
             }else{
-                return SCHEMATIC.getReader(is).read();
+                return new MCEditSchematicReader(new NBTInputStream(new GZIPInputStream(is))).read();
             }
         } catch (NullPointerException e) {
             throw new NoClipboardException();
@@ -73,4 +89,265 @@ class Schematic_14 {
     private static WorldEditPlugin getWorldEditPlugin() {
         return (WorldEditPlugin) Bukkit.getPluginManager().getPlugin("WorldEdit");
     }
+
+    private static class MCEditSchematicReader extends NBTSchematicReader {
+
+        private final NBTInputStream inputStream;
+        private final DataFixer fixer;
+        private boolean faweSchem = false;
+        private static final ImmutableList<NBTCompatibilityHandler> COMPATIBILITY_HANDLERS
+                = ImmutableList.of(
+                new SignCompatibilityHandler(),
+                new FlowerPotCompatibilityHandler(),
+                new NoteBlockCompatibilityHandler(),
+                new SkullBlockCompatibilityHandler()
+        );
+
+        /**
+         * Create a new instance.
+         *
+         * @param inputStream the input stream to read from
+         */
+        MCEditSchematicReader(NBTInputStream inputStream) {
+            checkNotNull(inputStream);
+            this.inputStream = inputStream;
+            this.fixer = null;
+        }
+
+        @Override
+        public Clipboard read() throws IOException {
+            // Schematic tag
+            NamedTag rootTag = inputStream.readNamedTag();
+            if (!rootTag.getName().equals("Schematic")) {
+                throw new IOException("Tag 'Schematic' does not exist or is not first");
+            }
+            CompoundTag schematicTag = (CompoundTag) rootTag.getTag();
+
+            // Check
+            Map<String, Tag> schematic = schematicTag.getValue();
+            if (!schematic.containsKey("Blocks")) {
+                throw new IOException("Schematic file is missing a 'Blocks' tag");
+            }
+
+            // Check type of Schematic
+            String materials = requireTag(schematic, "Materials", StringTag.class).getValue();
+            if (!materials.equals("Alpha")) {
+                throw new IOException("Schematic file is not an Alpha schematic");
+            }
+
+            // ====================================================================
+            // Metadata
+            // ====================================================================
+
+            BlockVector3 origin;
+            Region region;
+
+            // Get information
+            short width = requireTag(schematic, "Width", ShortTag.class).getValue();
+            short height = requireTag(schematic, "Height", ShortTag.class).getValue();
+            short length = requireTag(schematic, "Length", ShortTag.class).getValue();
+
+            int originX = 0;
+            int originY = 0;
+            int originZ = 0;
+            try {
+                originX = requireTag(schematic, "WEOriginX", IntTag.class).getValue();
+                originY = requireTag(schematic, "WEOriginY", IntTag.class).getValue();
+                originZ = requireTag(schematic, "WEOriginZ", IntTag.class).getValue();
+                BlockVector3 min = BlockVector3.at(originX, originY, originZ);
+
+                int offsetX = requireTag(schematic, "WEOffsetX", IntTag.class).getValue();
+                int offsetY = requireTag(schematic, "WEOffsetY", IntTag.class).getValue();
+                int offsetZ = requireTag(schematic, "WEOffsetZ", IntTag.class).getValue();
+                BlockVector3 offset = BlockVector3.at(offsetX, offsetY, offsetZ);
+
+                origin = min.subtract(offset);
+                region = new CuboidRegion(min, min.add(width, height, length).subtract(BlockVector3.ONE));
+            } catch (IOException ignored) {
+                origin = BlockVector3.ZERO;
+                region = new CuboidRegion(origin, origin.add(width, height, length).subtract(BlockVector3.ONE));
+            }
+
+            // ====================================================================
+            // Blocks
+            // ====================================================================
+
+            // Get blocks
+            byte[] blockId = requireTag(schematic, "Blocks", ByteArrayTag.class).getValue();
+            byte[] blockData = requireTag(schematic, "Data", ByteArrayTag.class).getValue();
+            byte[] addId = new byte[0];
+            short[] blocks = new short[blockId.length]; // Have to later combine IDs
+
+            // We support 4096 block IDs using the same method as vanilla Minecraft, where
+            // the highest 4 bits are stored in a separate byte array.
+            if (schematic.containsKey("AddBlocks")) {
+                addId = requireTag(schematic, "AddBlocks", ByteArrayTag.class).getValue();
+            }
+
+            // Combine the AddBlocks data with the first 8-bit block ID
+            for (int index = 0; index < blockId.length; index++) {
+                if ((index >> 1) >= addId.length) { // No corresponding AddBlocks index
+                    blocks[index] = (short) (blockId[index] & 0xFF);
+                } else {
+                    if ((index & 1) == 0) {
+                        blocks[index] = (short) (((addId[index >> 1] & 0x0F) << 8) + (blockId[index] & 0xFF));
+                    } else {
+                        blocks[index] = (short) (((addId[index >> 1] & 0xF0) << 4) + (blockId[index] & 0xFF));
+                    }
+                }
+            }
+
+            // Need to pull out tile entities
+            final ListTag tileEntityTag = getTag(schematic, "TileEntities", ListTag.class);
+            List<Tag> tileEntities = tileEntityTag == null ? new ArrayList<>() : tileEntityTag.getValue();
+            Map<BlockVector3, Map<String, Tag>> tileEntitiesMap = new HashMap<>();
+            Map<BlockVector3, BlockState> blockStates = new HashMap<>();
+
+            for (Tag tag : tileEntities) {
+                if (!(tag instanceof CompoundTag)) continue;
+                CompoundTag t = (CompoundTag) tag;
+                int x = t.getInt("x");
+                int y = t.getInt("y");
+                int z = t.getInt("z");
+                int index = y * width * length + z * width + x;
+                if(index < 0 || index >= blocks.length)
+                    faweSchem = true;
+            }
+
+            for (Tag tag : tileEntities) {
+                if (!(tag instanceof CompoundTag)) continue;
+                CompoundTag t = (CompoundTag) tag;
+                Map<String, Tag> values = new HashMap<>(t.getValue());
+                String id = t.getString("id");
+                values.put("id", new StringTag(convertBlockEntityId(id)));
+                int x = t.getInt("x");
+                int y = t.getInt("y");
+                int z = t.getInt("z");
+                if(faweSchem){
+                    x -= originX;
+                    y -= originY;
+                    z -= originZ;
+                }
+
+                int index = y * width * length + z * width + x;
+
+                BlockState block = getBlockState(blocks[index], blockData[index]);
+                System.out.println(block.getBlockType().getId());
+                BlockState newBlock = block;
+                if (newBlock != null) {
+                    for (NBTCompatibilityHandler handler : COMPATIBILITY_HANDLERS) {
+                        if (handler.isAffectedBlock(newBlock)) {
+                            newBlock = handler.updateNBT(block, values);
+                            if (newBlock == null || values.isEmpty()) {
+                                break;
+                            }
+                        }
+                    }
+                }
+                if (values.isEmpty()) {
+                    t = null;
+                } else {
+                    t = new CompoundTag(values);
+                }
+
+                if (fixer != null && t != null) {
+                    t = fixer.fixUp(DataFixer.FixTypes.BLOCK_ENTITY, t, -1);
+                }
+
+                BlockVector3 vec = BlockVector3.at(x, y, z);
+                if (t != null) {
+                    tileEntitiesMap.put(vec, t.getValue());
+                }
+                blockStates.put(vec, newBlock);
+            }
+
+            BlockArrayClipboard clipboard = new BlockArrayClipboard(region);
+            clipboard.setOrigin(origin);
+
+
+            for (int x = 0; x < width; ++x) {
+                for (int y = 0; y < height; ++y) {
+                    for (int z = 0; z < length; ++z) {
+                        int index = y * width * length + z * width + x;
+                        BlockVector3 pt = BlockVector3.at(x, y, z);
+                        BlockState state = blockStates.computeIfAbsent(pt, p -> getBlockState(blocks[index], blockData[index]));
+
+                        try {
+                            if (state != null) {
+                                if (tileEntitiesMap.containsKey(pt)) {
+                                    clipboard.setBlock(region.getMinimumPoint().add(pt), state.toBaseBlock(new CompoundTag(tileEntitiesMap.get(pt))));
+                                } else {
+                                    clipboard.setBlock(region.getMinimumPoint().add(pt), state);
+                                }
+                            }
+                        } catch (WorldEditException ignored) { // BlockArrayClipboard won't throw this
+                        }
+                    }
+                }
+            }
+
+            return clipboard;
+        }
+
+        private String convertBlockEntityId(String id) {
+            switch (id) {
+                case "Cauldron":
+                    return "brewing_stand";
+                case "Control":
+                    return "command_block";
+                case "DLDetector":
+                    return "daylight_detector";
+                case "Trap":
+                    return "dispenser";
+                case "EnchantTable":
+                    return "enchanting_table";
+                case "EndGateway":
+                    return "end_gateway";
+                case "AirPortal":
+                    return "end_portal";
+                case "EnderChest":
+                    return "ender_chest";
+                case "FlowerPot":
+                    return "flower_pot";
+                case "RecordPlayer":
+                    return "jukebox";
+                case "MobSpawner":
+                    return "mob_spawner";
+                case "Music":
+                case "noteblock":
+                    return "note_block";
+                case "Structure":
+                    return "structure_block";
+                case "Chest":
+                    return "chest";
+                case "Sign":
+                    return "sign";
+                case "Banner":
+                    return "banner";
+                case "Beacon":
+                    return "beacon";
+                case "Comparator":
+                    return "comparator";
+                case "Dropper":
+                    return "dropper";
+                case "Furnace":
+                    return "furnace";
+                case "Hopper":
+                    return "hopper";
+                case "Skull":
+                    return "skull";
+                default:
+                    return id;
+            }
+        }
+
+        private BlockState getBlockState(int id, int data) {
+            return LegacyMapper.getInstance().getBlockFromLegacy(id, data);
+        }
+
+        @Override
+        public void close() throws IOException {
+            inputStream.close();
+        }
+    }
 }