From ab6edecdf7f86f21afd2a885b92890c77770131f Mon Sep 17 00:00:00 2001
From: Leon Friedrich <60421075+ElectroJr@users.noreply.github.com>
Date: Tue, 23 May 2023 09:57:30 +1200
Subject: [PATCH] Fix device links saving deleted entities. (#16675)

---
 .../DeviceNetwork/Systems/DeviceListSystem.cs | 47 +++++++++++++++++++
 .../DeviceLinking/SharedDeviceLinkSystem.cs   | 41 ++++++++++++++--
 2 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/Content.Server/DeviceNetwork/Systems/DeviceListSystem.cs b/Content.Server/DeviceNetwork/Systems/DeviceListSystem.cs
index e91cce1e99..edc65764da 100644
--- a/Content.Server/DeviceNetwork/Systems/DeviceListSystem.cs
+++ b/Content.Server/DeviceNetwork/Systems/DeviceListSystem.cs
@@ -4,18 +4,65 @@ using Content.Shared.DeviceNetwork;
 using Content.Shared.DeviceNetwork.Components;
 using Content.Shared.Interaction;
 using JetBrains.Annotations;
+using Robust.Shared.Map.Events;
 
 namespace Content.Server.DeviceNetwork.Systems;
 
 [UsedImplicitly]
 public sealed class DeviceListSystem : SharedDeviceListSystem
 {
+    private ISawmill _sawmill = default!;
+
     public override void Initialize()
     {
         base.Initialize();
         SubscribeLocalEvent<DeviceListComponent, ComponentInit>(OnInit);
         SubscribeLocalEvent<DeviceListComponent, BeforeBroadcastAttemptEvent>(OnBeforeBroadcast);
         SubscribeLocalEvent<DeviceListComponent, BeforePacketSentEvent>(OnBeforePacketSent);
+        SubscribeLocalEvent<BeforeSaveEvent>(OnMapSave);
+        _sawmill = Logger.GetSawmill("devicelist");
+    }
+
+    private void OnMapSave(BeforeSaveEvent ev)
+    {
+        List<EntityUid> toRemove = new();
+        var query = GetEntityQuery<TransformComponent>();
+        var enumerator = AllEntityQuery<DeviceListComponent, TransformComponent>();
+        while (enumerator.MoveNext(out var uid, out var device, out var xform))
+        {
+            if (xform.MapUid != ev.Map)
+                continue;
+
+            foreach (var ent in device.Devices)
+            {
+                if (!query.TryGetComponent(ent, out var linkedXform))
+                {
+                    // Entity was deleted.
+                    // TODO remove these on deletion instead of on-save.
+                    toRemove.Add(ent);
+                    continue;
+                }
+
+                if (linkedXform.MapUid == ev.Map)
+                    continue;
+
+                toRemove.Add(ent);
+                // TODO full game saves.
+                // when full saves are supported, this should instead add data to the BeforeSaveEvent informing the
+                // saving system that this map (or null-space entity) also needs to be included in the save.
+                _sawmill.Error(
+                    $"Saving a device list ({ToPrettyString(uid)}) that has a reference to an entity on another map ({ToPrettyString(ent)}). Removing entity from list.");
+            }
+
+            if (toRemove.Count == 0)
+                continue;
+
+            var old = device.Devices.ToList();
+            device.Devices.ExceptWith(toRemove);
+            RaiseLocalEvent(uid, new DeviceListUpdateEvent(old, device.Devices.ToList()));
+            Dirty(device);
+            toRemove.Clear();
+        }
     }
 
     public void OnInit(EntityUid uid, DeviceListComponent component, ComponentInit args)
diff --git a/Content.Shared/DeviceLinking/SharedDeviceLinkSystem.cs b/Content.Shared/DeviceLinking/SharedDeviceLinkSystem.cs
index 648596ffc1..79de58dcc2 100644
--- a/Content.Shared/DeviceLinking/SharedDeviceLinkSystem.cs
+++ b/Content.Shared/DeviceLinking/SharedDeviceLinkSystem.cs
@@ -9,6 +9,7 @@ public abstract class SharedDeviceLinkSystem : EntitySystem
 {
     [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
     [Dependency] private readonly SharedPopupSystem _popupSystem = default!;
+    private ISawmill _sawmill = default!;
 
     public const string InvokedPort = "link_port";
 
@@ -19,6 +20,7 @@ public abstract class SharedDeviceLinkSystem : EntitySystem
         SubscribeLocalEvent<DeviceLinkSinkComponent, ComponentStartup>(OnSinkStartup);
         SubscribeLocalEvent<DeviceLinkSourceComponent, ComponentRemove>(OnSourceRemoved);
         SubscribeLocalEvent<DeviceLinkSinkComponent, ComponentRemove>(OnSinkRemoved);
+        _sawmill = Logger.GetSawmill("devicelink");
     }
 
     #region Link Validation
@@ -102,9 +104,11 @@ public abstract class SharedDeviceLinkSystem : EntitySystem
     /// </summary>
     private void OnSourceRemoved(EntityUid uid, DeviceLinkSourceComponent component, ComponentRemove args)
     {
+        var query = GetEntityQuery<DeviceLinkSinkComponent>();
         foreach (var sinkUid in component.LinkedPorts.Keys)
         {
-            RemoveSinkFromSource(uid, sinkUid, component);
+            if (query.TryGetComponent(sinkUid, out var sink))
+                RemoveSinkFromSourceInternal(uid, sinkUid, component, sink);
         }
     }
 
@@ -113,9 +117,11 @@ public abstract class SharedDeviceLinkSystem : EntitySystem
     /// </summary>
     private void OnSinkRemoved(EntityUid sinkUid, DeviceLinkSinkComponent sinkComponent, ComponentRemove args)
     {
+        var query = GetEntityQuery<DeviceLinkSourceComponent>();
         foreach (var linkedSource in sinkComponent.LinkedSources)
         {
-            RemoveSinkFromSource(linkedSource, sinkUid, null, sinkComponent);
+            if (query.TryGetComponent(sinkUid, out var source))
+                RemoveSinkFromSourceInternal(linkedSource, sinkUid, source, sinkComponent);
         }
     }
 
@@ -318,8 +324,37 @@ public abstract class SharedDeviceLinkSystem : EntitySystem
         DeviceLinkSourceComponent? sourceComponent = null,
         DeviceLinkSinkComponent? sinkComponent = null)
     {
-        if (!Resolve(sourceUid, ref sourceComponent, false) || !Resolve(sinkUid, ref sinkComponent, false))
+        if (Resolve(sourceUid, ref sourceComponent, false) && Resolve(sinkUid, ref sinkComponent, false))
+        {
+            RemoveSinkFromSourceInternal(sourceUid, sinkUid, sourceComponent, sinkComponent);
             return;
+        }
+
+        if (sourceComponent == null && sinkComponent == null)
+        {
+            // Both were delted?
+            return;
+        }
+
+        if (sourceComponent == null)
+        {
+            _sawmill.Error($"Attempted to remove link between {ToPrettyString(sourceUid)} and {ToPrettyString(sinkUid)}, but the source component was missing.");
+            sinkComponent!.LinkedSources.Remove(sourceUid);
+        }
+        else
+        {
+            _sawmill.Error($"Attempted to remove link between {ToPrettyString(sourceUid)} and {ToPrettyString(sinkUid)}, but the sink component was missing.");
+            sourceComponent.LinkedPorts.Remove(sourceUid);
+        }
+    }
+
+    private void RemoveSinkFromSourceInternal(
+        EntityUid sourceUid,
+        EntityUid sinkUid,
+        DeviceLinkSourceComponent sourceComponent,
+        DeviceLinkSinkComponent sinkComponent)
+    {
+        // This function gets called on component removal. Beware that TryComp & Resolve may return false.
 
         if (sourceComponent.LinkedPorts.TryGetValue(sinkUid, out var ports))
         {