MULTIDAMAGE: make check entity safe

ApplyMultiDamage: Fixed potential crash when in TakeDamage hook causes another additional damage
rehlds · May 31, 2024 · c7be8bf · c7be8bf
1 parent 7372573
commit c7be8bf
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/regamedll/dlls/weapons.cpp b/regamedll/dlls/weapons.cpp
@@ -79,7 +79,7 @@ LINK_HOOK_VOID_CHAIN2(ClearMultiDamage)
 // Resets the global multi damage accumulator
 void EXT_FUNC __API_HOOK(ClearMultiDamage)()
 {
-	gMultiDamage.pEntity = nullptr;
+	gMultiDamage.hEntity = nullptr;
 	gMultiDamage.amount = 0;
 	gMultiDamage.type = 0;
 }
@@ -89,11 +89,15 @@ LINK_HOOK_VOID_CHAIN(ApplyMultiDamage, (entvars_t *pevInflictor, entvars_t *pevA
 // Inflicts contents of global multi damage register on gMultiDamage.pEntity
 void EXT_FUNC __API_HOOK(ApplyMultiDamage)(entvars_t *pevInflictor, entvars_t *pevAttacker)
 {
-	if (!gMultiDamage.pEntity)
+	EntityHandle<CBaseEntity> hEnt = gMultiDamage.hEntity;
+	if (!hEnt)
 		return;
 
-	gMultiDamage.pEntity->TakeDamage(pevInflictor, pevAttacker, gMultiDamage.amount, gMultiDamage.type);
-	gMultiDamage.pEntity->ResetDmgPenetrationLevel();
+	hEnt->TakeDamage(pevInflictor, pevAttacker, gMultiDamage.amount, gMultiDamage.type);
+
+	// check again, the entity may be removed after taking damage
+	if (hEnt)
+		hEnt->ResetDmgPenetrationLevel();
 }
 
 LINK_HOOK_VOID_CHAIN(AddMultiDamage, (entvars_t *pevInflictor, CBaseEntity *pEntity, float flDamage, int bitsDamageType), pevInflictor, pEntity, flDamage, bitsDamageType)
@@ -105,17 +109,17 @@ void EXT_FUNC __API_HOOK(AddMultiDamage)(entvars_t *pevInflictor, CBaseEntity *p
 
 	gMultiDamage.type |= bitsDamageType;
 
-	if (pEntity != gMultiDamage.pEntity)
+	if (pEntity != gMultiDamage.hEntity)
 	{
 #ifdef REGAMEDLL_FIXES
-		if (gMultiDamage.pEntity) // avoid api calls with null default pEntity
+		if (gMultiDamage.hEntity) // avoid api calls with null default pEntity
 #endif
 		{
 			// UNDONE: wrong attacker!
 			ApplyMultiDamage(pevInflictor, pevInflictor);
 		}
 
-		gMultiDamage.pEntity = pEntity;
+		gMultiDamage.hEntity = pEntity;
 		gMultiDamage.amount = 0;
 	}
 

diff --git a/regamedll/dlls/weapons.h b/regamedll/dlls/weapons.h
@@ -133,7 +133,7 @@ struct AmmoInfo
 
 struct MULTIDAMAGE
 {
-	CBaseEntity *pEntity;
+	EntityHandle<CBaseEntity> hEntity;
 	float amount;
 	int type;
 };