[HOTFIX] Prevent Payload Enumeration Failure. (#41376)

* Prevent payload enumeration failure when spawning entities

* hate it here

* boop

---------

Co-authored-by: Princess Cheeseballs <66055347+Pronana@users.noreply.github.com>

diff --git a/Content.Server/Payload/EntitySystems/PayloadSystem.cs b/Content.Server/Payload/EntitySystems/PayloadSystem.cs
index 542c3559d525ff55c426f05e3edbae33c35222d5..aed72a1b2869ca174884ee0d63272dcb54d620d8 100644 (file)
--- a/Content.Server/Payload/EntitySystems/PayloadSystem.cs
+++ b/Content.Server/Payload/EntitySystems/PayloadSystem.cs
@@ -17,50 +17,49 @@ namespace Content.Server.Payload.EntitySystems;
 
 public sealed class PayloadSystem : EntitySystem
 {
-    [Dependency] private readonly TagSystem _tagSystem = default!;
-    [Dependency] private readonly SharedSolutionContainerSystem _solutionContainerSystem = default!;
-    [Dependency] private readonly TransformSystem _transform = default!;
     [Dependency] private readonly IAdminLogManager _adminLogger = default!;
     [Dependency] private readonly ISerializationManager _serializationManager = default!;
+    [Dependency] private readonly SharedContainerSystem _container = default!;
+    [Dependency] private readonly SharedSolutionContainerSystem _solutionContainerSystem = default!;
+    [Dependency] private readonly TagSystem _tagSystem = default!;
+    [Dependency] private readonly TransformSystem _transform = default!;
 
     private static readonly ProtoId<TagPrototype> PayloadTag = "Payload";
 
+    // TODO: Construction System Integration tests and remove the EnsureContainer from ConstructionSystem. :(
+    private static readonly string PayloadContainer = "payload";
+    private static readonly string TriggerContainer = "payloadTrigger";
+
     public override void Initialize()
     {
         base.Initialize();
 
         SubscribeLocalEvent<PayloadCaseComponent, TriggerEvent>(OnCaseTriggered);
         SubscribeLocalEvent<PayloadTriggerComponent, TriggerEvent>(OnTriggerTriggered);
+        SubscribeLocalEvent<PayloadCaseComponent, ContainerIsInsertingAttemptEvent>(OnInsertAttempt);
         SubscribeLocalEvent<PayloadCaseComponent, EntInsertedIntoContainerMessage>(OnEntityInserted);
         SubscribeLocalEvent<PayloadCaseComponent, EntRemovedFromContainerMessage>(OnEntityRemoved);
         SubscribeLocalEvent<PayloadCaseComponent, ExaminedEvent>(OnExamined);
         SubscribeLocalEvent<ChemicalPayloadComponent, TriggerEvent>(HandleChemicalPayloadTrigger);
     }
 
-    public IEnumerable<EntityUid> GetAllPayloads(EntityUid uid, ContainerManagerComponent? contMan = null)
+    public IEnumerable<EntityUid> GetAllPayloads(EntityUid uid)
     {
-        if (!Resolve(uid, ref contMan, false))
+        if (!_container.TryGetContainer(uid, PayloadContainer, out var container))
             yield break;
 
-        foreach (var container in contMan.Containers.Values)
+        foreach (var entity in container.ContainedEntities)
         {
-            foreach (var entity in container.ContainedEntities)
-            {
-                if (_tagSystem.HasTag(entity, PayloadTag))
-                    yield return entity;
-            }
+            if (_tagSystem.HasTag(entity, PayloadTag))
+                yield return entity;
         }
     }
 
     private void OnCaseTriggered(EntityUid uid, PayloadCaseComponent component, TriggerEvent args)
     {
         // TODO: Adjust to the new trigger system
-
-        if (!TryComp(uid, out ContainerManagerComponent? contMan))
-            return;
-
         // Pass trigger event onto all contained payloads. Payload capacity configurable by construction graphs.
-        foreach (var ent in GetAllPayloads(uid, contMan))
+        foreach (var ent in GetAllPayloads(uid))
         {
             RaiseLocalEvent(ent, ref args, false);
         }
@@ -82,9 +81,18 @@ public sealed class PayloadSystem : EntitySystem
         RaiseLocalEvent(parent, ref args);
     }
 
+    private void OnInsertAttempt(Entity<PayloadCaseComponent> ent, ref ContainerIsInsertingAttemptEvent args)
+    {
+        if (args.Container.ID == PayloadContainer && !_tagSystem.HasTag(args.EntityUid, PayloadTag))
+            args.Cancel();
+
+        if (args.Container.ID == TriggerContainer && !HasComp<PayloadTriggerComponent>(args.EntityUid))
+            args.Cancel();
+    }
+
     private void OnEntityInserted(EntityUid uid, PayloadCaseComponent _, EntInsertedIntoContainerMessage args)
     {
-        if (!TryComp(args.Entity, out PayloadTriggerComponent? trigger))
+        if (args.Container.ID != TriggerContainer || !TryComp(args.Entity, out PayloadTriggerComponent? trigger))
             return;
 
         trigger.Active = true;
@@ -114,7 +122,7 @@ public sealed class PayloadSystem : EntitySystem
 
     private void OnEntityRemoved(EntityUid uid, PayloadCaseComponent component, EntRemovedFromContainerMessage args)
     {
-        if (!TryComp(args.Entity, out PayloadTriggerComponent? trigger))
+        if (args.Container.ID != TriggerContainer || !TryComp(args.Entity, out PayloadTriggerComponent? trigger))
             return;
 
         trigger.Active = false;