Rewritten notification system

Moonlight/App/Http/Controllers/Api/Moonlight/Notifications/ListenController.cs

@@ -1,7 +1,9 @@
 using System.Net.WebSockets;
 using System.Text;
+using Logging.Net;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.EntityFrameworkCore;
+using Moonlight.App.Database.Entities;
 using Moonlight.App.Database.Entities.Notification;
 using Moonlight.App.Models.Notifications;
 using Moonlight.App.Repositories;
@@ -12,135 +14,156 @@ using Newtonsoft.Json;
 
 namespace Moonlight.App.Http.Controllers.Api.Moonlight.Notifications;
 
-public class ListenController : ControllerBase
+[ApiController]
+[Route("api/moonlight/notification/listen")]
+public class ListenController : Controller
 {
-    internal WebSocket ws;
+    private WebSocket WebSocket;
-    private bool active = true;
-    private bool isAuth = false;
     private NotificationClient Client;
-    
+    private CancellationTokenSource CancellationTokenSource = new();
-    private readonly IdentityService IdentityService;
-    private readonly NotificationRepository NotificationRepository;
-    private readonly OneTimeJwtService OneTimeJwtService;
-    private readonly NotificationClientService NotificationClientService;
-    private readonly NotificationServerService NotificationServerService;
 
-    public ListenController(IdentityService identityService, 
+    private User? CurrentUser;
-        NotificationRepository notificationRepository, 
+
-        OneTimeJwtService oneTimeJwtService, 
+    private readonly OneTimeJwtService OneTimeJwtService;
-        NotificationClientService notificationClientService,
+    private readonly NotificationServerService NotificationServerService;
-        NotificationServerService notificationServerService)
+    private readonly Repository<NotificationClient> NotificationClientRepository;
+
+    public ListenController(
+        OneTimeJwtService oneTimeJwtService,
+        NotificationServerService notificationServerService, Repository<NotificationClient> notificationClientRepository)
     {
-        IdentityService = identityService;
-        NotificationRepository = notificationRepository;
         OneTimeJwtService = oneTimeJwtService;
-        NotificationClientService = notificationClientService;
         NotificationServerService = notificationServerService;
+        NotificationClientRepository = notificationClientRepository;
     }
-    
+
     [Route("/api/moonlight/notifications/listen")]
-    public async Task Get()
+    public async Task<ActionResult> Get()
     {
         if (HttpContext.WebSockets.IsWebSocketRequest)
         {
-            using var webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
+            WebSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
-            ws = webSocket;
+
-            await Echo();
+            await ProcessWebsocket();
+
+            return new EmptyResult();
         }
         else
         {
-            HttpContext.Response.StatusCode = StatusCodes.Status403Forbidden;
+            return StatusCode(400);
         }
     }
 
-    private async Task Echo()
+    private async Task ProcessWebsocket()
     {
-        while (active)
+        while (!CancellationTokenSource.Token.IsCancellationRequested && WebSocket.State == WebSocketState.Open)
         {
-            byte[] bytes = new byte[1024 * 16]; 
+            try
-            var asg = new ArraySegment<byte>(bytes);
-            var res = await ws.ReceiveAsync(asg, CancellationToken.None);
-
-            var text = Encoding.UTF8.GetString(bytes).Trim('\0');
-
-            var obj = JsonConvert.DeserializeObject<BasicWSModel>(text);
-
-            if (!string.IsNullOrWhiteSpace(obj.Action))
             {
-                await HandleRequest(text, obj.Action);
+                byte[] buffer = new byte[1024 * 16];
+                _ = await WebSocket.ReceiveAsync(buffer, CancellationTokenSource.Token);
+                var text = Encoding.UTF8.GetString(buffer).Trim('\0');
+
+                var basicWsModel = JsonConvert.DeserializeObject<BasicWSModel>(text) ?? new();
+
+                if (!string.IsNullOrWhiteSpace(basicWsModel.Action))
+                {
+                    await HandleRequest(text, basicWsModel.Action);
+                }
+
+                if (WebSocket.State != WebSocketState.Open)
+                {
+                    CancellationTokenSource.Cancel();
+                }
+            }
+            catch (WebSocketException e)
+            {
+                CancellationTokenSource.Cancel();
             }
-            
-            active = ws.State == WebSocketState.Open;
         }
+
+        await NotificationServerService.UnRegisterClient(Client);
     }
 
     private async Task HandleRequest(string text, string action)
     {
-        if (!isAuth && action == "login")
+        if (CurrentUser == null && action != "login")
-            await Login(text);
-        else if (!isAuth)
-            await ws.SendAsync(Encoding.UTF8.GetBytes("{\"error\": \"Unauthorised\"}"), WebSocketMessageType.Text,
-                WebSocketMessageFlags.EndOfMessage, CancellationToken.None);
-        else switch (action)
         {
+            await Send("{\"error\": \"Unauthorised\"}");
+        }
+
+        switch (action)
+        {
+            case "login":
+                await Login(text);
+                break;
             case "received":
                 await Received(text);
                 break;
             case "read":
                 await Read(text);
                 break;
-            default:
-                break;
         }
     }
 
+    private async Task Send(string text)
+    {
+        await WebSocket.SendAsync(
+            Encoding.UTF8.GetBytes(text),
+            WebSocketMessageType.Text,
+            WebSocketMessageFlags.EndOfMessage, CancellationTokenSource.Token
+        );
+    }
+
     private async Task Login(string json)
     {
-        var jwt = JsonConvert.DeserializeObject<Login>(json).token;
+        var loginModel = JsonConvert.DeserializeObject<Login>(json) ?? new();
-        
+
-        var dict = await OneTimeJwtService.Validate(jwt);
+        var dict = await OneTimeJwtService.Validate(loginModel.Token);
 
         if (dict == null)
         {
-            string error = "{\"status\":false}";
+            await Send("{\"status\":false}");
-            var bytes = Encoding.UTF8.GetBytes(error);
-            await ws.SendAsync(bytes, WebSocketMessageType.Text, WebSocketMessageFlags.EndOfMessage, CancellationToken.None);
             return;
         }
 
-        var _clientId = dict["clientId"];
+        if (!int.TryParse(dict["clientId"], out int clientId))
-        var clientId = int.Parse(_clientId);
+        {
+            await Send("{\"status\":false}");
+            return;
+        }
 
-        var client = NotificationRepository.GetClients().Include(x => x.User).First(x => x.Id == clientId);
+        Client = NotificationClientRepository
+            .Get()
+            .Include(x => x.User)
+            .First(x => x.Id == clientId);
 
-        Client = client;
+        CurrentUser = Client.User;
-        await InitWebsocket();
-        
-        string success = "{\"status\":true}";
-        var byt = Encoding.UTF8.GetBytes(success);
-        await ws.SendAsync(byt, WebSocketMessageType.Text, WebSocketMessageFlags.EndOfMessage, CancellationToken.None);
-    }
 
-    private async Task InitWebsocket()
+        await NotificationServerService.RegisterClient(WebSocket, Client);
-    {
-        NotificationClientService.listenController = this;
-        NotificationClientService.WebsocketReady(Client);
 
-        isAuth = true;
+        await Send("{\"status\":true}");
     }
 
     private async Task Received(string json)
     {
-        var id = JsonConvert.DeserializeObject<NotificationById>(json).notification;
+        var id = JsonConvert.DeserializeObject<NotificationById>(json).Notification;
-        
+
         //TODO: Implement ws notification received
     }
 
     private async Task Read(string json)
     {
-        var id = JsonConvert.DeserializeObject<NotificationById>(json).notification;
+        var model = JsonConvert.DeserializeObject<NotificationById>(json) ?? new();
 
-        await NotificationServerService.SendAction(NotificationClientService.User,
+        await NotificationServerService.SendAction(
-            JsonConvert.SerializeObject(new NotificationById() {Action = "hide", notification = id}));
+            CurrentUser!,
+            JsonConvert.SerializeObject(
+                new NotificationById()
+                {
+                    Action = "hide", Notification = model.Notification
+                }
+            )
+        );
     }
 }

Moonlight/App/Models/Misc/ActiveNotificationClient.cs

@@ -0,0 +1,19 @@
+using System.Net.WebSockets;
+using System.Text;
+using Moonlight.App.Database.Entities.Notification;
+
+namespace Moonlight.App.Models.Misc;
+
+public class ActiveNotificationClient
+{
+    public WebSocket WebSocket { get; set; }
+    public NotificationClient Client { get; set; }
+    
+    public async Task SendAction(string action)
+    {
+        await WebSocket.SendAsync(
+            Encoding.UTF8.GetBytes(action),
+            WebSocketMessageType.Text,
+            WebSocketMessageFlags.EndOfMessage, CancellationToken.None);
+    }
+}

Moonlight/App/Models/Notifications/Login.cs

@@ -1,6 +1,8 @@
-namespace Moonlight.App.Models.Notifications;
+using Newtonsoft.Json;
+
+namespace Moonlight.App.Models.Notifications;
 
 public class Login : BasicWSModel
 {
-    public string token { get; set; }
+    [JsonProperty("token")] public string Token { get; set; } = "";
 }

Moonlight/App/Models/Notifications/NotificationById.cs

@@ -1,6 +1,9 @@
-namespace Moonlight.App.Models.Notifications;
+using Newtonsoft.Json;
+
+namespace Moonlight.App.Models.Notifications;
 
 public class NotificationById : BasicWSModel
 {
-    public int notification { get; set; }
+    [JsonProperty("notification")]
+    public int Notification { get; set; }
 }

Moonlight/App/Services/Notifications/NotificationClientService.cs

@@ -1,43 +0,0 @@
-using System.Net.WebSockets;
-using System.Text;
-using Moonlight.App.Database.Entities;
-using Moonlight.App.Database.Entities.Notification;
-using Moonlight.App.Http.Controllers.Api.Moonlight.Notifications;
-using Moonlight.App.Repositories;
-using Moonlight.App.Services.Sessions;
-
-namespace Moonlight.App.Services.Notifications;
-
-public class NotificationClientService
-{
-    private readonly NotificationRepository NotificationRepository;
-    private readonly NotificationServerService NotificationServerService;
-    internal ListenController listenController;
-    
-    public NotificationClientService(NotificationRepository notificationRepository, NotificationServerService notificationServerService)
-    {
-        NotificationRepository = notificationRepository;
-        NotificationServerService = notificationServerService;
-    }
-
-    public User User => NotificationClient.User;
-    
-    public NotificationClient NotificationClient { get; set; }
-
-    public async Task SendAction(string action)
-    {
-        await listenController.ws.SendAsync(Encoding.UTF8.GetBytes(action), WebSocketMessageType.Text,
-            WebSocketMessageFlags.EndOfMessage, CancellationToken.None);
-    }
-
-    public void WebsocketReady(NotificationClient client)
-    {
-        NotificationClient = client;
-        NotificationServerService.AddClient(this);
-    }
-
-    public void WebsocketClosed()
-    {
-        NotificationServerService.RemoveClient(this);
-    }
-}

Moonlight/App/Services/Notifications/NotificationServerService.cs

@@ -1,84 +1,113 @@
-using Microsoft.EntityFrameworkCore;
+using System.Net.WebSockets;
+using System.Text;
+using Microsoft.EntityFrameworkCore;
 using Moonlight.App.Database.Entities;
 using Moonlight.App.Database.Entities.Notification;
+using Moonlight.App.Events;
+using Moonlight.App.Models.Misc;
 using Moonlight.App.Repositories;
 
 namespace Moonlight.App.Services.Notifications;
 
 public class NotificationServerService
 {
-    private UserRepository UserRepository;
+    private readonly List<ActiveNotificationClient> ActiveClients = new();
-    private NotificationRepository NotificationRepository;
 
     private readonly IServiceScopeFactory ServiceScopeFactory;
-    private IServiceScope ServiceScope;
+    private readonly EventSystem Event;
-    
+
-    public NotificationServerService(IServiceScopeFactory serviceScopeFactory)
+    public NotificationServerService(IServiceScopeFactory serviceScopeFactory, EventSystem eventSystem)
     {
         ServiceScopeFactory = serviceScopeFactory;
-        Task.Run(Run);
+        Event = eventSystem;
-    }
-    
-    private Task Run()
-    {
-        ServiceScope = ServiceScopeFactory.CreateScope();
-
-        UserRepository = ServiceScope
-            .ServiceProvider
-            .GetRequiredService<UserRepository>();
-
-        NotificationRepository = ServiceScope
-            .ServiceProvider
-            .GetRequiredService<NotificationRepository>();
-        
-        return Task.CompletedTask;
     }
 
-    private List<NotificationClientService> connectedClients = new();
+    public Task<ActiveNotificationClient[]> GetActiveClients()
-
-    public List<NotificationClientService> GetConnectedClients()
     {
-        return connectedClients.ToList();
+        lock (ActiveClients)
+        {
+            return Task.FromResult(ActiveClients.ToArray());
+        }
     }
 
-    public List<NotificationClientService> GetConnectedClients(User user)
+    public Task<ActiveNotificationClient[]> GetUserClients(User user)
     {
-        return connectedClients.Where(x => x.User == user).ToList();
+        lock (ActiveClients)
+        {
+            return Task.FromResult(
+                ActiveClients
+                    .Where(x => x.Client.User.Id == user.Id)
+                    .ToArray()
+            );
+        }
     }
 
     public async Task SendAction(User user, string action)
     {
-        var clients = NotificationRepository.GetClients().Include(x => x.User).Where(x => x.User == user).ToList();
+        using var scope = ServiceScopeFactory.CreateScope();
+        var notificationClientRepository =
+            scope.ServiceProvider.GetRequiredService<Repository<NotificationClient>>();
+        
+        var clients = notificationClientRepository
+            .Get()
+            .Include(x => x.User)
+            .Where(x => x.User == user)
+            .ToList();
 
         foreach (var client in clients)
         {
-            var notificationAction = new NotificationAction()
+            ActiveNotificationClient[] connectedUserClients;
-            {
-                Action = action,
-                NotificationClient = client
-            };
-            
-            var connected = connectedClients.Where(x => x.NotificationClient.Id == client.Id).ToList();
 
-            if (connected.Count > 0)
+            lock (ActiveClients)
             {
-                var clientService = connected[0];
+                connectedUserClients = ActiveClients
-                await clientService.SendAction(action);
+                    .Where(x => x.Client.Id == user.Id)
+                    .ToArray();
+            }
+            
+            if (connectedUserClients.Length > 0)
+            {
+                await connectedUserClients[0].SendAction(action);
             }
             else
             {
-                NotificationRepository.AddAction(notificationAction);
+                var notificationAction = new NotificationAction()
+                {
+                    Action = action,
+                    NotificationClient = client
+                };
+                
+                var notificationActionsRepository =
+                    scope.ServiceProvider.GetRequiredService<Repository<NotificationAction>>();
+
+                notificationActionsRepository.Add(notificationAction);
             }
         }
     }
 
-    public void AddClient(NotificationClientService notificationClientService)
+    public async Task RegisterClient(WebSocket webSocket, NotificationClient notificationClient)
     {
-        connectedClients.Add(notificationClientService);
+        var newClient = new ActiveNotificationClient()
+        {
+            WebSocket = webSocket,
+            Client = notificationClient
+        };
+        
+        lock (ActiveClients)
+        {
+            ActiveClients.Add(newClient);
+        }
+
+        await Event.Emit("notifications.addClient", notificationClient);
     }
 
-    public void RemoveClient(NotificationClientService notificationClientService)
+    public async Task UnRegisterClient(NotificationClient client)
     {
-        connectedClients.Remove(notificationClientService);
+        lock (ActiveClients)
+        {
+            ActiveClients.RemoveAll(x => x.Client == client);
+        }
+        
+        await Event.Emit("notifications.removeClient", client);
     }
 }

Moonlight/Program.cs

@@ -118,7 +118,6 @@ namespace Moonlight
             builder.Services.AddScoped<OneTimeJwtService>();
             builder.Services.AddSingleton<NotificationServerService>();
             builder.Services.AddScoped<NotificationAdminService>();
-            builder.Services.AddScoped<NotificationClientService>();
             builder.Services.AddScoped<ModalService>();
             builder.Services.AddScoped<SmartDeployService>();
             builder.Services.AddScoped<WebSpaceService>();

Moonlight/Shared/Views/Admin/Notifications/Debugging.razor

@@ -1,33 +1,74 @@
 @page "/admin/notifications/debugging"
 @using Moonlight.App.Services.Notifications
+@using Moonlight.App.Models.Misc
+@using Moonlight.App.Events
+@using BlazorTable
+@using Moonlight.App.Database.Entities.Notification
+@using Moonlight.App.Services
 
 @inject NotificationServerService NotificationServerService
+@inject SmartTranslateService SmartTranslateService
+@inject EventSystem Event
+
+@implements IDisposable
 
 <OnlyAdmin>
     <LazyLoader Load="Load">
-        <h1>Notification Debugging</h1>
+        <div class="card card-body">
-        @foreach (var client in Clients)
+            <Table TableItem="ActiveNotificationClient" Items="Clients" PageSize="25" TableClass="table table-row-bordered table-row-gray-100 align-middle gs-0 gy-3" TableHeadClass="fw-bold text-muted">
-        {
+                <Column TableItem="ActiveNotificationClient" Title="@(SmartTranslateService.Translate("Id"))" Field="@(x => x.Client.Id)" Sortable="false" Filterable="true"/>
-            <hr/>
+                <Column TableItem="ActiveNotificationClient" Title="@(SmartTranslateService.Translate("User"))" Field="@(x => x.Client.User.Email)" Sortable="false" Filterable="true"/>
-            <div>
+                <Column TableItem="ActiveNotificationClient" Title="" Field="@(x => x.Client.Id)" Sortable="false" Filterable="false">
-                <p>Id: @client.NotificationClient.Id User: @client.User.Email</p>
+                    <Template>
-                <button @onclick="async () => await SendSampleNotification(client)"></button>
+                        <WButton Text="@(SmartTranslateService.Translate("Send notification"))"
-            </div>
+                                 WorkingText="@(SmartTranslateService.Translate("Working"))"
-        }
+                                 CssClasses="btn-primary"
+                                 OnClick="() => SendSampleNotification(context)">
+                        </WButton>
+                    </Template>
+                </Column>
+                <Pager ShowPageNumber="true" ShowTotalCount="true"/>
+            </Table>
+        </div>
     </LazyLoader>
 </OnlyAdmin>
 
-    
+
-@code {
+@code
-    private List<NotificationClientService> Clients;
+{
+    private ActiveNotificationClient[] Clients;
+
+    protected override async Task OnAfterRenderAsync(bool firstRender)
+    {
+        if (firstRender)
+        {
+            await Event.On<NotificationClient>("notifications.addClient", this, async client =>
+            {
+                Clients = await NotificationServerService.GetActiveClients();
+                await InvokeAsync(StateHasChanged);
+            });
+
+            await Event.On<NotificationClient>("notifications.removeClient", this, async client =>
+            {
+                Clients = await NotificationServerService.GetActiveClients();
+                await InvokeAsync(StateHasChanged);
+            });
+        }
+    }
 
     private async Task Load(LazyLoader loader)
     {
-        Clients = NotificationServerService.GetConnectedClients();
+        Clients = await NotificationServerService.GetActiveClients();
     }
 
-    private async Task SendSampleNotification(NotificationClientService client)
+    private async Task SendSampleNotification(ActiveNotificationClient client)
     {
         await client.SendAction(@"{""action"": ""notify"",""notification"":{""id"":999,""channel"":""Sample Channel"",""content"":""This is a sample Notification"",""title"":""Sample Notification"",""url"":""server/9b724fe2-d882-49c9-8c34-3414c7e4a17e""}}");
     }
+
+    public async void Dispose()
+    {
+        await Event.Off("notifications.addClient", this);
+        await Event.Off("notifications.removeClient", this);
+    }
 }