Core/Movement: Prevent dynamic changes of movement flags for charmed creatures

Closes #10075
Closes #10079

1 files changed, 21 insertions, 1 deletions

diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 541c5d9f107..9f6b5f76cb4 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -17406,9 +17406,19 @@ bool Unit::SetDisableGravity(bool disable, bool /*packetOnly = false*/)
         return false;
 
     if (disable)
+    {
         AddUnitMovementFlag(MOVEMENTFLAG_DISABLE_GRAVITY);
+        RemoveUnitMovementFlag(MOVEMENTFLAG_FALLING);
+    }
     else
+    {
         RemoveUnitMovementFlag(MOVEMENTFLAG_DISABLE_GRAVITY);
+        if (!HasUnitMovementFlag(MOVEMENTFLAG_CAN_FLY))
+        {
+            m_movementInfo.SetFallTime(0);
+            AddUnitMovementFlag(MOVEMENTFLAG_FALLING);
+        }
+    }
 
     return true;
 }
@@ -17432,9 +17442,19 @@ bool Unit::SetCanFly(bool enable)
         return false;
 
     if (enable)
+    {
         AddUnitMovementFlag(MOVEMENTFLAG_CAN_FLY);
+        RemoveUnitMovementFlag(MOVEMENTFLAG_FALLING);
+    }
     else
-        RemoveUnitMovementFlag(MOVEMENTFLAG_CAN_FLY);
+    {
+        RemoveUnitMovementFlag(MOVEMENTFLAG_CAN_FLY | MOVEMENTFLAG_MASK_MOVING_FLY);
+        if (!IsLevitating())
+        {
+            m_movementInfo.SetFallTime(0);
+            AddUnitMovementFlag(MOVEMENTFLAG_FALLING);
+        }
+    }
 
     return true;
 }