From 702ee5e59664c1351762f975a62844da05c9ba23 Mon Sep 17 00:00:00 2001
From: Ovahlord <dreadkiller@gmx.de>
Date: Mon, 5 Nov 2018 08:55:22 +0100
Subject: [PATCH] Core/Player: update Mount Capabilities uppon changing liquid
 status in order to keep mount auras up to date

---
 src/server/game/Entities/Player/Player.cpp | 45 ++++++++++++----------
 src/server/game/Entities/Player/Player.h   |  3 ++
 src/server/game/Entities/Unit/Unit.cpp     |  7 ++++
 3 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/src/server/game/Entities/Player/Player.cpp b/src/server/game/Entities/Player/Player.cpp
index 48fe22f25f..003c3f896d 100644
--- a/src/server/game/Entities/Player/Player.cpp
+++ b/src/server/game/Entities/Player/Player.cpp
@@ -18846,6 +18846,30 @@ void Player::AddToPlayerPetDataStore(PlayerPetData* playerPetData)
     PlayerPetDataStore.push_back(playerPetData);
 }
 
+void Player::UpdateMountCapabilities()
+{
+    if (HasAuraType(SPELL_AURA_MOUNTED))
+    {
+        Unit::AuraEffectList const& auraList = GetAuraEffectsByType(SPELL_AURA_MOUNTED);
+        if (!auraList.empty())
+        {
+            AuraEffect* mountEffect = auraList.front();
+            if (MountCapabilityEntry const* mountCapability = mountEffect->GetBase()->GetUnitOwner()->GetMountCapability(uint32(mountEffect->GetMiscValueB())))
+            {
+                if (mountCapability->Id != mountEffect->GetAmount())
+                {
+                    Unit* target = mountEffect->GetBase()->GetApplicationOfTarget(GetGUID())->GetTarget();
+                    if (MountCapabilityEntry const* oldMountCapability = sMountCapabilityStore.LookupEntry(mountEffect->GetAmount()))
+                        RemoveAurasDueToSpell(oldMountCapability->SpeedModSpell, target->GetGUID());
+
+                    CastSpell(target, mountCapability->SpeedModSpell, true);
+                    mountEffect->SetAmount(mountCapability->Id);
+                }
+            }
+        }
+    }
+}
+
 void Player::_LoadQuestStatus(PreparedQueryResult result)
 {
     uint16 slot = 0;
@@ -23707,26 +23731,7 @@ void Player::SendInitialPacketsAfterAddToMap()
             auraList.front()->HandleEffect(this, AURA_EFFECT_HANDLE_SEND_FOR_CLIENT, true);
     }
 
-    if (HasAuraType(SPELL_AURA_MOUNTED))
-    {
-        Unit::AuraEffectList const& auraList = GetAuraEffectsByType(SPELL_AURA_MOUNTED);
-        if (!auraList.empty())
-        {
-            AuraEffect* mountEffect = auraList.front();
-            if (MountCapabilityEntry const* mountCapability = mountEffect->GetBase()->GetUnitOwner()->GetMountCapability(uint32(mountEffect->GetMiscValueB())))
-            {
-                if (mountCapability->Id != mountEffect->GetAmount())
-                {
-                    Unit* target = mountEffect->GetBase()->GetApplicationOfTarget(GetGUID())->GetTarget();
-                    if (MountCapabilityEntry const* oldMountCapability = sMountCapabilityStore.LookupEntry(mountEffect->GetAmount()))
-                        RemoveAurasDueToSpell(oldMountCapability->SpeedModSpell, target->GetGUID());
-
-                    CastSpell(target, mountCapability->SpeedModSpell, true);
-                    mountEffect->SetAmount(mountCapability->Id);
-                }
-            }
-        }
-    }
+    UpdateMountCapabilities();
 
     if (HasAuraType(SPELL_AURA_MOD_STUN))
         SetRooted(true);
diff --git a/src/server/game/Entities/Player/Player.h b/src/server/game/Entities/Player/Player.h
index 0221df8c5e..72dfd658c0 100644
--- a/src/server/game/Entities/Player/Player.h
+++ b/src/server/game/Entities/Player/Player.h
@@ -2373,6 +2373,9 @@ class TC_GAME_API Player : public Unit, public GridObject<Player>
         void DeleteFromPlayerPetDataStore(uint32 petNumber);
         void AddToPlayerPetDataStore(PlayerPetData* playerPetData);
 
+        // Mount Capabilites
+        void UpdateMountCapabilities();
+
     protected:
         // Gamemaster whisper whitelist
         GuidList WhisperList;
diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 50ae035f3d..42a9804a0b 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -3304,6 +3304,11 @@ void Unit::ProcessTerrainStatusUpdate(ZLiquidStatus status, Optional<LiquidData>
         Player* player = GetCharmerOrOwnerPlayerOrPlayerItself();
         if (curLiquid && curLiquid->SpellId && (!player || !player->IsGameMaster()))
             CastSpell(this, curLiquid->SpellId, true);
+
+        // Update mount capabilities when changing liquidstatus (enabling / disabling flight auras for example)
+        if (player)
+            player->UpdateMountCapabilities();
+
         _lastLiquid = curLiquid;
     }
 }
@@ -8341,6 +8346,8 @@ MountCapabilityEntry const* Unit::GetMountCapability(uint32 mountType) const
                 if (!(mountCapability->Flags & MOUNT_CAPABILITY_FLAG_GROUND))
                     continue;
             }
+            else if (!(mountCapability->Flags & MOUNT_CAPABILITY_FLAG_UNDERWATER))
+                continue;
         }
         else if (isInWater)
         {