From 9d2a76e2d4cb251c434171b5552eafd3b3f31b92 Mon Sep 17 00:00:00 2001
From: wrexbe <81056464+wrexbe@users.noreply.github.com>
Date: Sat, 27 Aug 2022 22:17:30 -0700
Subject: [PATCH] Fix Client EuiManager (#10898)

* Fix EuiManager cleanup

* cleanup

* This seems to work better
---
 Content.Client/Eui/EuiManager.cs              | 45 +++++++++++--------
 .../Tests/Cleanup/EuiManagerTest.cs           | 34 ++++++++++++++
 2 files changed, 61 insertions(+), 18 deletions(-)
 create mode 100644 Content.IntegrationTests/Tests/Cleanup/EuiManagerTest.cs

diff --git a/Content.Client/Eui/EuiManager.cs b/Content.Client/Eui/EuiManager.cs
index 43e57c8bd3..001493c537 100644
--- a/Content.Client/Eui/EuiManager.cs
+++ b/Content.Client/Eui/EuiManager.cs
@@ -1,9 +1,12 @@
 using System;
 using System.Collections.Generic;
 using Content.Shared.Eui;
+using Robust.Client.GameStates;
+using Robust.Client.State;
 using Robust.Shared.IoC;
 using Robust.Shared.Network;
 using Robust.Shared.Reflection;
+using Robust.Shared.Utility;
 
 namespace Content.Client.Eui
 {
@@ -20,6 +23,16 @@ namespace Content.Client.Eui
             _net.RegisterNetMessage<MsgEuiCtl>(RxMsgCtl);
             _net.RegisterNetMessage<MsgEuiState>(RxMsgState);
             _net.RegisterNetMessage<MsgEuiMessage>(RxMsgMessage);
+            _net.Disconnect += NetOnDisconnect;
+        }
+
+        private void NetOnDisconnect(object? sender, NetDisconnectedArgs e)
+        {
+            foreach (var openUi in _openUis)
+            {
+                openUi.Value.Eui.Closed();
+            }
+            _openUis.Clear();
         }
 
         private void RxMsgMessage(MsgEuiMessage message)
@@ -36,26 +49,22 @@ namespace Content.Client.Eui
 
         private void RxMsgCtl(MsgEuiCtl message)
         {
-            switch (message.Type)
+            // Will always close the window first when getting a control message
+            if (_openUis.TryGetValue(message.Id, out var openEui))
             {
-                case MsgEuiCtl.CtlType.Open:
-                    var euiType = _refl.LooseGetType(message.OpenType);
-                    var instance = _dtf.CreateInstance<BaseEui>(euiType);
-                    instance.Initialize(this, message.Id);
-                    instance.Opened();
-                    _openUis.Add(message.Id, new EuiData(instance));
-                    break;
-
-                case MsgEuiCtl.CtlType.Close:
-                    var dat = _openUis[message.Id];
-                    dat.Eui.Closed();
-                    _openUis.Remove(message.Id);
-
-                    break;
-
-                default:
-                    throw new ArgumentOutOfRangeException();
+                openEui.Eui.Closed();
+                _openUis.Remove(message.Id);
             }
+
+            if (message.Type != MsgEuiCtl.CtlType.Open)
+                return;
+
+            // Will open/re-open the window if the server wants the eui opened.
+            var euiType = _refl.LooseGetType(message.OpenType);
+            var instance = _dtf.CreateInstance<BaseEui>(euiType);
+            instance.Initialize(this, message.Id);
+            instance.Opened();
+            _openUis.Add(message.Id, new EuiData(instance));
         }
 
         private sealed class EuiData
diff --git a/Content.IntegrationTests/Tests/Cleanup/EuiManagerTest.cs b/Content.IntegrationTests/Tests/Cleanup/EuiManagerTest.cs
new file mode 100644
index 0000000000..e3c392450e
--- /dev/null
+++ b/Content.IntegrationTests/Tests/Cleanup/EuiManagerTest.cs
@@ -0,0 +1,34 @@
+using System.Linq;
+using System.Threading.Tasks;
+using Content.Server.Administration.UI;
+using Content.Server.EUI;
+using NUnit.Framework;
+using Robust.Server.Player;
+using Robust.Shared.IoC;
+
+namespace Content.IntegrationTests.Tests.Cleanup;
+
+public sealed class EuiManagerTest
+{
+    [Test]
+    public async Task EuiManagerRecycleWithOpenWindowTest()
+    {
+        // Even though we are using the server EUI here, we actually want to see if the client EUIManager crashes
+        for (int i = 0; i < 2; i++)
+        {
+            await using var pairTracker = await PoolManager.GetServerClient(new PoolSettings{Dirty = true});
+            var server = pairTracker.Pair.Server;
+
+            var sPlayerManager = server.ResolveDependency<IPlayerManager>();
+
+            await server.WaitAssertion(async () =>
+            {
+                var clientSession = sPlayerManager.ServerSessions.Single();
+                var eui = IoCManager.Resolve<EuiManager>();
+                var ui = new AdminAnnounceEui();
+                eui.OpenEui(ui, clientSession);
+            });
+            await pairTracker.CleanReturnAsync();
+        }
+    }
+}