From 6073d9a87caba440001cd91cc6d55398be494465 Mon Sep 17 00:00:00 2001
From: MJSailor <92106367+kurokoTurbo@users.noreply.github.com>
Date: Mon, 12 Aug 2024 10:46:30 +0000
Subject: [PATCH] fix charged spells vulnerability (#612)

---
 .../Systems/Actions/ChargeActionSystem.cs          |  2 +-
 .../_White/Wizard/Charging/ChargingSystem.cs       | 14 ++++++++++++--
 Content.Shared/_White/Wizard/WizardEvents.cs       |  6 +++---
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/Content.Client/UserInterface/Systems/Actions/ChargeActionSystem.cs b/Content.Client/UserInterface/Systems/Actions/ChargeActionSystem.cs
index 29a00acff3..e8798bcfcd 100644
--- a/Content.Client/UserInterface/Systems/Actions/ChargeActionSystem.cs
+++ b/Content.Client/UserInterface/Systems/Actions/ChargeActionSystem.cs
@@ -132,7 +132,7 @@ public sealed class ChargeActionSystem : SharedChargingSystem
         {
             if (_chargeLevel > 0 && _charging)
             {
-                RaiseNetworkEvent(new AddWizardChargeEvent(action.ChargeProto));
+                RaiseNetworkEvent(new AddWizardChargeEvent(GetNetEntity(actionId)));
             }
             _prevChargeLevel = _chargeLevel;
         }
diff --git a/Content.Server/_White/Wizard/Charging/ChargingSystem.cs b/Content.Server/_White/Wizard/Charging/ChargingSystem.cs
index 63b36002bc..e8ae12a962 100644
--- a/Content.Server/_White/Wizard/Charging/ChargingSystem.cs
+++ b/Content.Server/_White/Wizard/Charging/ChargingSystem.cs
@@ -1,5 +1,7 @@
-using Content.Shared._White.Wizard;
+using Content.Server.Actions;
+using Content.Shared._White.Wizard;
 using Content.Shared._White.Wizard.Charging;
+using Content.Shared.Actions;
 using Content.Shared.Follower;
 using Content.Shared.Mobs;
 using Robust.Shared.Audio;
@@ -12,6 +14,7 @@ public sealed class ChargingSystem : SharedChargingSystem
 {
     [Dependency] private readonly SharedAudioSystem _audio = default!;
     [Dependency] private readonly FollowerSystem _followerSystem = default!;
+    [Dependency] private readonly ActionsSystem _actionsSystem = default!;
 
     private readonly Dictionary<EntityUid, List<EntityUid>> _charges = new();
 
@@ -132,8 +135,15 @@ public sealed class ChargingSystem : SharedChargingSystem
 
     private void Add(AddWizardChargeEvent msg, EntitySessionEventArgs args)
     {
+        var spell = GetEntity(msg.Spell);
+        if (!_actionsSystem.TryGetActionData(spell, out var baseAction) ||
+            baseAction is not BaseTargetActionComponent action || !action.IsChargeEnabled)
+        {
+            return;
+        }
+
         if (args.SenderSession.AttachedEntity != null)
-            AddCharge(args.SenderSession.AttachedEntity.Value, msg.ChargeProto);
+            AddCharge(args.SenderSession.AttachedEntity.Value, action.ChargeProto);
     }
 
     private void Remove(RemoveWizardChargeEvent msg, EntitySessionEventArgs args)
diff --git a/Content.Shared/_White/Wizard/WizardEvents.cs b/Content.Shared/_White/Wizard/WizardEvents.cs
index 1627b7ec42..a3c6496aa3 100644
--- a/Content.Shared/_White/Wizard/WizardEvents.cs
+++ b/Content.Shared/_White/Wizard/WizardEvents.cs
@@ -31,11 +31,11 @@ public struct BeforeCastSpellEvent
 [Serializable, NetSerializable]
 public sealed partial class AddWizardChargeEvent : EntityEventArgs
 {
-    public string ChargeProto;
+    public NetEntity Spell;
 
-    public AddWizardChargeEvent(string chargeProto)
+    public AddWizardChargeEvent(NetEntity spell)
     {
-        ChargeProto = chargeProto;
+        Spell = spell;
     }
 }