From 4c409f98f397ad97d160de9155ca728d3afadd6e Mon Sep 17 00:00:00 2001
From: Camotoy <20743703+Camotoy@users.noreply.github.com>
Date: Mon, 20 Dec 2021 22:54:34 -0500
Subject: [PATCH] Properly implement piglin/hoglin dimension safety

Whether they shake or not is now controlled by the server, as implemented in Java Edition.
---
 .../geyser/entity/type/living/animal/HoglinEntity.java      | 5 ++---
 .../geyser/entity/type/living/monster/BasePiglinEntity.java | 3 +--
 .../java/org/geysermc/geyser/session/GeyserSession.java     | 6 ++++++
 core/src/main/java/org/geysermc/geyser/util/ChunkUtils.java | 3 +++
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/HoglinEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/HoglinEntity.java
index 1cbbf5d1d..5ac9e8bbd 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/HoglinEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/HoglinEntity.java
@@ -29,9 +29,8 @@ import com.github.steveice10.mc.protocol.data.game.entity.metadata.type.BooleanE
 import com.nukkitx.math.vector.Vector3f;
 import com.nukkitx.protocol.bedrock.data.entity.EntityFlag;
 import org.geysermc.geyser.entity.EntityDefinition;
-import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.registry.type.ItemMapping;
-import org.geysermc.geyser.util.DimensionUtils;
+import org.geysermc.geyser.session.GeyserSession;
 
 import java.util.UUID;
 
@@ -50,7 +49,7 @@ public class HoglinEntity extends AnimalEntity {
 
     @Override
     protected boolean isShaking() {
-        return (!isImmuneToZombification && !session.getDimension().equals(DimensionUtils.NETHER)) || super.isShaking();
+        return (!isImmuneToZombification && !session.isDimensionPiglinSafe()) || super.isShaking();
     }
 
     @Override
diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/monster/BasePiglinEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/monster/BasePiglinEntity.java
index c90839ef9..526ff41ac 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/monster/BasePiglinEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/monster/BasePiglinEntity.java
@@ -30,7 +30,6 @@ import com.nukkitx.math.vector.Vector3f;
 import com.nukkitx.protocol.bedrock.data.entity.EntityFlag;
 import org.geysermc.geyser.entity.EntityDefinition;
 import org.geysermc.geyser.session.GeyserSession;
-import org.geysermc.geyser.util.DimensionUtils;
 
 import java.util.UUID;
 
@@ -49,6 +48,6 @@ public class BasePiglinEntity extends MonsterEntity {
 
     @Override
     protected boolean isShaking() {
-        return (!isImmuneToZombification && !session.getDimension().equals(DimensionUtils.NETHER)) || super.isShaking();
+        return (!isImmuneToZombification && !session.isDimensionPiglinSafe()) || super.isShaking();
     }
 }
diff --git a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
index c67299466..4cd9cd69f 100644
--- a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
+++ b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
@@ -302,6 +302,12 @@ public class GeyserSession implements GeyserConnection, CommandSender {
      */
     @Setter
     private String dimension = DimensionUtils.OVERWORLD;
+    /**
+     * Whether piglins and hoglins are safe from conversion in this dimension.
+     * This controls if they have the shaking effect applied in the dimension.
+     */
+    @Setter
+    private boolean dimensionPiglinSafe;
 
     @Setter
     private int breakingBlock;
diff --git a/core/src/main/java/org/geysermc/geyser/util/ChunkUtils.java b/core/src/main/java/org/geysermc/geyser/util/ChunkUtils.java
index 3e56bc7e8..9dcb9f303 100644
--- a/core/src/main/java/org/geysermc/geyser/util/ChunkUtils.java
+++ b/core/src/main/java/org/geysermc/geyser/util/ChunkUtils.java
@@ -265,5 +265,8 @@ public class ChunkUtils {
         // Load world coordinate scale for the world border
         double coordinateScale = ((Number) dimensionTag.get("coordinate_scale").getValue()).doubleValue();
         session.getWorldBorder().setWorldCoordinateScale(coordinateScale);
+
+        // Set if piglins/hoglins should shake
+        session.setDimensionPiglinSafe(((Number) dimensionTag.get("piglin_safe").getValue()).byteValue() != (byte) 0);
     }
 }