From a5aa9e7281b0471078d5b8e2d9e9d5f402612c50 Mon Sep 17 00:00:00 2001
From: Matsv <mats@matsv.nl>
Date: Wed, 21 Jun 2017 19:30:48 +0200
Subject: [PATCH] Handle nulls in EntityRewriter (#11)

---
 .../viabackwards/api/entities/storage/EntityTracker.java | 7 ++++---
 .../matsv/viabackwards/api/rewriters/EntityRewriter.java | 9 ++++++++-
 .../protocol1_10to1_11/packets/BlockItemPackets1_11.java | 8 +++++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/nl/matsv/viabackwards/api/entities/storage/EntityTracker.java b/core/src/main/java/nl/matsv/viabackwards/api/entities/storage/EntityTracker.java
index 88295ed5..55906747 100644
--- a/core/src/main/java/nl/matsv/viabackwards/api/entities/storage/EntityTracker.java
+++ b/core/src/main/java/nl/matsv/viabackwards/api/entities/storage/EntityTracker.java
@@ -19,6 +19,7 @@ import us.myles.ViaVersion.api.data.StoredObject;
 import us.myles.ViaVersion.api.data.UserConnection;
 
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class EntityTracker extends StoredObject {
@@ -51,12 +52,12 @@ public class EntityTracker extends StoredObject {
 
         public AbstractEntityType getEntityType(int id) {
             if (containsEntity(id))
-                return getEntity(id).getType();
+                return getEntity(id).get().getType();
             return null;
         }
 
-        public StoredEntity getEntity(int id) {
-            return entityMap.get(id);
+        public Optional<StoredEntity> getEntity(int id) {
+            return Optional.ofNullable(entityMap.get(id));
         }
 
         public boolean containsEntity(int id) {
diff --git a/core/src/main/java/nl/matsv/viabackwards/api/rewriters/EntityRewriter.java b/core/src/main/java/nl/matsv/viabackwards/api/rewriters/EntityRewriter.java
index 24b974d8..6a314ac9 100644
--- a/core/src/main/java/nl/matsv/viabackwards/api/rewriters/EntityRewriter.java
+++ b/core/src/main/java/nl/matsv/viabackwards/api/rewriters/EntityRewriter.java
@@ -29,6 +29,7 @@ import us.myles.ViaVersion.api.data.UserConnection;
 import us.myles.ViaVersion.api.minecraft.metadata.MetaType;
 import us.myles.ViaVersion.api.minecraft.metadata.Metadata;
 import us.myles.ViaVersion.api.minecraft.metadata.types.MetaType1_9;
+import us.myles.ViaVersion.exception.CancelException;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -98,7 +99,13 @@ public abstract class EntityRewriter<T extends BackwardsProtocol> extends Rewrit
     }
 
     protected MetaStorage handleMeta(UserConnection user, int entityId, MetaStorage storage) throws Exception {
-        EntityTracker.StoredEntity entity = getEntityTracker(user).getEntity(entityId);
+        Optional<EntityTracker.StoredEntity> optEntity = getEntityTracker(user).getEntity(entityId);
+        if (!optEntity.isPresent()) {
+            ViaBackwards.getPlatform().getLogger().warning("Metadata for entity id: " + entityId + " not sent because the entity doesn't exist. " + storage);
+            throw new CancelException();
+        }
+
+        EntityTracker.StoredEntity entity = optEntity.get();
         AbstractEntityType type = entity.getType();
 
         List<Metadata> newList = new CopyOnWriteArrayList<>();
diff --git a/core/src/main/java/nl/matsv/viabackwards/protocol/protocol1_10to1_11/packets/BlockItemPackets1_11.java b/core/src/main/java/nl/matsv/viabackwards/protocol/protocol1_10to1_11/packets/BlockItemPackets1_11.java
index cca1eb2a..06d16847 100644
--- a/core/src/main/java/nl/matsv/viabackwards/protocol/protocol1_10to1_11/packets/BlockItemPackets1_11.java
+++ b/core/src/main/java/nl/matsv/viabackwards/protocol/protocol1_10to1_11/packets/BlockItemPackets1_11.java
@@ -428,7 +428,8 @@ public class BlockItemPackets1_11 extends BlockItemRewriter<Protocol1_10To1_11>
         WindowTracker tracker = user.get(WindowTracker.class);
         if (tracker.getInventory() != null && tracker.getInventory().equals("EntityHorse")) {
             EntityTracker.ProtocolEntityTracker entTracker = user.get(EntityTracker.class).get(getProtocol());
-            if (tracker.getEntityId() != -1 && entTracker.getEntity(tracker.getEntityId()).getType().is(EntityType1_11.EntityType.LIAMA))
+            Optional<EntityTracker.StoredEntity> optEntity = entTracker.getEntity(tracker.getEntityId());
+            if (optEntity.isPresent() && optEntity.get().getType().is(EntityType1_11.EntityType.LIAMA))
                 return true;
         }
         return false;
@@ -438,8 +439,9 @@ public class BlockItemPackets1_11 extends BlockItemRewriter<Protocol1_10To1_11>
         WindowTracker tracker = user.get(WindowTracker.class);
         if (tracker.getInventory() != null && tracker.getInventory().equals("EntityHorse")) {
             EntityTracker.ProtocolEntityTracker entTracker = user.get(EntityTracker.class).get(getProtocol());
-            if (tracker.getEntityId() != -1)
-                return Optional.of(entTracker.getEntity(tracker.getEntityId()).get(ChestedHorseStorage.class));
+            Optional<EntityTracker.StoredEntity> optEntity = entTracker.getEntity(tracker.getEntityId());
+            if (optEntity.isPresent())
+                return Optional.of(optEntity.get().get(ChestedHorseStorage.class));
         }
         return Optional.empty();
     }