From 86d45754c3842f810b59fce59322a75a1e1e3f96 Mon Sep 17 00:00:00 2001
From: libraryaddict <github@lib.co.nz>
Date: Fri, 24 May 2019 19:20:55 +1200
Subject: [PATCH] Readded set/get horse armor. Added sanitizing of saved
 disguises loading + logging.

---
 .../disguisetypes/watchers/HorseWatcher.java  |  8 +++++++
 .../utilities/json/SerializerFlagWatcher.java | 24 ++++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/main/java/me/libraryaddict/disguise/disguisetypes/watchers/HorseWatcher.java b/src/main/java/me/libraryaddict/disguise/disguisetypes/watchers/HorseWatcher.java
index 74604818..59d9cf7f 100644
--- a/src/main/java/me/libraryaddict/disguise/disguisetypes/watchers/HorseWatcher.java
+++ b/src/main/java/me/libraryaddict/disguise/disguisetypes/watchers/HorseWatcher.java
@@ -34,4 +34,12 @@ public class HorseWatcher extends AbstractHorseWatcher {
         setData(MetaIndex.HORSE_COLOR, getColor().ordinal() & 0xFF | style.ordinal() << 8);
         sendData(MetaIndex.HORSE_COLOR);
     }
+
+    public void setHorseArmor(ItemStack item) {
+        getEquipment().setChestplate(item);
+    }
+
+    public ItemStack getHorseArmor() {
+        return getEquipment().getChestplate();
+    }
 }
diff --git a/src/main/java/me/libraryaddict/disguise/utilities/json/SerializerFlagWatcher.java b/src/main/java/me/libraryaddict/disguise/utilities/json/SerializerFlagWatcher.java
index ec6ad1dd..2415b299 100644
--- a/src/main/java/me/libraryaddict/disguise/utilities/json/SerializerFlagWatcher.java
+++ b/src/main/java/me/libraryaddict/disguise/utilities/json/SerializerFlagWatcher.java
@@ -4,10 +4,12 @@ import com.comphenix.protocol.wrappers.WrappedBlockData;
 import com.comphenix.protocol.wrappers.WrappedChatComponent;
 import com.comphenix.protocol.wrappers.WrappedGameProfile;
 import com.google.gson.*;
+import com.google.gson.internal.LinkedTreeMap;
 import com.google.gson.reflect.TypeToken;
 import com.mojang.authlib.properties.PropertyMap;
 import me.libraryaddict.disguise.disguisetypes.*;
 import me.libraryaddict.disguise.disguisetypes.watchers.ArrowWatcher;
+import me.libraryaddict.disguise.utilities.DisguiseUtilities;
 import org.bukkit.inventory.ItemStack;
 
 import java.lang.reflect.Field;
@@ -61,11 +63,12 @@ public class SerializerFlagWatcher implements JsonDeserializer<FlagWatcher>, Jso
         Field field = FlagWatcher.class.getDeclaredField(name);
         field.setAccessible(true);
         HashMap<Integer, Object> map = (HashMap<Integer, Object>) field.get(watcher);
+        int count = 0;
 
         for (Map.Entry<Integer, Object> entry : map.entrySet()) {
-            if (entry.getValue() instanceof Double) {
-                MetaIndex index = MetaIndex.getMetaIndex(flagWatcher, entry.getKey());
+            MetaIndex index = MetaIndex.getMetaIndex(flagWatcher, entry.getKey());
 
+            if (entry.getValue() instanceof Double) {
                 Object def = index.getDefault();
 
                 if (def instanceof Long)
@@ -78,15 +81,24 @@ public class SerializerFlagWatcher implements JsonDeserializer<FlagWatcher>, Jso
                     entry.setValue(((Double) entry.getValue()).shortValue());
                 else if (def instanceof Byte)
                     entry.setValue(((Double) entry.getValue()).byteValue());
-            } else if (entry.getValue() instanceof Map) {
-                MetaIndex index = MetaIndex.getMetaIndex(flagWatcher, entry.getKey());
-
+            } else if (entry.getValue() instanceof LinkedTreeMap) { // If it's deserialized incorrectly as a map
+                // If the default value is not VillagerData
                 if (!(index.getDefault() instanceof VillagerData)) {
                     continue;
                 }
 
-                entry.setValue(new Gson().fromJson(new Gson().toJson(entry.getValue()),VillagerData.class));
+                entry.setValue(new Gson().fromJson(new Gson().toJson(entry.getValue()), VillagerData.class));
             }
+
+            // If the deserialized class is not the same class type as the default
+            if (!index.getDefault().getClass().isInstance(entry.getValue())) {
+                entry.setValue(index.getDefault());
+                count++;
+            }
+        }
+
+        if (count > 0) {
+            DisguiseUtilities.getLogger().info("Fixed " + count + " incorrect disguise flags on saved disguise");
         }
     }