Fix several issues with ender tank covers (#3377)

Co-authored-by: Martin Robertz <dream-master@gmx.net>
Co-authored-by: Ethryan <3237986+Ethryan@users.noreply.github.com>

5 files changed, 278 insertions, 76 deletions

diff --git a/src/main/java/tectech/mechanics/enderStorage/EnderFluidContainer.java b/src/main/java/tectech/mechanics/enderStorage/EnderFluidContainer.java
index f77a5a1b51..9995ec5da0 100644
--- a/src/main/java/tectech/mechanics/enderStorage/EnderFluidContainer.java
+++ b/src/main/java/tectech/mechanics/enderStorage/EnderFluidContainer.java
@@ -131,4 +131,26 @@ public class EnderFluidContainer implements IFluidHandler, Serializable {
         }
         if (fluidStack != null) fluidStack.tag = tag;
     }
+
+    public NBTTagCompound save() {
+        NBTTagCompound data = new NBTTagCompound();
+
+        if (fluidStack != null) {
+            fluidStack.writeToNBT(data);
+        }
+
+        return data;
+    }
+
+    public static EnderFluidContainer load(NBTTagCompound data) {
+        EnderFluidContainer container = new EnderFluidContainer();
+
+        if (!"".equals(data.getString("FluidName"))) {
+            container.fluidStack = FluidStack.loadFluidStackFromNBT(data);
+        } else {
+            container.fluidStack = null;
+        }
+
+        return container;
+    }
 }
diff --git a/src/main/java/tectech/mechanics/enderStorage/EnderLinkTag.java b/src/main/java/tectech/mechanics/enderStorage/EnderLinkTag.java
index 516b7fd637..78b3c47c10 100644
--- a/src/main/java/tectech/mechanics/enderStorage/EnderLinkTag.java
+++ b/src/main/java/tectech/mechanics/enderStorage/EnderLinkTag.java
@@ -3,6 +3,8 @@ package tectech.mechanics.enderStorage;
 import java.io.Serializable;
 import java.util.UUID;
 
+import net.minecraft.nbt.NBTTagCompound;
+
 import com.google.common.base.Objects;
 
 public class EnderLinkTag implements Serializable {
@@ -36,4 +38,23 @@ public class EnderLinkTag implements Serializable {
     public int hashCode() {
         return Objects.hashCode(frequency, player);
     }
+
+    public NBTTagCompound save() {
+        NBTTagCompound data = new NBTTagCompound();
+
+        data.setString("Freq", frequency);
+        if (player != null) {
+            data.setLong("ID1", player.getLeastSignificantBits());
+            data.setLong("ID2", player.getMostSignificantBits());
+        }
+
+        return data;
+    }
+
+    public static EnderLinkTag load(NBTTagCompound data) {
+        String freq = data.getString("Freq");
+        UUID player = data.hasKey("ID1") ? new UUID(data.getLong("ID2"), data.getLong("ID1")) : null;
+
+        return new EnderLinkTag(freq, player);
+    }
 }
diff --git a/src/main/java/tectech/mechanics/enderStorage/EnderLinkTank.java b/src/main/java/tectech/mechanics/enderStorage/EnderLinkTank.java
index 1b14c1371a..fe363b2058 100644
--- a/src/main/java/tectech/mechanics/enderStorage/EnderLinkTank.java
+++ b/src/main/java/tectech/mechanics/enderStorage/EnderLinkTank.java
@@ -2,11 +2,14 @@ package tectech.mechanics.enderStorage;
 
 import java.io.Serializable;
 
+import net.minecraft.nbt.NBTTagCompound;
 import net.minecraft.tileentity.TileEntity;
+import net.minecraft.world.World;
 import net.minecraftforge.common.DimensionManager;
 import net.minecraftforge.fluids.IFluidHandler;
 
 import com.google.common.base.Objects;
+import com.gtnewhorizon.gtnhlib.util.CoordinatePacker;
 
 public class EnderLinkTank implements Serializable {
 
@@ -24,12 +27,24 @@ public class EnderLinkTank implements Serializable {
         D = tile.getWorldObj().provider.dimensionId;
     }
 
+    private EnderLinkTank(int x, int y, int z, int d) {
+        X = x;
+        Y = y;
+        Z = z;
+        D = d;
+    }
+
     public IFluidHandler getFluidHandler() {
-        IFluidHandler fluidHandler = null;
-        TileEntity tile = DimensionManager.getWorld(D)
-            .getTileEntity(X, Y, Z);
-        if (tile instanceof IFluidHandler) fluidHandler = (IFluidHandler) tile;
-        return fluidHandler;
+        World world = DimensionManager.getWorld(D);
+
+        if (world == null) return null;
+
+        TileEntity tile = world.getTileEntity(X, Y, Z);
+        if (tile instanceof IFluidHandler fluidHandler) {
+            return fluidHandler;
+        } else {
+            return null;
+        }
     }
 
     @Override
@@ -44,4 +59,24 @@ public class EnderLinkTank implements Serializable {
     public int hashCode() {
         return Objects.hashCode(X, Y, Z, D);
     }
+
+    public NBTTagCompound save() {
+        NBTTagCompound data = new NBTTagCompound();
+
+        data.setLong("coords", CoordinatePacker.pack(X, Y, Z));
+        data.setInteger("dim", D);
+
+        return data;
+    }
+
+    public static EnderLinkTank load(NBTTagCompound data) {
+        long coords = data.getLong("coords");
+        int dim = data.getInteger("dim");
+
+        return new EnderLinkTank(
+            CoordinatePacker.unpackX(coords),
+            CoordinatePacker.unpackY(coords),
+            CoordinatePacker.unpackZ(coords),
+            dim);
+    }
 }
diff --git a/src/main/java/tectech/mechanics/enderStorage/EnderWorldSavedData.java b/src/main/java/tectech/mechanics/enderStorage/EnderWorldSavedData.java
index 515812548e..62625ed435 100644
--- a/src/main/java/tectech/mechanics/enderStorage/EnderWorldSavedData.java
+++ b/src/main/java/tectech/mechanics/enderStorage/EnderWorldSavedData.java
@@ -1,15 +1,18 @@
 package tectech.mechanics.enderStorage;
 
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
+import java.io.ObjectStreamClass;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import net.minecraft.nbt.NBTTagCompound;
+import net.minecraft.nbt.NBTTagList;
 import net.minecraft.world.WorldSavedData;
 import net.minecraft.world.storage.MapStorage;
 import net.minecraftforge.common.DimensionManager;
@@ -26,7 +29,7 @@ public class EnderWorldSavedData extends WorldSavedData {
     private static final String DATA_NAME = Reference.MODID + "_EnderWorldSavedData";
     private static final String ENDER_LIQUID_TAG_LINK = DATA_NAME + "_EnderLiquidTagLink";
     private static final String ENDER_LIQUID_TANK_LINK = DATA_NAME + "_EnderLiquidTankLink";
-    private static final EnderLinkTag DEFAULT_LINK_TAG = new EnderLinkTag("", null);
+    private static final String ENDER_LIQUID_TAG_VERSION = DATA_NAME + "_Version";
 
     private Map<EnderLinkTag, EnderFluidContainer> EnderLiquidTagLink = new HashMap<>();
     private Map<EnderLinkTank, EnderLinkTag> EnderLiquidTankLink = new HashMap<>();
@@ -42,54 +45,91 @@ public class EnderWorldSavedData extends WorldSavedData {
     @SuppressWarnings("unchecked")
     @Override
     public void readFromNBT(NBTTagCompound nbtTagCompound) {
-        try {
-            byte[] ba = nbtTagCompound.getByteArray(ENDER_LIQUID_TAG_LINK);
-            InputStream is = new ByteArrayInputStream(ba);
-            ObjectInputStream ois = new ObjectInputStream(is);
-            Object data = ois.readObject();
-            EnderLiquidTagLink = (Map<EnderLinkTag, EnderFluidContainer>) data;
-        } catch (IOException | ClassNotFoundException e) {
-            System.out.println("ENDER_LIQUID_TAG_LINK LOAD FAILED");
-            e.printStackTrace();
-        }
+        int version = nbtTagCompound.getInteger(ENDER_LIQUID_TAG_VERSION);
+
+        switch (version) {
+            case 0: {
+                try {
+                    byte[] bytes = nbtTagCompound.getByteArray(ENDER_LIQUID_TAG_LINK);
+
+                    try (InputStream is = new ByteArrayInputStream(bytes);
+                        ObjectInputStream objectStream = new MigratingObjectInputStream(is);) {
+                        EnderLiquidTagLink = (Map<EnderLinkTag, EnderFluidContainer>) objectStream.readObject();
+                    }
+                } catch (IOException | ClassNotFoundException e) {
+                    System.out.println("ENDER_LIQUID_TAG_LINK LOAD FAILED");
+                    e.printStackTrace();
+                }
+
+                try {
+                    byte[] bytes = nbtTagCompound.getByteArray(ENDER_LIQUID_TANK_LINK);
+                    try (InputStream is = new ByteArrayInputStream(bytes);
+                        ObjectInputStream objectStream = new MigratingObjectInputStream(is);) {
+                        EnderLiquidTankLink = (Map<EnderLinkTank, EnderLinkTag>) objectStream.readObject();
+                    }
+                } catch (IOException | ClassNotFoundException e) {
+                    System.out.println("ENDER_LIQUID_TANK_LINK LOAD FAILED");
+                    e.printStackTrace();
+                }
+                break;
+            }
+            case 1: {
+                List<NBTTagCompound> tags = ((NBTTagList) nbtTagCompound.getTag(ENDER_LIQUID_TAG_LINK)).tagList;
+
+                EnderLiquidTagLink = new HashMap<>();
+
+                for (NBTTagCompound tagLink : tags) {
+                    EnderLinkTag tag = EnderLinkTag.load((NBTTagCompound) tagLink.getTag("k"));
+                    EnderFluidContainer container = EnderFluidContainer.load((NBTTagCompound) tagLink.getTag("v"));
+                    EnderLiquidTagLink.put(tag, container);
+                }
+
+                List<NBTTagCompound> tanks = ((NBTTagList) nbtTagCompound.getTag(ENDER_LIQUID_TANK_LINK)).tagList;
 
-        try {
-            byte[] ba = nbtTagCompound.getByteArray(ENDER_LIQUID_TANK_LINK);
-            InputStream is = new ByteArrayInputStream(ba);
-            ObjectInputStream ois = new ObjectInputStream(is);
-            Object data = ois.readObject();
-            EnderLiquidTankLink = (Map<EnderLinkTank, EnderLinkTag>) data;
-        } catch (IOException | ClassNotFoundException e) {
-            System.out.println("ENDER_LIQUID_TANK_LINK LOAD FAILED");
-            e.printStackTrace();
+                EnderLiquidTankLink = new HashMap<>();
+
+                for (NBTTagCompound tankLink : tanks) {
+                    EnderLinkTank tank = EnderLinkTank.load((NBTTagCompound) tankLink.getTag("k"));
+                    EnderLinkTag tag = EnderLinkTag.load((NBTTagCompound) tankLink.getTag("v"));
+                    EnderLiquidTankLink.put(tank, tag);
+                }
+
+                break;
+            }
         }
     }
 
     @Override
     public void writeToNBT(NBTTagCompound nbtTagCompound) {
-        try {
-            ByteArrayOutputStream bos = new ByteArrayOutputStream();
-            ObjectOutputStream oos = new ObjectOutputStream(bos);
-            oos.writeObject(EnderLiquidTagLink);
-            oos.flush();
-            byte[] data = bos.toByteArray();
-            nbtTagCompound.setByteArray(ENDER_LIQUID_TAG_LINK, data);
-        } catch (IOException e) {
-            System.out.println("ENDER_LIQUID_TAG_LINK SAVE FAILED");
-            e.printStackTrace();
-        }
+        nbtTagCompound.setInteger(ENDER_LIQUID_TAG_VERSION, 1);
 
-        try {
-            ByteArrayOutputStream bos = new ByteArrayOutputStream();
-            ObjectOutputStream oos = new ObjectOutputStream(bos);
-            oos.writeObject(EnderLiquidTankLink);
-            oos.flush();
-            byte[] data = bos.toByteArray();
-            nbtTagCompound.setByteArray(ENDER_LIQUID_TANK_LINK, data);
-        } catch (IOException e) {
-            System.out.println("ENDER_LIQUID_TANK_LINK SAVE FAILED");
-            e.printStackTrace();
-        }
+        HashSet<EnderLinkTag> usedTags = new HashSet<>();
+
+        NBTTagList tanks = new NBTTagList();
+
+        EnderLiquidTankLink.forEach((tank, tag) -> {
+            if (tank.getFluidHandler() != null) {
+                usedTags.add(tag);
+                NBTTagCompound entry = new NBTTagCompound();
+                entry.setTag("k", tank.save());
+                entry.setTag("v", tag.save());
+                tanks.appendTag(entry);
+            }
+        });
+
+        nbtTagCompound.setTag(ENDER_LIQUID_TANK_LINK, tanks);
+        NBTTagList tags = new NBTTagList();
+
+        EnderLiquidTagLink.forEach((tag, container) -> {
+            if (usedTags.contains(tag)) {
+                NBTTagCompound entry = new NBTTagCompound();
+                entry.setTag("k", tag.save());
+                entry.setTag("v", container.save());
+                tags.appendTag(entry);
+            }
+        });
+
+        nbtTagCompound.setTag(ENDER_LIQUID_TAG_LINK, tags);
     }
 
     private static void loadInstance() {
@@ -123,16 +163,38 @@ public class EnderWorldSavedData extends WorldSavedData {
 
     public static EnderLinkTag getEnderLinkTag(IFluidHandler handler) {
         EnderLinkTank tank = new EnderLinkTank(handler);
-        if (!getEnderLiquidTankLink().containsKey(tank)) {
-            getEnderLiquidTankLink().put(tank, DEFAULT_LINK_TAG);
+        EnderLinkTag tag = getEnderLiquidTankLink().get(tank);
+
+        if (tag != null && Objects.equals(tag.getFrequency(), "")) {
+            tag = null;
         }
-        return getEnderLiquidTankLink().get(tank);
+
+        return tag;
     }
 
     public static void bindEnderLinkTag(IFluidHandler handler, EnderLinkTag tag) {
         EnderLinkTank tank = new EnderLinkTank(handler);
-        getEnderLiquidTankLink().remove(tank);
-        getEnderLiquidTankLink().put(tank, tag);
+
+        if (!tag.equals(getEnderLiquidTankLink().get(tank))) {
+            unbindTank(handler);
+
+            getEnderLiquidTankLink().remove(tank);
+            getEnderLiquidTankLink().put(tank, tag);
+        }
+    }
+
+    public static void unbindTank(IFluidHandler handler) {
+        EnderLinkTank tank = new EnderLinkTank(handler);
+        EnderLinkTag oldTag = getEnderLiquidTankLink().remove(tank);
+
+        if (oldTag != null) {
+            boolean isReferenced = getEnderLiquidTankLink().values()
+                .contains(oldTag);
+
+            if (!isReferenced) {
+                getEnderLiquidLink().remove(oldTag);
+            }
+        }
     }
 
     @SubscribeEvent
@@ -141,4 +203,25 @@ public class EnderWorldSavedData extends WorldSavedData {
             INSTANCE = null;
         }
     }
+
+    private static class MigratingObjectInputStream extends ObjectInputStream {
+
+        public MigratingObjectInputStream(InputStream in) throws IOException {
+            super(in);
+        }
+
+        @Override
+        protected ObjectStreamClass readClassDescriptor() throws IOException, ClassNotFoundException {
+            ObjectStreamClass resultClassDescriptor = super.readClassDescriptor();
+
+            String className = resultClassDescriptor.getName();
+
+            if (className.startsWith("com.github.technus.")) {
+                resultClassDescriptor = ObjectStreamClass
+                    .lookup(Class.forName(className.replace("com.github.technus.", "")));
+            }
+
+            return resultClassDescriptor;
+        }
+    }
 }
diff --git a/src/main/java/tectech/thing/cover/CoverEnderFluidLink.java b/src/main/java/tectech/thing/cover/CoverEnderFluidLink.java
index bae4e26b6c..9e63c85221 100644
--- a/src/main/java/tectech/thing/cover/CoverEnderFluidLink.java
+++ b/src/main/java/tectech/thing/cover/CoverEnderFluidLink.java
@@ -1,5 +1,6 @@
 package tectech.thing.cover;
 
+import java.util.Objects;
 import java.util.UUID;
 
 import net.minecraft.entity.player.EntityPlayer;
@@ -45,31 +46,50 @@ public class CoverEnderFluidLink extends CoverBehavior {
         }
     }
 
-    private boolean testBit(int aCoverVariable, int bitMask) {
+    private static boolean testBit(int aCoverVariable, int bitMask) {
         return (aCoverVariable & bitMask) != 0;
     }
 
-    private int toggleBit(int aCoverVariable, int bitMask) {
+    private static int toggleBit(int aCoverVariable, int bitMask) {
         return (aCoverVariable ^ bitMask);
     }
 
     @Override
     public int doCoverThings(ForgeDirection side, byte aInputRedstone, int aCoverID, int aCoverVariable,
         ICoverable aTileEntity, long aTimer) {
-        if ((aTileEntity instanceof IFluidHandler fluidHandlerSelf)) {
-            IFluidHandler fluidHandlerEnder = EnderWorldSavedData
-                .getEnderFluidContainer(EnderWorldSavedData.getEnderLinkTag((IFluidHandler) aTileEntity));
-
-            if (testBit(aCoverVariable, IMPORT_EXPORT_MASK)) {
-                transferFluid(fluidHandlerEnder, ForgeDirection.UNKNOWN, fluidHandlerSelf, side, L_PER_TICK);
-            } else {
-                transferFluid(fluidHandlerSelf, side, fluidHandlerEnder, ForgeDirection.UNKNOWN, L_PER_TICK);
+        if ((aTileEntity instanceof IFluidHandler teTank)) {
+            EnderLinkTag tag = EnderWorldSavedData.getEnderLinkTag((IFluidHandler) aTileEntity);
+
+            if (tag != null) {
+                IFluidHandler enderTank = EnderWorldSavedData.getEnderFluidContainer(tag);
+
+                if (testBit(aCoverVariable, IMPORT_EXPORT_MASK)) {
+                    transferFluid(enderTank, ForgeDirection.UNKNOWN, teTank, side, L_PER_TICK);
+                } else {
+                    transferFluid(teTank, side, enderTank, ForgeDirection.UNKNOWN, L_PER_TICK);
+                }
             }
         }
         return aCoverVariable;
     }
 
     @Override
+    public void onBaseTEDestroyed(ForgeDirection side, int aCoverID, int aCoverVariable, ICoverable aTileEntity) {
+        if (aTileEntity instanceof IFluidHandler fluidHandlerSelf) {
+            EnderWorldSavedData.unbindTank(fluidHandlerSelf);
+        }
+    }
+
+    @Override
+    public boolean onCoverRemoval(ForgeDirection side, int aCoverID, int aCoverVariable, ICoverable aTileEntity,
+        boolean aForced) {
+        if (aTileEntity instanceof IFluidHandler fluidHandlerSelf) {
+            EnderWorldSavedData.unbindTank(fluidHandlerSelf);
+        }
+        return true;
+    }
+
+    @Override
     public String getDescription(ForgeDirection side, int aCoverID, int aCoverVariable, ICoverable aTileEntity) {
         return "";
     }
@@ -143,23 +163,23 @@ public class CoverEnderFluidLink extends CoverBehavior {
             builder.widget(frequencyField.setGetter(() -> {
                 ICoverable te = getUIBuildContext().getTile();
                 if (!frequencyField.isClient() && te instanceof IFluidHandler) {
-                    return EnderWorldSavedData.getEnderLinkTag((IFluidHandler) te)
-                        .getFrequency();
+                    EnderLinkTag tag = EnderWorldSavedData.getEnderLinkTag((IFluidHandler) te);
+
+                    return tag == null ? "" : tag.getFrequency();
                 }
                 return "";
             })
                 .setSetter(val -> {
-                    ICoverable te = getUIBuildContext().getTile();
-                    if (!frequencyField.isClient() && te instanceof IFluidHandler) {
-                        UUID uuid;
+                    if (!frequencyField.isClient() && getUIBuildContext().getTile() instanceof IFluidHandler tank) {
+                        UUID uuid = null;
+
                         if (testBit(convert(getCoverData()), PUBLIC_PRIVATE_MASK)) {
                             uuid = getUUID();
-                            if (!(te instanceof IGregTechTileEntity)) return;
-                            if (!uuid.equals(((IGregTechTileEntity) te).getOwnerUuid())) return;
-                        } else {
-                            uuid = null;
+                            if (!(tank instanceof IGregTechTileEntity gte)) return;
+                            if (!uuid.equals(gte.getOwnerUuid())) return;
                         }
-                        EnderWorldSavedData.bindEnderLinkTag((IFluidHandler) te, new EnderLinkTag(val, uuid));
+
+                        EnderWorldSavedData.bindEnderLinkTag(tank, new EnderLinkTag(val, uuid));
                     }
                 })
                 .setTextColor(Color.WHITE.dark(1))
@@ -213,9 +233,30 @@ public class CoverEnderFluidLink extends CoverBehavior {
 
         private int getNewCoverVariable(int id, int coverVariable) {
             switch (id) {
-                case PUBLIC_BUTTON_ID:
-                case PRIVATE_BUTTON_ID:
+                case PUBLIC_BUTTON_ID: {
+                    if (getUIBuildContext().getTile() instanceof IGregTechTileEntity gte) {
+                        EnderLinkTag tag = EnderWorldSavedData.getEnderLinkTag(gte);
+
+                        if (tag != null) {
+                            EnderWorldSavedData.bindEnderLinkTag(gte, new EnderLinkTag(tag.getFrequency(), null));
+                        }
+                    }
+
+                    return toggleBit(coverVariable, PUBLIC_PRIVATE_MASK);
+                }
+                case PRIVATE_BUTTON_ID: {
+                    if (getUIBuildContext().getTile() instanceof IGregTechTileEntity gte) {
+                        EnderLinkTag tag = EnderWorldSavedData.getEnderLinkTag(gte);
+
+                        UUID player = getUUID();
+
+                        if ((gte.getOwnerUuid() == null || Objects.equals(player, gte.getOwnerUuid())) && tag != null) {
+                            EnderWorldSavedData.bindEnderLinkTag(gte, new EnderLinkTag(tag.getFrequency(), player));
+                        }
+                    }
+
                     return toggleBit(coverVariable, PUBLIC_PRIVATE_MASK);
+                }
                 case IMPORT_BUTTON_ID:
                 case EXPORT_BUTTON_ID:
                     return toggleBit(coverVariable, IMPORT_EXPORT_MASK);