From a8d787d811e3eaa5785f4e63349c7d1d05e4866c Mon Sep 17 00:00:00 2001
From: Oliver Booth <me@olivr.me>
Date: Wed, 18 May 2022 16:45:07 +0100
Subject: [PATCH] Support multiple packet handlers

---
 TcpDotNet/BaseClientNode.cs          |  4 ++--
 TcpDotNet/Node.cs                    | 25 ++++++++++++++++++++-----
 TcpDotNet/ProtocolListener.Client.cs |  3 ++-
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/TcpDotNet/BaseClientNode.cs b/TcpDotNet/BaseClientNode.cs
index 2fbb723..3bb6801 100644
--- a/TcpDotNet/BaseClientNode.cs
+++ b/TcpDotNet/BaseClientNode.cs
@@ -96,8 +96,8 @@ public abstract class BaseClientNode : Node
         await packet.DeserializeAsync(bufferReader);
         await targetStream.DisposeAsync();
 
-        if (RegisteredPacketHandlers.TryGetValue(packetType, out PacketHandler? packetHandler))
-            await packetHandler.HandleAsync(this, packet);
+        if (RegisteredPacketHandlers.TryGetValue(packetType, out IReadOnlyCollection<PacketHandler>? handlers))
+            await Task.WhenAll(handlers.Select(h => h.HandleAsync(this, packet)));
 
         return packet;
     }
diff --git a/TcpDotNet/Node.cs b/TcpDotNet/Node.cs
index 436123a..9e4d25f 100644
--- a/TcpDotNet/Node.cs
+++ b/TcpDotNet/Node.cs
@@ -12,7 +12,7 @@ namespace TcpDotNet;
 public abstract class Node : IDisposable
 {
     private readonly ConcurrentDictionary<int, Type> _registeredPackets = new();
-    private readonly ConcurrentDictionary<Type, PacketHandler> _registeredPacketHandlers = new();
+    private readonly ConcurrentDictionary<Type, List<PacketHandler>> _registeredPacketHandlers = new();
 
     /// <summary>
     ///     Gets the underlying socket for this node.
@@ -31,8 +31,9 @@ public abstract class Node : IDisposable
     ///     Gets the registered packets for this node.
     /// </summary>
     /// <value>The registered packets.</value>
-    public IReadOnlyDictionary<Type, PacketHandler> RegisteredPacketHandlers =>
-        new ReadOnlyDictionary<Type, PacketHandler>(_registeredPacketHandlers);
+    public IReadOnlyDictionary<Type, IReadOnlyCollection<PacketHandler>> RegisteredPacketHandlers =>
+        new ReadOnlyDictionary<Type, IReadOnlyCollection<PacketHandler>>(
+            _registeredPacketHandlers.ToDictionary(p => p.Key, p => (IReadOnlyCollection<PacketHandler>) p.Value.AsReadOnly()));
 
     /// <inheritdoc />
     public void Dispose()
@@ -54,7 +55,14 @@ public abstract class Node : IDisposable
         if (packetType is null) throw new ArgumentNullException(nameof(packetType));
         if (handler is null) throw new ArgumentNullException(nameof(handler));
         RegisterPacket(packetType);
-        _registeredPacketHandlers.TryAdd(packetType, handler);
+
+        if (!_registeredPacketHandlers.TryGetValue(packetType, out List<PacketHandler>? handlers))
+        {
+            handlers = new List<PacketHandler>();
+            _registeredPacketHandlers.TryAdd(packetType, handlers);
+        }
+
+        handlers.Add(handler);
     }
 
     /// <summary>
@@ -69,7 +77,14 @@ public abstract class Node : IDisposable
     {
         if (handler is null) throw new ArgumentNullException(nameof(handler));
         RegisterPacket<TPacket>();
-        _registeredPacketHandlers.TryAdd(typeof(TPacket), handler);
+
+        if (!_registeredPacketHandlers.TryGetValue(typeof(TPacket), out List<PacketHandler>? handlers))
+        {
+            handlers = new List<PacketHandler>();
+            _registeredPacketHandlers.TryAdd(typeof(TPacket), handlers);
+        }
+
+        handlers.Add(handler);
     }
 
     /// <summary>
diff --git a/TcpDotNet/ProtocolListener.Client.cs b/TcpDotNet/ProtocolListener.Client.cs
index f8d09ea..2ea3bf1 100644
--- a/TcpDotNet/ProtocolListener.Client.cs
+++ b/TcpDotNet/ProtocolListener.Client.cs
@@ -32,7 +32,8 @@ public sealed partial class ProtocolListener
             foreach (Type packetType in ParentListener.RegisteredPackets.Values)
                 RegisterPacket(packetType);
 
-            foreach ((Type packetType, PacketHandler handler) in ParentListener.RegisteredPacketHandlers)
+            foreach ((Type packetType, IReadOnlyCollection<PacketHandler>? handlers) in ParentListener.RegisteredPacketHandlers)
+            foreach (PacketHandler handler in handlers)
                 RegisterPacketHandler(packetType, handler);
 
             Task.Run(ReadLoopAsync);