From 5fee3f5ab1e42df49a6d8ba2c7f10fae0a7aba2e Mon Sep 17 00:00:00 2001
From: Julien Lemaire <julien.lemaire@sticmac.fr>
Date: Fri, 8 Jul 2022 23:39:33 +0200
Subject: [PATCH] Several callbacks can be added to a same message type and
 disposed (#197)

---
 .../Runtime/Scripts/Room/ColyseusRoom.cs      | 81 +++++++++++++------
 1 file changed, 56 insertions(+), 25 deletions(-)

diff --git a/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs b/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs
index 2aae26d8..6601a7c8 100644
--- a/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs
+++ b/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs
@@ -282,18 +282,53 @@ public async Task Send(string type, object message)
             await colyseusConnection.Send(bytes);
         }
 
+        /// <summary>
+        /// Disposable class to return with OnMessage function, so that a message callback can be removed.
+        /// </summary>
+        /// <typeparam name="MessageType">The type of object the message responds with</typeparam>
+        private class MessageRemover<MessageType> : IDisposable
+        {
+            private string _type;
+            private Action<MessageType> _targetHandler;
+            private Dictionary<string, IColyseusMessageHandler> _onMessageHandlers;
+
+            public MessageRemover(string type,
+                Action<MessageType> targetHandler,
+                Dictionary<string, IColyseusMessageHandler> onMessageHandlers)
+            {
+                _type = type;
+                _targetHandler = targetHandler;
+                _onMessageHandlers = onMessageHandlers;
+            }
+
+            public void Dispose()
+            {
+                var messageHandler = _onMessageHandlers[_type] as ColyseusMessageHandler<MessageType>;
+                messageHandler.Action -= _targetHandler;
+                if (messageHandler.Action == null) {
+                    _onMessageHandlers.Remove(_type);
+                }
+            }
+        }
+
         /// <summary>
         ///     Method to add new message handlers to the room
         /// </summary>
         /// <param name="type">The type of message received</param>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(string type, Action<MessageType> handler)
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(string type, Action<MessageType> handler)
         {
-            OnMessageHandlers.Add(type, new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            if (OnMessageHandlers.ContainsKey(type)) {
+                (OnMessageHandlers[type] as ColyseusMessageHandler<MessageType>).Action += handler;
+            } else {
+                OnMessageHandlers.Add(type, new ColyseusMessageHandler<MessageType>
+                {
+                    Action = handler
+                });
+            }
+            return new MessageRemover<MessageType>(type, handler, OnMessageHandlers);
         }
 
         /// <summary>
@@ -302,12 +337,10 @@ public void OnMessage<MessageType>(string type, Action<MessageType> handler)
         /// <param name="type">The type of message received</param>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(byte type, Action<MessageType> handler)
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(byte type, Action<MessageType> handler)
         {
-            OnMessageHandlers.Add("i" + type, new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            return this.OnMessage<MessageType>("i" + type, handler);
         }
 
         /// <summary>
@@ -315,12 +348,10 @@ public void OnMessage<MessageType>(byte type, Action<MessageType> handler)
         /// </summary>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(Action<MessageType> handler) where MessageType : Schema.Schema, new()
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(Action<MessageType> handler) where MessageType : Schema.Schema, new()
         {
-            OnMessageHandlers.Add("s" + typeof(MessageType), new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            return this.OnMessage<MessageType>("s" + typeof(MessageType), handler);
         }
 
         /// <summary>
@@ -372,15 +403,15 @@ protected async void ParseMessage(byte[] bytes)
 
                 if (bytes.Length > offset)
                 {
-	                try {
-		                serializer.Handshake(bytes, offset);
-	                }
-	                catch (Exception e)
-	                {
-		                await Leave(false);
-		                OnError?.Invoke(ColyseusErrorCode.SCHEMA_MISMATCH, e.Message);
-		                return;
-	                }
+                    try {
+                        serializer.Handshake(bytes, offset);
+                    }
+                    catch (Exception e)
+                    {
+                        await Leave(false);
+                        OnError?.Invoke(ColyseusErrorCode.SCHEMA_MISMATCH, e.Message);
+                        return;
+                    }
                 }
 
                 OnJoin?.Invoke();
@@ -423,7 +454,7 @@ protected async void ParseMessage(byte[] bytes)
             }
             else if (code == ColyseusProtocol.ROOM_STATE)
             {
-	            SetState(bytes, 1);
+                SetState(bytes, 1);
             }
             else if (code == ColyseusProtocol.ROOM_STATE_PATCH)
             {