Core/Units: Fixed crash happening when charm was removed by its own charmed AI during update

(cherry picked from commit c7445669e8577884fd512fe009f7d3da4cf0429a)

2 files changed, 9 insertions, 3 deletions

diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 24686921572..0ea5d683234 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -9076,7 +9076,6 @@ void Unit::SetAI(UnitAI* newAI)
 
 void Unit::ScheduleAIChange()
 {
-    ASSERT(!m_aiLocked, "Attempt to schedule AI change during AI update tick");
     bool const charmed = IsCharmed();
     // if charm is applied, we can't have disabled AI already, and vice versa
     if (charmed)
@@ -9086,15 +9085,20 @@ void Unit::ScheduleAIChange()
 
     if (charmed)
         i_disabledAI = std::move(i_AI);
+    else if (m_aiLocked)
+    {
+        ASSERT(!i_lockedAILifetimeExtension, "Attempt to schedule multiple charm AI changes during one update");
+        i_lockedAILifetimeExtension = std::move(i_AI); // AI needs to live just a bit longer to finish its UpdateAI
+    }
     else
         i_AI.reset();
 }
 
 void Unit::RestoreDisabledAI()
 {
-    ASSERT(!m_aiLocked, "Attempt to restore AI during UpdateAI tick");
     ASSERT((GetTypeId() == TYPEID_PLAYER) || i_disabledAI, "Attempt to restore disabled AI on creature without disabled AI");
     i_AI = std::move(i_disabledAI);
+    i_lockedAILifetimeExtension.reset();
 }
 
 void Unit::AddToWorld()
diff --git a/src/server/game/Entities/Unit/Unit.h b/src/server/game/Entities/Unit/Unit.h
index ab1f5896009..93065c19921 100644
--- a/src/server/game/Entities/Unit/Unit.h
+++ b/src/server/game/Entities/Unit/Unit.h
@@ -1970,7 +1970,9 @@ class TC_GAME_API Unit : public WorldObject
 
         void UpdateCharmAI();
         void RestoreDisabledAI();
-        std::unique_ptr<UnitAI> i_AI, i_disabledAI;
+        std::unique_ptr<UnitAI> i_AI;
+        std::unique_ptr<UnitAI> i_disabledAI;
+        std::unique_ptr<UnitAI> i_lockedAILifetimeExtension; // yes, this lifetime extension is terrible
         bool m_aiLocked;
 
         std::unordered_set<AbstractFollower*> m_followingMe;