inabox/inabox.server.websocket/WebSocketServer.cs

using InABox.Core;
using InABox.WebSocket.Shared;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Runtime.CompilerServices;
using System.Text;
using System.Threading.Tasks;
using WebSocketSharp;
using WebSocketSharp.Server;
using static InABox.Server.WebSocket.NotifyState;
using InternalServer = WebSocketSharp.Server.WebSocketServer;
using Logger = InABox.Core.Logger;

namespace InABox.Server.WebSocket
{
    public delegate void PollEvent(Session session);

    public class NotifyState
    {
        public class Session
        {
            public string ID { get; set; }

            public Guid SessionID { get; set; }

            public Platform Platform { get; set; }

            public Session(string id, Guid sessionID, Platform platform)
            {
                ID = id;
                SessionID = sessionID;
                Platform = platform;
            }
        }

        public ConcurrentDictionary<Guid, Session> SessionMap = new();

        public event PollEvent? OnPoll;

        public void Poll(Session session)
        {
            OnPoll?.Invoke(session);
        }
    }

    public class NotifyHandler : WebSocketBehavior
    {
        private NotifyState NotifyState;

        public NotifyHandler(NotifyState state)
        {
            NotifyState = state;
        }

        protected override void OnOpen()
        {
            Logger.Send(LogType.Information, "", $"WebSocket client connected");
        }

        protected override void OnClose(CloseEventArgs e)
        {
            Logger.Send(LogType.Information, "", $"WebSocket client disconnected");

            var sessionID = NotifyState.SessionMap.Where(x => x.Value.ID == ID).FirstOrDefault().Key;
            NotifyState.SessionMap.TryRemove(sessionID, out var session);
        }

        protected override void OnError(WebSocketSharp.ErrorEventArgs e)
        {
            Logger.Send(LogType.Error, "", $"WebSocket Error: {e.Message}");
        }

        private void DoInitial(InitialMessage initial)
        {
            var newSession = new Session(ID, initial.SessionID, initial.Platform);
            NotifyState.SessionMap[initial.SessionID] = newSession;
            NotifyState.Poll(newSession);
        }

        protected override void OnMessage(MessageEventArgs e)
        {
            Logger.Send(LogType.Information, "", "Message received");
            var message = SocketMessage.ReadMessage(e.RawData);
            if (message is InitialMessage initial)
            {
                DoInitial(initial);
            }
        }
    }

    public class WebSocketServer
    {
        private InternalServer Server;
        private NotifyState NotifyState = new();

        public event PollEvent? Poll;

        public int Port => Server.Port;

        public WebSocketServer(int port)
        {
            Server = new InternalServer(IPAddress.Any, port);
            Server.AddWebSocketService("/notify", NewNotify);

            NotifyState.OnPoll += NotifyState_Poll;
        }

        private void NotifyState_Poll(Session session)
        {
            Poll?.Invoke(session);
        }

        public NotifyHandler NewNotify()
        {
            return new NotifyHandler(NotifyState);
        }

        public IEnumerable<Guid> GetSessions(Platform platform)
        {
            return NotifyState.SessionMap.Where(x => x.Value.Platform == platform).Select(x => x.Key);
        }

        public void Push(Guid sessionID, SocketMessage message)
        {
            if(NotifyState.SessionMap.TryGetValue(sessionID, out var session))
            {
                using(var stream = new MemoryStream())
                {
                    message.Write(stream);
                    Server.WebSocketServices["/notify"].Sessions.SendToAsync(stream, (int)stream.Length, session.ID, (succ) => { });
                }
            }
        }

        private void PushMessage(SocketMessage message)
        {
            Server.WebSocketServices["/notify"].Sessions.Broadcast(message.WriteToBytes());
        }
        private void PushMessage(SocketMessage message, string session)
        {
            Server.WebSocketServices["/notify"].Sessions.SendTo(message.WriteToBytes(), session);
        }

        public void Push(Type TNotification, object notification)
        {
            PushMessage(NotifyMessage.Notify(TNotification, notification));
        }
        public void Push<TNotification>(TNotification notification)
        {
            PushMessage(NotifyMessage.Notify(notification));
        }

        public void Push(Guid sessionID, Type TNotification, object notification)
        {
            if(NotifyState.SessionMap.TryGetValue(sessionID, out var session))
            {
                PushMessage(NotifyMessage.Notify(TNotification, notification), session.ID);
            }
        }
        public void Push<TNotification>(Guid sessionID, TNotification notification)
        {
            if(NotifyState.SessionMap.TryGetValue(sessionID, out var session))
            {
                PushMessage(NotifyMessage.Notify(notification), session.ID);
            }
        }

        public void Start()
        {
            Server.Start();
        }

        public void Stop()
        {
            Server.Stop();
        }
    }
}