[GEN][ZH] Fix AIGroup memory management and game crash when a player is selected in Replay playback due heap-use-after-free in GameLogic::logicMessageDispatcher()

Generals/Code/GameEngine/Include/Common/Player.h

@@ -608,7 +608,7 @@ class Player : public Snapshot
 	// add to the player's current selection this hotkey team.
 	void processAddTeamGameMessage(Int hotkeyNum, GameMessage *msg);
 
-	// returns an AIGroup object that is the currently selected group.
+	// fills an AIGroup object that is the currently selected group.
 	void getCurrentSelectionAsAIGroup(AIGroup *group);
 
 	// sets the currently selected group to be the given AIGroup

Generals/Code/GameEngine/Include/GameLogic/AI.h

@@ -37,6 +37,8 @@
 #include "Common/GameType.h"
 #include "GameLogic/Damage.h"
 #include "Common/STLTypedefs.h"
+#include "refcount.h"
+#include "ref_ptr.h"
 
 class AIGroup;
 class AttackPriorityInfo;
@@ -265,7 +267,7 @@ class AI : public SubsystemInterface, public Snapshot
 	void loadPostProcess( void );
 
 	// AI Groups -----------------------------------------------------------------------------------------------
-	AIGroup *createGroup( void );					///< instantiate a new AI Group
+	RefCountPtr<AIGroup> createGroup( void );	///< instantiate a new AI Group
 	void destroyGroup( AIGroup *group );	///< destroy the given AI Group
 	AIGroup *findGroup( UnsignedInt id );	///< return the AI Group with the given ID
 
@@ -851,6 +853,10 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	void xfer( Xfer *xfer );
 	void loadPostProcess( void );
 
+	void Add_Ref() const { m_refCount.Add_Ref(); }
+	void Release_Ref() const { m_refCount.Release_Ref(destroy, this); }
+	UnsignedShort Num_Refs() const { return m_refCount.Num_Refs(); }
+
 	void groupMoveToPosition( const Coord3D *pos, Bool addWaypoint, CommandSourceType cmdSource );
 	void groupMoveToAndEvacuate( const Coord3D *pos, CommandSourceType cmdSource );			///< move to given position(s)
 	void groupMoveToAndEvacuateAndExit( const Coord3D *pos, CommandSourceType cmdSource );			///< move to given position & unload transport.
@@ -938,13 +944,15 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 
 	void add( Object *obj );								///< add object to group
 
-	// returns true if remove destroyed the group for us.
+	// Returns true if the group was emptied.
 	Bool remove( Object *obj);
+
+	void removeAll( void );
 
 	// If the group contains any objects not owned by ownerPlayer, return TRUE.
 	Bool containsAnyObjectsNotOwnedByPlayer( const Player *ownerPlayer );
 
-	// Remove any objects that aren't owned by the player, and return true if the group was destroyed due to emptiness
+	// Removes any objects that aren't owned by the player, and returns true if the group was emptied.
 	Bool removeAnyObjectsNotOwnedByPlayer( const Player *ownerPlayer );
 
 	UnsignedInt getID( void );
@@ -974,6 +982,11 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	friend class AI;
 	AIGroup( void );
 
+	static void destroy(AIGroup* self)
+	{
+		TheAI->destroyGroup(self);
+	}
+
 	void recompute( void );									///< recompute various group info, such as speed, leader, etc
 
 	ListObjectPtr m_memberList;							///< the list of member Objects
@@ -982,6 +995,8 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	Real m_speed;														///< maximum speed of group (slowest member)
 	Bool m_dirty;														///< "dirty bit" - if true then group speed, leader, needs recompute
 
+	RefCountValue<UnsignedShort> m_refCount;	///< the reference counter
+
 	UnsignedInt m_id;												///< the unique ID of this group
 	Path *m_groundPath;											///< Group ground path.
 

Generals/Code/GameEngine/Include/GameLogic/Object.h

@@ -31,6 +31,7 @@
 #define _OBJECT_H_
 
 #include "Lib/BaseType.h"
+#include "ref_ptr.h"
 
 #include "Common/Geometry.h"
 #include "Common/Snapshot.h"
@@ -663,7 +664,7 @@ class Object : public Thing, public Snapshot
 
 	GeometryInfo	m_geometryInfo;
 
-	AIGroup*			m_group;								///< if non-NULL, we are part of this group of agents
+	RefCountPtr<AIGroup> m_group;								///< if non-NULL, we are part of this group of agents
 
 	// These will last for my lifetime.  I will reuse them and reset them.  The truly dynamic ones are in PartitionManager
 	SightingInfo	*m_partitionLastLook;		///< Where and for whom I last looked, so I can undo its effects when I stop

Generals/Code/GameEngine/Source/GameLogic/AI/AI.cpp

@@ -330,18 +330,11 @@ void AI::reset( void )
 		m_aiData = m_aiData->m_next;
 		delete cur;
 	}
-	while (m_groupList.size())
-	{
-		AIGroup *groupToRemove = m_groupList.front();
-		if (groupToRemove)
-		{
-			destroyGroup(groupToRemove);
-		}
-		else
-		{
-			m_groupList.pop_front(); // NULL group, just kill from list.  Shouldn't really happen, but just in case.
-		}
-	}
+
+	DEBUG_ASSERTCRASH(m_groupList.empty(), ("AI::m_groupList is expected empty already\n"));
+
+	m_groupList.clear(); // Clear just in case...
+
 	m_nextGroupID = 0;
 	m_nextFormationID = NO_FORMATION_ID;
 	getNextFormationID(); // increment once past NO_FORMATION_ID.  jba.
@@ -440,14 +433,14 @@ void AI::parseAiDataDefinition( INI* ini )
 /**
  * Create a new AI Group
  */
-AIGroup *AI::createGroup( void )
+RefCountPtr<AIGroup> AI::createGroup( void )
 {
 	// create a new instance
-	AIGroup *group = newInstance(AIGroup);
+	RefCountPtr<AIGroup> group = RefCountPtr<AIGroup>::Create_NoAddRef(newInstance(AIGroup));
 
 	// add it to the list
 //	DEBUG_LOG(("***AIGROUP %x is being added to m_groupList.\n", group ));
-	m_groupList.push_back( group );
+	m_groupList.push_back( group.Peek() );
 
 	return group;
 }

Generals/Code/GameEngine/Source/GameLogic/AI/AIGroup.cpp

@@ -92,20 +92,8 @@ AIGroup::~AIGroup()
 {
 //	DEBUG_LOG(("***AIGROUP %x is being destructed.\n", this));
 	// disassociate each member from the group
-	std::list<Object *>::iterator i;
-	for( i = m_memberList.begin(); i != m_memberList.end(); /* empty */ )
-	{
-		Object *member = *i;
-		if (member)
-		{
-			member->leaveGroup();
-			i = m_memberList.begin();	// jump back to the beginning, cause ai->leaveGroup will remove this element. 
-		}
-		else
-		{
-			i = m_memberList.erase(i);
-		}
-	}
+	removeAll();
+
 	if (m_groundPath) {
 		deleteInstance(m_groundPath);
 		m_groundPath = NULL;
@@ -223,15 +211,33 @@ Bool AIGroup::remove( Object *obj )
 	// list has changed, properties need recomputation
 	m_dirty = true;
 
-	// if the group is empty, no-one is using it any longer, so destroy it
 	if (isEmpty()) {
-		TheAI->destroyGroup( this );
 		return TRUE;
 	}
 
 	return FALSE;
 }
 
+/**
+ * Remove all objects from group
+ */
+void AIGroup::removeAll( void )
+{
+	std::list<Object *> memberList;
+	memberList.swap(m_memberList);
+	m_memberListSize = 0;
+
+	std::list<Object *>::iterator i;
+	for ( i = memberList.begin(); i != memberList.end(); ++i )
+	{
+		Object *member = *i;
+		if (member)
+			member->leaveGroup();
+	}
+
+	m_dirty = true;
+}
+
 /**
  * If the group contains any objects not owned by ownerPlayer, return TRUE.
  */

Generals/Code/GameEngine/Source/GameLogic/AI/AIPlayer.cpp

@@ -893,11 +893,11 @@ void AIPlayer::guardSupplyCenter( Team *team, Int minSupplies )
 	}
 	if (warehouse) {
 
-		AIGroup* theGroup = TheAI->createGroup();
+		RefCountPtr<AIGroup> theGroup = TheAI->createGroup();
 		if (!theGroup) {
 			return;
 		}
-		team->getTeamAsAIGroup(theGroup);
+		team->getTeamAsAIGroup(theGroup.Peek());
 		Coord3D location = *warehouse->getPosition();
 		// It's probably a defensive move - position towards the enemy.
 		Region2D bounds;
@@ -2467,9 +2467,9 @@ void AIPlayer::checkReadyTeams( void )
 					/*
 					if (team->m_team->getPrototype()->getTemplateInfo()->m_hasHomeLocation && 
 							!team->m_reinforcement) {
- 						AIGroup* theGroup = TheAI->createGroup();
+						RefCountPtr<AIGroup> theGroup = TheAI->createGroup();
 						if (theGroup) {
-							team->m_team->getTeamAsAIGroup(theGroup);
+							team->m_team->getTeamAsAIGroup(theGroup.Peek());
 							Coord3D destination = team->m_team->getPrototype()->getTemplateInfo()->m_homeLocation;
 							theGroup->groupTightenToPosition( &destination, false, CMD_FROM_AI );
 							team->m_frameStarted = TheGameLogic->getFrame();

Generals/Code/GameEngine/Source/GameLogic/AI/AIStates.cpp

@@ -4050,8 +4050,8 @@ StateReturnType AIFollowWaypointPathState::update()
 	if (m_moveAsGroup) {
 		if (obj->getControllingPlayer()->isSkirmishAIPlayer()) {
 			Team *team = obj->getTeam();
-			AIGroup *group = TheAI->createGroup();
-			team->getTeamAsAIGroup(group);
+			RefCountPtr<AIGroup> group = TheAI->createGroup();
+			team->getTeamAsAIGroup(group.Peek());
 
 			Coord3D pos;
 			group->getCenter(&pos);

Generals/Code/GameEngine/Source/GameLogic/Object/Object.cpp

@@ -563,9 +563,8 @@ Object::~Object()
 		ThePartitionManager->unRegisterObject( this );
 
 	// if we are in a group, remove us
-	if (m_group)
-		m_group->remove( this );
-
+	leaveGroup();
+
 	// note, do NOT free these, there are just a shadow copy!
 	m_ai = NULL;
 	m_physics = NULL;
@@ -3841,7 +3840,7 @@ void Object::xfer( Xfer *xfer )
 	}
 
 	// Doesn't need to be saved.  These are created as needed.  jba.
-	//AIGroup*		m_group;															///< if non-NULL, we are part of this group of agents
+	//m_group;
 
 	// don't need to save m_partitionData.
 	DEBUG_ASSERTCRASH(!(xfer->getXferMode() == XFER_LOAD && m_partitionData == NULL), ("should not be in partitionmgr yet"));
@@ -5446,7 +5445,7 @@ RadarPriorityType Object::getRadarPriority( void ) const
 // ------------------------------------------------------------------------------------------------
 AIGroup *Object::getGroup(void)
 {
-	return m_group;
+	return m_group.Peek();
 }
 
 //-------------------------------------------------------------------------------------------------
@@ -5456,7 +5455,7 @@ void Object::enterGroup( AIGroup *group )
 	// if we are in another group, remove ourselves from it first
 	leaveGroup();
 
-	m_group = group;
+	m_group = RefCountPtr<AIGroup>::Create_AddRef(group);
 }
 
 //-------------------------------------------------------------------------------------------------
@@ -5467,7 +5466,7 @@ void Object::leaveGroup( void )
 	if (m_group)
 	{
 		// to avoid recursion, set m_group to NULL before removing
-		AIGroup *group = m_group;
+		RefCountPtr<AIGroup> group = m_group;
 		m_group = NULL;
 		group->remove( this );
 	}