From 68b840bfb8b1e363cb679f87b4b0f8cf76bbff29 Mon Sep 17 00:00:00 2001
From: Derek S <44935661+Derek-R-S@users.noreply.github.com>
Date: Sat, 10 Apr 2021 00:19:05 -0500
Subject: [PATCH] QoL improvements unity side

---
 .../LRM/Program/ProgramExtra.cs                |  2 +-
 .../LRM/RelayHandler/RelayHandlerCallbacks.cs  |  7 ++++++-
 .../Transport/LRM/Editor/LRMInspector.cs       | 18 +++++++++++++++++-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/Program/ProgramExtra.cs b/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/Program/ProgramExtra.cs
index b75784d..347a850 100644
--- a/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/Program/ProgramExtra.cs
+++ b/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/Program/ProgramExtra.cs
@@ -10,7 +10,7 @@ namespace LightReflectiveMirror
         public TimeSpan GetUptime() => DateTime.Now - _startupTime;
         public int GetPublicRoomCount() => _relay.rooms.Where(x => x.isPublic).Count();
 
-        static void WriteLogMessage(string message, ConsoleColor color = ConsoleColor.White, bool oneLine = false)
+        public static void WriteLogMessage(string message, ConsoleColor color = ConsoleColor.White, bool oneLine = false)
         {
             Console.ForegroundColor = color;
             if (oneLine)
diff --git a/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/RelayHandler/RelayHandlerCallbacks.cs b/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/RelayHandler/RelayHandlerCallbacks.cs
index 99c6b2d..e8a6ef3 100644
--- a/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/RelayHandler/RelayHandlerCallbacks.cs
+++ b/ServerProject-DONT-IMPORT-INTO-UNITY/LRM/RelayHandler/RelayHandlerCallbacks.cs
@@ -46,6 +46,11 @@ namespace LightReflectiveMirror
                             sendBuffer.WriteByte(ref writePos, (byte)OpCodes.Authenticated);
                             Program.transport.ServerSend(clientId, 0, new ArraySegment<byte>(sendBuffer, 0, writePos));
                         }
+                        else
+                        {
+                            Program.WriteLogMessage($"Client {clientId} sent wrong auth key! Removing from LRM node.");
+                            Program.transport.ServerDisconnect(clientId);
+                        }
                     }
                     return;
                 }
@@ -73,7 +78,7 @@ namespace LightReflectiveMirror
                     case OpCodes.UpdateRoomData:
                         var plyRoom = _cachedClientRooms[clientId];
 
-                        if (plyRoom == null)
+                        if (plyRoom == null || plyRoom.hostId != clientId)
                             return;
 
                         bool newName = data.ReadBool(ref pos);
diff --git a/UnityProject/Assets/Mirror/Runtime/Transport/LRM/Editor/LRMInspector.cs b/UnityProject/Assets/Mirror/Runtime/Transport/LRM/Editor/LRMInspector.cs
index 33009a9..bc48329 100644
--- a/UnityProject/Assets/Mirror/Runtime/Transport/LRM/Editor/LRMInspector.cs
+++ b/UnityProject/Assets/Mirror/Runtime/Transport/LRM/Editor/LRMInspector.cs
@@ -192,6 +192,20 @@ namespace LightReflectiveMirror
             {
                 // They have completed the "setup guide" Show them the main UI
 
+                // Remove unused transports...
+                foreach(var transport in lrm.GetComponentsInChildren<Transport>())
+                {
+                    if(!(transport is LightReflectiveMirrorTransport))
+                    {
+                        if(transport != lrm.clientToServerTransport && (directModule == null ? true : directModule.directConnectTransport != transport))
+                        {
+                            if (transport.gameObject == lrm.gameObject)
+                                DestroyImmediate(transport);
+                            else
+                                DestroyImmediate(transport.gameObject);
+                        }
+                    }
+                }
 #if !NET_4_6
                 EditorGUILayout.BeginVertical("Box");
                 EditorGUILayout.HelpBox("For LRM to function properly, it needs the API Compatibility to be at 4.x", MessageType.Error);
@@ -231,7 +245,9 @@ namespace LightReflectiveMirror
                         // NAT punch tab.
                         if (directModule == null)
                         {
-                            EditorGUILayout.HelpBox("If you wish to use NAT punch, you will need to add a \"Direct Connect Module\" to this gameobject.", MessageType.Info);
+                            EditorGUILayout.HelpBox("NAT Punchthrough disabled, missing Direct Connect.", MessageType.Info);
+                            if (GUILayout.Button("Add Direct Connect"))
+                                lrm.gameObject.AddComponent<LRMDirectConnectModule>();
                         }
                         else
                         {