From 2a289484188187d64b61e1388757e2d193ea57f4 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Fri, 27 Jun 2025 22:55:00 +0800
Subject: [PATCH] :sparkles: Chat room subscribe

---
 DysonNetwork.Sphere/Chat/ChatService.cs       |  1 +
 .../Handlers/MessagesSubscribeHandler.cs      | 53 +++++++++++++++++++
 .../Handlers/MessagesUnsubscribeHandler.cs    | 21 ++++++++
 .../Connection/WebSocketService.cs            | 26 +++++++++
 .../Startup/ServiceCollectionExtensions.cs    |  2 +
 5 files changed, 103 insertions(+)
 create mode 100644 DysonNetwork.Sphere/Connection/Handlers/MessagesSubscribeHandler.cs
 create mode 100644 DysonNetwork.Sphere/Connection/Handlers/MessagesUnsubscribeHandler.cs

diff --git a/DysonNetwork.Sphere/Chat/ChatService.cs b/DysonNetwork.Sphere/Chat/ChatService.cs
index dc675aa..8e6daa3 100644
--- a/DysonNetwork.Sphere/Chat/ChatService.cs
+++ b/DysonNetwork.Sphere/Chat/ChatService.cs
@@ -252,6 +252,7 @@ public partial class ChatService(
 
             if (member.Account.Id == sender.AccountId) continue;
             if (member.Notify == ChatMemberNotify.None) continue;
+            if (scopedWs.IsUserSubscribedToChatRoom(member.AccountId, room.Id.ToString())) continue;
             if (message.MembersMentioned is null || !message.MembersMentioned.Contains(member.Account.Id))
             {
                 var now = SystemClock.Instance.GetCurrentInstant();
diff --git a/DysonNetwork.Sphere/Connection/Handlers/MessagesSubscribeHandler.cs b/DysonNetwork.Sphere/Connection/Handlers/MessagesSubscribeHandler.cs
new file mode 100644
index 0000000..549e75f
--- /dev/null
+++ b/DysonNetwork.Sphere/Connection/Handlers/MessagesSubscribeHandler.cs
@@ -0,0 +1,53 @@
+
+using System.Net.WebSockets;
+using DysonNetwork.Sphere.Chat;
+
+namespace DysonNetwork.Sphere.Connection.Handlers;
+
+public class MessagesSubscribeHandler(ChatRoomService crs, WebSocketService webSocketService) : IWebSocketPacketHandler
+{
+    public string PacketType => "messages.subscribe";
+
+    public async Task HandleAsync(
+        Account.Account currentUser,
+        string deviceId,
+        WebSocketPacket packet,
+        WebSocket socket,
+        WebSocketService srv
+    )
+    {
+        var request = packet.GetData<ChatController.TypingMessageRequest>();
+        if (request is null)
+        {
+            await socket.SendAsync(
+                new ArraySegment<byte>(new WebSocketPacket
+                {
+                    Type = WebSocketPacketType.Error,
+                    ErrorMessage = "messages.subscribe requires you provide the ChatRoomId"
+                }.ToBytes()),
+                WebSocketMessageType.Binary,
+                true,
+                CancellationToken.None
+            );
+            return;
+        }
+
+        var sender = await crs.GetRoomMember(currentUser.Id, request.ChatRoomId);
+        if (sender is null)
+        {
+            await socket.SendAsync(
+                new ArraySegment<byte>(new WebSocketPacket
+                {
+                    Type = WebSocketPacketType.Error,
+                    ErrorMessage = "User is not a member of the chat room."
+                }.ToBytes()),
+                WebSocketMessageType.Binary,
+                true,
+                CancellationToken.None
+            );
+            return;
+        }
+
+        webSocketService.SubscribeToChatRoom(sender.ChatRoomId.ToString(), deviceId);
+    }
+}
diff --git a/DysonNetwork.Sphere/Connection/Handlers/MessagesUnsubscribeHandler.cs b/DysonNetwork.Sphere/Connection/Handlers/MessagesUnsubscribeHandler.cs
new file mode 100644
index 0000000..4b353aa
--- /dev/null
+++ b/DysonNetwork.Sphere/Connection/Handlers/MessagesUnsubscribeHandler.cs
@@ -0,0 +1,21 @@
+using System.Net.WebSockets;
+using DysonNetwork.Sphere.Chat;
+
+namespace DysonNetwork.Sphere.Connection.Handlers;
+
+public class MessagesUnsubscribeHandler(WebSocketService webSocketService) : IWebSocketPacketHandler
+{
+    public string PacketType => "messages.unsubscribe";
+
+    public Task HandleAsync(
+        Account.Account currentUser,
+        string deviceId,
+        WebSocketPacket packet,
+        WebSocket socket,
+        WebSocketService srv
+    )
+    {
+        webSocketService.UnsubscribeFromChatRoom(deviceId);
+        return Task.CompletedTask;
+    }
+}
diff --git a/DysonNetwork.Sphere/Connection/WebSocketService.cs b/DysonNetwork.Sphere/Connection/WebSocketService.cs
index b8b70a6..c63575c 100644
--- a/DysonNetwork.Sphere/Connection/WebSocketService.cs
+++ b/DysonNetwork.Sphere/Connection/WebSocketService.cs
@@ -17,6 +17,31 @@ public class WebSocketService
         (WebSocket Socket, CancellationTokenSource Cts)
     > ActiveConnections = new();
 
+    private static readonly ConcurrentDictionary<string, string> ActiveSubscriptions = new(); // deviceId -> chatRoomId
+
+    public void SubscribeToChatRoom(string chatRoomId, string deviceId)
+    {
+        ActiveSubscriptions[deviceId] = chatRoomId;
+    }
+
+    public void UnsubscribeFromChatRoom(string deviceId)
+    {
+        ActiveSubscriptions.TryRemove(deviceId, out _);
+    }
+
+    public bool IsUserSubscribedToChatRoom(Guid accountId, string chatRoomId)
+    {
+        var userDeviceIds = ActiveConnections.Keys.Where(k => k.AccountId == accountId).Select(k => k.DeviceId);
+        foreach (var deviceId in userDeviceIds)
+        {
+            if (ActiveSubscriptions.TryGetValue(deviceId, out var subscribedChatRoomId) && subscribedChatRoomId == chatRoomId)
+            {
+                return true;
+            }
+        }
+        return false;
+    }
+
     public bool TryAdd(
         (Guid AccountId, string DeviceId) key,
         WebSocket socket,
@@ -39,6 +64,7 @@ public class WebSocketService
         );
         data.Cts.Cancel();
         ActiveConnections.TryRemove(key, out _);
+        UnsubscribeFromChatRoom(key.DeviceId);
     }
 
     public bool GetAccountIsConnected(Guid accountId)
diff --git a/DysonNetwork.Sphere/Startup/ServiceCollectionExtensions.cs b/DysonNetwork.Sphere/Startup/ServiceCollectionExtensions.cs
index d24b147..ab71e5d 100644
--- a/DysonNetwork.Sphere/Startup/ServiceCollectionExtensions.cs
+++ b/DysonNetwork.Sphere/Startup/ServiceCollectionExtensions.cs
@@ -187,6 +187,8 @@ public static class ServiceCollectionExtensions
         // The handlers for websocket
         services.AddScoped<IWebSocketPacketHandler, MessageReadHandler>();
         services.AddScoped<IWebSocketPacketHandler, MessageTypingHandler>();
+        services.AddScoped<IWebSocketPacketHandler, MessagesSubscribeHandler>();
+        services.AddScoped<IWebSocketPacketHandler, MessagesUnsubscribeHandler>();
 
         return services;
     }