Fix restartround crash with the DoAfter system (#1658)

Content.Client/GameObjects/EntitySystems/DoAfter/DoAfterGui.cs

@@ -20,20 +20,20 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
     {
         [Dependency] private readonly IEyeManager _eyeManager = default!;
         [Dependency] private readonly IGameTiming _gameTiming = default!;
-        
+
         private Dictionary<byte, PanelContainer> _doAfterControls = new Dictionary<byte, PanelContainer>();
         private Dictionary<byte, DoAfterBar> _doAfterBars = new Dictionary<byte, DoAfterBar>();
-        
+
         // We'll store cancellations for a little bit just so we can flash the graphic to indicate it's cancelled
         private Dictionary<byte, TimeSpan> _cancelledDoAfters = new Dictionary<byte, TimeSpan>();
 
         public IEntity? AttachedEntity { get; set; }
         private ScreenCoordinates _playerPosition;
-        
+
         // This behavior probably shouldn't be happening; so for whatever reason the control position is set the frame after
         // I got NFI why because I don't know the UI internals
         private bool _firstDraw = true;
-        
+
         public DoAfterGui()
         {
             IoCManager.InjectDependencies(this);
@@ -71,7 +71,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
                         TextureScale = Vector2.One * DoAfterBar.DoAfterBarScale,
                         SizeFlagsVertical = SizeFlags.ShrinkCenter,
                     },
-                    
+
                     doAfterBar
                 }
             };
@@ -79,9 +79,9 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             AddChild(control);
             _doAfterControls.Add(message.ID, control);
         }
-        
+
         // NOTE THAT THE BELOW ONLY HANDLES THE UI SIDE
-        
+
         /// <summary>
         ///     Removes a DoAfter without showing a cancel graphic.
         /// </summary>
@@ -92,7 +92,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             {
                 return;
             }
-            
+
             var control = _doAfterControls[id];
             RemoveChild(control);
             _doAfterControls.Remove(id);
@@ -114,7 +114,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             {
                 return;
             }
-            
+
             var control = _doAfterControls[id];
             _doAfterBars[id].Cancelled = true;
             _cancelledDoAfters.Add(id, _gameTiming.CurTime);
@@ -124,13 +124,13 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         {
             base.FrameUpdate(args);
 
-            if (AttachedEntity == null || !AttachedEntity.TryGetComponent(out DoAfterComponent doAfterComponent))
+            if (AttachedEntity?.IsValid() != true || !AttachedEntity.TryGetComponent(out DoAfterComponent doAfterComponent))
             {
                 return;
             }
-            
+
             var doAfters = doAfterComponent.DoAfters;
-            
+
             // Nothing to render so we'll hide.
             if (doAfters.Count == 0 && _cancelledDoAfters.Count == 0)
             {
@@ -138,7 +138,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
                 Visible = false;
                 return;
             }
-            
+
             // Set position ready for 2nd+ frames.
             _playerPosition = _eyeManager.WorldToScreen(AttachedEntity.Transform.GridPosition);
             LayoutContainer.SetPosition(this, new Vector2(_playerPosition.X - Width / 2, _playerPosition.Y - Height - 30.0f));
@@ -152,7 +152,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             Visible = true;
             var currentTime = _gameTiming.CurTime;
             var toCancel = new List<byte>();
-            
+
             // Cleanup cancelled DoAfters
             foreach (var (id, cancelTime) in _cancelledDoAfters)
             {
@@ -166,7 +166,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             {
                 RemoveDoAfter(id);
             }
-            
+
             // Update 0 -> 1.0f of the things
             foreach (var (id, message) in doAfters)
             {
@@ -174,11 +174,11 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
                 {
                     continue;
                 }
-                
+
                 var doAfterBar = _doAfterBars[id];
                 doAfterBar.Ratio = MathF.Min(1.0f,
                     (float) (currentTime - message.StartTime).TotalSeconds / message.Delay);
             }
         }
     }
-}
+}

Content.Client/GameObjects/EntitySystems/DoAfter/DoAfterSystem.cs

@@ -50,19 +50,20 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             base.Shutdown();
             Gui?.Dispose();
             Gui = null;
+            _player = null;
         }
-        
+
         private void HandlePlayerAttached(IEntity? entity)
         {
             _player = entity;
             // Setup the GUI and pass the new data to it if applicable.
             Gui?.Dispose();
-            
+
             if (entity == null)
             {
                 return;
             }
-            
+
             Gui ??= new DoAfterGui();
             Gui.AttachedEntity = entity;
 
@@ -81,24 +82,29 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
 
             var currentTime = _gameTiming.CurTime;
 
-            if (_player == null || !_player.TryGetComponent(out DoAfterComponent doAfterComponent))
+            if (_player?.IsValid() != true)
             {
                 return;
             }
-            
+
+            if (!_player.TryGetComponent(out DoAfterComponent doAfterComponent))
+            {
+                return;
+            }
+
             var doAfters = doAfterComponent.DoAfters.ToList();
             if (doAfters.Count == 0)
             {
                 return;
             }
-                
+
             var userGrid = _player.Transform.GridPosition;
-                
+
             // Check cancellations / finishes
             foreach (var (id, doAfter) in doAfters)
             {
                 var elapsedTime = (currentTime - doAfter.StartTime).TotalSeconds;
-                    
+
                 // If we've passed the final time (after the excess to show completion graphic) then remove.
                 if (elapsedTime > doAfter.Delay + ExcessTime)
                 {
@@ -106,7 +112,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
                     doAfterComponent.Remove(doAfter);
                     continue;
                 }
-                    
+
                 // Don't predict cancellation if it's already finished.
                 if (elapsedTime > doAfter.Delay)
                 {
@@ -147,4 +153,4 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             }
         }
     }
-}
+}