From d1a830cd74f2ba65ef71350513058e738ddf3748 Mon Sep 17 00:00:00 2001
From: Julian Giebel <juliangiebel@live.de>
Date: Mon, 25 Oct 2021 21:55:00 +0200
Subject: [PATCH] Fix ApcPowerReceiverComponent deletion (#5031)

Co-authored-by: Julian Giebel <j.giebel@netrocks.info>
---
 .../Components/ApcPowerReceiverComponent.cs   |  9 +++++++++
 .../EntitySystems/ExtensionCableSystem.cs     |  2 +-
 .../EntitySystems/PowerReceiverSystem.cs      | 20 +++++++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/Content.Server/Power/Components/ApcPowerReceiverComponent.cs b/Content.Server/Power/Components/ApcPowerReceiverComponent.cs
index 7f0cc03e1b..474c414ab7 100644
--- a/Content.Server/Power/Components/ApcPowerReceiverComponent.cs
+++ b/Content.Server/Power/Components/ApcPowerReceiverComponent.cs
@@ -35,6 +35,8 @@ namespace Content.Server.Power.Components
         [DataField("powerLoad")]
         public float Load { get => NetworkLoad.DesiredPower; set => NetworkLoad.DesiredPower = value; }
 
+        public ApcPowerProviderComponent? Provider = null;
+
         /// <summary>
         ///     When false, causes this to appear powered even if not receiving power from an Apc.
         /// </summary>
@@ -68,6 +70,13 @@ namespace Content.Server.Power.Components
             DesiredPower = 5
         };
 
+        protected override void OnRemove()
+        {
+            Provider?.RemoveReceiver(this);
+
+            base.OnRemove();
+        }
+
         public void ApcPowerChanged()
         {
             OnNewPowerState();
diff --git a/Content.Server/Power/EntitySystems/ExtensionCableSystem.cs b/Content.Server/Power/EntitySystems/ExtensionCableSystem.cs
index 83847eff56..1bbce9041f 100644
--- a/Content.Server/Power/EntitySystems/ExtensionCableSystem.cs
+++ b/Content.Server/Power/EntitySystems/ExtensionCableSystem.cs
@@ -79,7 +79,7 @@ namespace Content.Server.Power.EntitySystems
 
             foreach (var entity in nearbyEntities)
             {
-                if (entity.TryGetComponent<ExtensionCableReceiverComponent>(out var receiver) &&
+                if (EntityManager.TryGetComponent<ExtensionCableReceiverComponent>(entity.Uid, out var receiver) &&
                     receiver.Connectable &&
                     receiver.Provider == null &&
                     entity.Transform.Coordinates.TryDistance(owner.EntityManager, owner.Transform.Coordinates, out var distance) &&
diff --git a/Content.Server/Power/EntitySystems/PowerReceiverSystem.cs b/Content.Server/Power/EntitySystems/PowerReceiverSystem.cs
index de39ec242d..8a77441cf0 100644
--- a/Content.Server/Power/EntitySystems/PowerReceiverSystem.cs
+++ b/Content.Server/Power/EntitySystems/PowerReceiverSystem.cs
@@ -12,17 +12,37 @@ namespace Content.Server.Power.EntitySystems
             SubscribeLocalEvent<ApcPowerReceiverComponent, ExtensionCableSystem.ProviderConnectedEvent>(OnProviderConnected);
             SubscribeLocalEvent<ApcPowerReceiverComponent, ExtensionCableSystem.ProviderDisconnectedEvent>(OnProviderDisconnected);
 
+            SubscribeLocalEvent<ApcPowerProviderComponent, ComponentShutdown>(OnProviderShutdown);
             SubscribeLocalEvent<ApcPowerProviderComponent, ExtensionCableSystem.ReceiverConnectedEvent>(OnReceiverConnected);
             SubscribeLocalEvent<ApcPowerProviderComponent, ExtensionCableSystem.ReceiverDisconnectedEvent>(OnReceiverDisconnected);
         }
 
+        private void OnProviderShutdown(EntityUid uid, ApcPowerProviderComponent component, ComponentShutdown args)
+        {
+            foreach (var receiver in component.LinkedReceivers)
+            {
+                receiver.NetworkLoad.LinkedNetwork = default;
+                component.Net?.QueueNetworkReconnect();
+            }
+
+            component.LinkedReceivers.Clear();
+        }
+
         private void OnProviderConnected(EntityUid uid, ApcPowerReceiverComponent receiver, ExtensionCableSystem.ProviderConnectedEvent args)
         {
+            var providerUid = args.Provider.Owner.Uid;
+            if (!EntityManager.TryGetComponent<ApcPowerProviderComponent>(providerUid, out var provider))
+                return;
+
+            receiver.Provider = provider;
+
             ProviderChanged(receiver);
         }
 
         private void OnProviderDisconnected(EntityUid uid, ApcPowerReceiverComponent receiver, ExtensionCableSystem.ProviderDisconnectedEvent args)
         {
+            receiver.Provider = null;
+
             ProviderChanged(receiver);
         }