radio implant refactor and fix (#2882)

Co-authored-by: deltanedas <@deltanedas:kde.org>

Content.Server/_DV/Implants/Radio/RadioImplantSystem.cs

@@ -13,20 +13,22 @@ namespace Content.Server._DV.Implants.Radio;
 /// <inheritdoc />
 public sealed class RadioImplantSystem : SharedRadioImplantSystem
 {
-    [Dependency] private readonly INetManager _netManager = default!;
-    [Dependency] private readonly RadioSystem _radioSystem = default!;
+    [Dependency] private readonly INetManager _net = default!;
+    [Dependency] private readonly RadioSystem _radio = default!;
 
-    private EntityQuery<ActorComponent> _actorQuery;
+    private EntityQuery<ActorComponent> _actor;
 
     public override void Initialize()
     {
         base.Initialize();
+
+        _actor = GetEntityQuery<ActorComponent>();
+
         SubscribeLocalEvent<RadioImplantComponent, MapInitEvent>(OnMapInit);
         SubscribeLocalEvent<RadioImplantComponent, EntInsertedIntoContainerMessage>(OnInsertEncryptionKey);
         SubscribeLocalEvent<RadioImplantComponent, EntRemovedFromContainerMessage>(OnRemoveEncryptionKey);
         SubscribeLocalEvent<RadioImplantComponent, RadioReceiveEvent>(OnRadioReceive);
         SubscribeLocalEvent<HasRadioImplantComponent, EntitySpokeEvent>(OnSpeak);
-        _actorQuery = GetEntityQuery<ActorComponent>();
     }
 
     /// <summary>
@@ -43,27 +45,29 @@ private void OnMapInit(Entity<RadioImplantComponent> ent, ref MapInitEvent args)
     private void OnSpeak(Entity<HasRadioImplantComponent> ent, ref EntitySpokeEvent args)
     {
         // not a radio message, or already handled by another radio
-        if (args.Channel is null)
+        if (args.Channel is not {} channel)
             return;
 
-        // does the implant have access to the channel the implantee is trying to speak on?
-        if (ent.Comp.Implant is {} implant
-            && TryComp<RadioImplantComponent>(implant, out var radioImplantComponent)
-            && radioImplantComponent.Channels.Contains(args.Channel.ID))
+        // does an implant have access to the channel the implantee is trying to speak on?
+        foreach (var implant in ent.Comp.Implants)
         {
-            _radioSystem.SendRadioMessage(ent, args.Message, args.Channel.ID, implant);
-            // prevent other radios they might be wearing from sending the message again
-            args.Channel = null;
+            if (TryComp<RadioImplantComponent>(implant, out var radioImplant) &&
+                radioImplant.Channels.Contains(channel.ID))
+            {
+                _radio.SendRadioMessage(ent, args.Message, channel.ID, implant);
+                // prevent other radios they might be wearing from sending the message again
+                args.Channel = null;
+            }
         }
     }
 
     /// <summary>
     /// Handles receiving radio messages and forwarding them to the implantee.
     /// </summary>
-    private void OnRadioReceive(EntityUid uid, RadioImplantComponent component, ref RadioReceiveEvent args)
+    private void OnRadioReceive(Entity<RadioImplantComponent> ent, ref RadioReceiveEvent args)
     {
-        if (_actorQuery.TryComp(component.Implantee, out var actorComponent))
-            _netManager.ServerSendMessage(args.ChatMsg, actorComponent.PlayerSession.Channel);
+        if (_actor.TryComp(ent.Comp.Implantee, out var actor))
+            _net.ServerSendMessage(args.ChatMsg, actor.PlayerSession.Channel);
     }
 
     /// <summary>
@@ -74,12 +78,12 @@ private void OnInsertEncryptionKey(Entity<RadioImplantComponent> ent, ref EntIns
         // check if the insertion is actually something getting inserted into the radio implant storage, since
         // this evt also fires when the radio implant is being inserted into a person.
         if (ent.Owner != args.Container.Owner
-            || !TryComp<EncryptionKeyComponent>(args.Entity, out var encryptionKeyComponent))
+            || !TryComp<EncryptionKeyComponent>(args.Entity, out var key))
             return;
 
         // copy over the radio channels that can be accessed
         ent.Comp.Channels.Clear();
-        foreach (var channel in encryptionKeyComponent.Channels)
+        foreach (var channel in key.Channels)
         {
             ent.Comp.Channels.Add(channel);
         }
@@ -99,6 +103,7 @@ private void OnRemoveEncryptionKey(Entity<RadioImplantComponent> ent, ref EntRem
             return;
 
         // clear the radio channels since there's no encryption key inserted anymore.
+        // if you ever make the storage have more than 1 key's space you will have to rebuild it instead
         ent.Comp.Channels.Clear();
         Dirty(ent);
         UpdateRadioReception(ent);
@@ -109,18 +114,17 @@ private void OnRemoveEncryptionKey(Entity<RadioImplantComponent> ent, ref EntRem
     /// </summary>
     private void UpdateRadioReception(Entity<RadioImplantComponent> ent)
     {
-        if (ent.Comp.Channels.Count != 0)
+        if (ent.Comp.Channels.Count == 0)
         {
-            // we need to add this comp to actually receive radio events.
-            var channels = EnsureComp<ActiveRadioComponent>(ent).Channels;
-            foreach (var channel in ent.Comp.Channels)
-            {
-                channels.Add(channel);
-            }
+            RemComp<ActiveRadioComponent>(ent);
+            return;
         }
-        else
+
+        // we need to add this comp to actually receive radio events.
+        var channels = EnsureComp<ActiveRadioComponent>(ent).Channels;
+        foreach (var channel in ent.Comp.Channels)
         {
-            RemComp<ActiveRadioComponent>(ent);
+            channels.Add(channel);
         }
     }
 }

Content.Shared/_DV/Implants/Radio/HasRadioImplantComponent.cs

@@ -5,12 +5,13 @@ namespace Content.Shared._DV.Implants.Radio;
 /// <summary>
 /// This indicates this entity has a radio implant implanted into themselves.
 /// </summary>
-[RegisterComponent, NetworkedComponent, AutoGenerateComponentState, Access(typeof(SharedRadioImplantSystem))]
+[RegisterComponent, NetworkedComponent, Access(typeof(SharedRadioImplantSystem))]
+[AutoGenerateComponentState]
 public sealed partial class HasRadioImplantComponent : Component
 {
     /// <summary>
-    /// The radio implant. We need this to be able to determine encryption keys.
+    /// A list of radio implants. We need this to be able to determine encryption keys.
     /// </summary>
     [DataField, AutoNetworkedField]
-    public EntityUid? Implant;
+    public HashSet<EntityUid> Implants = new();
 }

Content.Shared/_DV/Implants/Radio/SharedRadioImplantSystem.cs

@@ -1,7 +1,6 @@
-using Content.Shared.Actions;
 using Content.Shared.Implants;
-using Content.Shared.Storage;
-using Content.Shared.Storage.EntitySystems;
+//using Content.Shared.Storage;
+//using Content.Shared.Storage.EntitySystems;
 using Robust.Shared.Containers;
 
 namespace Content.Shared._DV.Implants.Radio;
@@ -15,42 +14,45 @@ public abstract class SharedRadioImplantSystem : EntitySystem
     public override void Initialize()
     {
         SubscribeLocalEvent<RadioImplantComponent, ImplantImplantedEvent>(OnImplanted);
-        SubscribeLocalEvent<RadioImplantComponent, EntGotRemovedFromContainerMessage>(OnPossiblyUnimplanted);
+        SubscribeLocalEvent<RadioImplantComponent, EntGotRemovedFromContainerMessage>(OnRemoved);
     }
 
     /// <summary>
     /// Handles implantation of the implant.
     /// </summary>
-    private void OnImplanted(EntityUid uid, RadioImplantComponent component, ImplantImplantedEvent args)
+    private void OnImplanted(Entity<RadioImplantComponent> ent, ref ImplantImplantedEvent args)
     {
-        if (args.Implanted is not { Valid: true })
+        if (args.Implanted is not {} user)
             return;
 
-        component.Implantee = args.Implanted.Value;
-        Dirty(uid, component);
+        ent.Comp.Implantee = user;
+        Dirty(ent, ent.Comp);
 
         // make sure the person entity gets slapped with a component so it can react to it talking.
-        var hasRadioImplantComponent = EnsureComp<HasRadioImplantComponent>(args.Implanted.Value);
-        hasRadioImplantComponent.Implant = uid;
-        Dirty(component.Implantee.Value, hasRadioImplantComponent);
+        var implanted = EnsureComp<HasRadioImplantComponent>(user);
+        implanted.Implants.Add(ent);
+        Dirty(user, implanted);
     }
 
-
     /// <summary>
     /// Handles removal of the implant from its containing mob.
     /// </summary>
     /// <remarks>Done via <see cref="EntGotRemovedFromContainerMessage"/> because there is no specific event for an implant being removed.</remarks>
-    private void OnPossiblyUnimplanted(EntityUid uid, RadioImplantComponent component, EntGotRemovedFromContainerMessage args)
+    private void OnRemoved(Entity<RadioImplantComponent> ent, ref EntGotRemovedFromContainerMessage args)
     {
-        if (Terminating(uid))
+        if (TerminatingOrDeleted(ent) ||
+            ent.Comp.Implantee is not {} user ||
+            // this gets fired if it gets removed from ANY container but really, we just want to know if it was removed from its owner...
+            // so check if the ent we got implanted into matches the container's owner (here, the container's owner is the entity)
+            user != args.Container.Owner ||
+            !TryComp<HasRadioImplantComponent>(user, out var implanted))
             return;
 
-        // this gets fired if it gets removed from ANY container but really, we just want to know if it was removed from its owner...
-        // so check if the ent we got implanted into matches the container's owner (here, the container's owner is the entity)
-        if (component.Implantee is not null && component.Implantee == args.Container.Owner)
-        {
-            RemComp<HasRadioImplantComponent>(component.Implantee.Value);
-            component.Implantee = null;
-        }
+        implanted.Implants.Remove(ent);
+        if (implanted.Implants.Count == 0)
+            RemComp<HasRadioImplantComponent>(user);
+        else
+            Dirty(user, implanted);
+        ent.Comp.Implantee = null;
     }
 }