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.PushState;
using InternalServer = WebSocketSharp.Server.WebSocketServer;
using Logger = InABox.Core.Logger;

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

    public class PushState
    {
        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 PushHandler : WebSocketBehavior
    {
        public PushState State { get; set; }

        public PushHandler()
        {
        }
        
        public PushHandler(PushState state)
        {
            State = 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 = State.SessionMap.FirstOrDefault(x => x.Value.ID == ID).Key;
            State.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);
            State.SessionMap[initial.SessionID] = newSession;
            State.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 PushState PushState = new();

        public event PollEvent? Poll;

        public int Port => Server.Port;

        public WebSocketServer(int port)
        {
            Server = new InternalServer(IPAddress.Any, port);
            Server.AddWebSocketService<PushHandler>("/push", (push) =>
            {
                push.State = PushState;
            });

            PushState.OnPoll += PushState_Poll;
        }

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

        public PushHandler NewPush()
        {
            return new PushHandler(PushState);
        }

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

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

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

        public void Push(Type TPush, object push)
        {
            PushMessage(InABox.WebSocket.Shared.PushMessage.Push(TPush, push));
        }
        public void Push<TPush>(TPush push)
            where TPush : BaseObject
        {
            PushMessage(InABox.WebSocket.Shared.PushMessage.Push(push));
        }

        public void Push(Guid sessionID, Type TPush, object push)
        {
            if(PushState.SessionMap.TryGetValue(sessionID, out var session))
            {
                PushMessage(InABox.WebSocket.Shared.PushMessage.Push(TPush, push), session.ID);
            }
        }
        public void Push<TPush>(Guid sessionID, TPush push)
            where TPush : BaseObject
        {
            if(PushState.SessionMap.TryGetValue(sessionID, out var session))
            {
                PushMessage(InABox.WebSocket.Shared.PushMessage.Push(push), session.ID);
            }
        }

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

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