tf2: fix shields not restoring charge with pickups when mid charge

Source: ValveSoftware/source-sdk-2013#1896

Author: Generalisk

src/game/server/tf/entity_ammopack.cpp

@@ -9,6 +9,7 @@
 #include "tf_shareddefs.h"
 #include "tf_player.h"
 #include "tf_team.h"
+#include "tf_wearable_weapons.h"
 #include "engine/IEngineSound.h"
 #include "entity_ammopack.h"
 #include "tf_gamestats.h"
@@ -104,7 +105,16 @@ bool CAmmoPack::MyTouch( CBasePlayer *pPlayer )
 				{
 					flPackRatio *= 0.2;
 				}
+
 				pTFPlayer->m_Shared.SetDemomanChargeMeter( flCurrentCharge + flPackRatio * 100.0f );
+
+				// Update the shield charge condition to reflect the ammo pack increase
+				if ( pTFPlayer->m_Shared.InCond( TF_COND_SHIELD_CHARGE ) )
+				{
+					float flChargeTime = ( pTFPlayer->m_Shared.GetDemomanChargeMeter() / 100.0f ) + flPackRatio;
+
+					pTFPlayer->m_Shared.SetConditionDuration( TF_COND_SHIELD_CHARGE, flChargeTime );
+				}
 				bSuccess = true;
 			}
 		}

src/game/server/tf/tf_ammo_pack.cpp

@@ -10,6 +10,7 @@
 #include "ammodef.h"
 #include "tf_gamerules.h"
 #include "explode.h"
+#include "tf_wearable_weapons.h"
 #include "tf_gamestats.h"
 
 // memdbgon must be the last include file in a .cpp file!!!
@@ -395,7 +396,17 @@ void CTFAmmoPack::PackTouch( CBaseEntity *pOther )
 			{
 				m_flAmmoRatio *= 0.2;
 			}
+
 			pPlayer->m_Shared.SetDemomanChargeMeter( flCurrentCharge + m_flAmmoRatio * 100.0f );
+
+			// Update the shield charge condition to reflect the ammo pack increase
+			if ( pPlayer->m_Shared.InCond( TF_COND_SHIELD_CHARGE ) )
+			{
+				float flChargeTime = (pPlayer->m_Shared.GetDemomanChargeMeter() / 100.0f ) + m_flAmmoRatio;
+
+				pPlayer->m_Shared.SetConditionDuration( TF_COND_SHIELD_CHARGE, flChargeTime );
+			}
+
 			iAmmoTaken++;
 		}
 	}

src/game/server/tf/tf_player.cpp

@@ -4100,6 +4100,17 @@ void CTFPlayer::Regenerate( bool bRefillHealthAndAmmo /*= true*/ )
 		m_Shared.SetScoutEnergyDrinkMeter( 100.0f );
 		m_Shared.SetDemomanChargeMeter( 100.0f );
 
+		// Update the shield charge condition to reflect the reset for charge
+		if ( m_Shared.InCond(TF_COND_SHIELD_CHARGE ) )
+		{
+			float flChargeTime = 1.5f;
+			CALL_ATTRIB_HOOK_FLOAT_ON_OTHER( this, flChargeTime, mod_charge_time );
+
+			flChargeTime = m_Shared.GetConditionDuration( TF_COND_SHIELD_CHARGE ) + flChargeTime;
+
+			m_Shared.SetConditionDuration( TF_COND_SHIELD_CHARGE, flChargeTime );
+		}
+
 		// Selectively refill the item effect meters if they're allowed
 		for( int i = FIRST_LOADOUT_SLOT_WITH_CHARGE_METER; i <= LAST_LOADOUT_SLOT_WITH_CHARGE_METER; ++i )
 		{
@@ -14851,6 +14862,17 @@ void CTFPlayer::CheatImpulseCommands( int iImpulse )
 				m_Shared.m_flRageMeter = 100.f;
 				m_Shared.SetDemomanChargeMeter( 100.f );
 
+				// Update the shield charge condition to reflect the reset for charge
+				if ( m_Shared.InCond( TF_COND_SHIELD_CHARGE ) )
+				{
+					float flChargeTime = 1.5f;
+					CALL_ATTRIB_HOOK_FLOAT_ON_OTHER( this, flChargeTime, mod_charge_time );
+
+					flChargeTime = m_Shared.GetConditionDuration( TF_COND_SHIELD_CHARGE ) + flChargeTime;
+
+					m_Shared.SetConditionDuration( TF_COND_SHIELD_CHARGE, flChargeTime );
+				}
+
 				for( int i = FIRST_LOADOUT_SLOT_WITH_CHARGE_METER; i <= LAST_LOADOUT_SLOT_WITH_CHARGE_METER; ++i )
 				{
 					m_Shared.SetItemChargeMeter( (loadout_positions_t)i, 100.f );