inabox/InABox.Server/RPC/Transports/RPCServerTransport.cs

using System.Collections.Concurrent;
using System.Text;
using InABox.Core;

namespace InABox.Rpc
{

    public abstract class RpcServerTransport<TConnection> : IPusher, IRpcServerTransport where TConnection : notnull
    {
        public abstract bool IsSecure();
        
        private ConcurrentDictionary<TConnection, RpcServerSession> _sessions = new ConcurrentDictionary<TConnection, RpcServerSession>();

        protected RpcServerSession CreateSession(TConnection connection)
        {
            var result = new RpcServerSession();
            _sessions[connection] = result;
            return result;
        }

        public RpcServerSession? GetSession(TConnection? connection)
        {
            if (connection == null)
                return null;
            _sessions.TryGetValue(connection, out RpcServerSession? result);
            return result;
        }

        protected RpcServerSession? DeleteSession(TConnection connection)
        {
            _sessions.Remove(connection, out RpcServerSession? session);
            return session;
        }
        
        private Dictionary<string, IRpcCommandHandler> _handlers = new Dictionary<string, IRpcCommandHandler>();

        public void AddHandler<TSender, TCommand, TProperties, TResult>(RpcCommandHandler<TSender, TCommand, TProperties, TResult> handler)
            where TSender : class
            where TCommand : IRpcCommand<TProperties, TResult>
            where TProperties : IRpcCommandParameters, new()
            where TResult : IRpcCommandResult, new()
        {
            _handlers[typeof(TCommand).Name] = handler;
        }
        
        public abstract void Start();

        public abstract void Stop();
        
        public event RpcTransportOpenEvent? OnOpen;

        protected void DoOpen(TConnection connection)
        {
            var session = CreateSession(connection);
            OnOpen?.Invoke(this, new RpcTransportOpenArgs(session) );
        }

        public event RpcTransportCloseEvent? OnClose;

        protected void DoClose(TConnection connection, RpcTransportCloseEventType type)
        {
            var session = DeleteSession(connection);
            OnClose?.Invoke(this, new RpcTransportCloseArgs(session, type));
        }

        public event RpcTransportExceptionEvent? OnException;
        protected void DoException(TConnection? connection, Exception e)
        {
            var session = GetSession(connection);
            OnException?.Invoke(this, new RpcTransportExceptionArgs(session, e));
        }

        public event RpcTransportMessageEvent? BeforeMessage;

        protected void DoBeforeMessage(RpcServerSession? session, RpcMessage? message)
        {
            BeforeMessage?.Invoke(this, new RpcTransportMessageArgs(session,message));
        }

        public event RpcTransportMessageEvent? AfterMessage;
        protected void DoAfterMessage(RpcServerSession? session, RpcMessage? message)
        {
            AfterMessage?.Invoke(this, new RpcTransportMessageArgs(session,message));
        }
        
        /// <summary>
        /// Handle a message from a client.
        /// </summary>
        /// <param name="connection">The client connection.</param>
        /// <param name="message">The message to be handled.</param>
        /// <returns>The response to be sent back to the client.</returns>
        public RpcMessage? DoMessage(TConnection? connection, RpcMessage? message)
        {
            if(message is null)
            {
                DoException(connection, new Exception("NULL Message Received"));
                return null;
            }
            var response = new RpcMessage() { Id = message.Id, Command = message.Command };
            try
            {
                var session = GetSession(connection);
                DoBeforeMessage(session, message);
                if (session != null)
                {
                    response = new RpcMessage() { Id = message.Id, Command = message.Command };

                    if (_handlers.TryGetValue(message.Command, out var command))
                    {
                        try
                        {
                            response.Payload = command.Execute(session, message.Payload);
                        }
                        catch (RpcException err)
                        {
                            response.Payload = Encoding.UTF8.GetBytes(err.Message);
                            response.Error = err.Error;
                        }
                    }
                    else
                    {
                        DoException(connection, new Exception("Command Not Found"));
                        response.Error = RpcError.COMMANDNOTFOUND;
                    }

                    DoAfterMessage(session, response);
                }
                else
                {
                    DoException(connection, new Exception("Session not Found"));
                    response.Error = RpcError.SESSIONNOTFOUND;
                }
            }
            catch(Exception e)
            {
                DoException(connection, e);
                response.Payload = Encoding.UTF8.GetBytes(e.Message);
                response.Error = RpcError.SERVERERROR;
            }

            return response;
        }

        /// <summary>
        /// Send a message to a particular client connection.
        /// </summary>
        /// <param name="connection">The connection to send to.</param>
        /// <param name="message">The message to send.</param>
        public abstract void Send(TConnection connection, RpcMessage message);

        #region Pusher Stuff

        public void PushToAll<TPush>(TPush push) where TPush : BaseObject
        {
            var message = new RpcMessage(Guid.NewGuid(), "Push", RpcPush.Create(push).WriteBinary(BinarySerializationSettings.Latest));
            foreach (var connection in _sessions.Keys)
            {
                Send(connection, message);
            }
        }

        public void PushToSession<TPush>(Guid session, TPush push) where TPush : BaseObject
        {
            var message = new RpcMessage(Guid.NewGuid(), "Push", RpcPush.Create(push).WriteBinary(BinarySerializationSettings.Latest));

            var sessionConnection = _sessions.FirstOrDefault(x => x.Value.ID == session).Key;
            if(sessionConnection is not null)
            {
                Send(sessionConnection, message);
            }
        }

        public void PushToSession(Guid session, Type TPush, BaseObject push)
        {
            var message = new RpcMessage(Guid.NewGuid(), "Push", new RpcPush
            {
                Object = push,
                Type = TPush
            }.WriteBinary(BinarySerializationSettings.Latest));

            var sessionConnection = _sessions.FirstOrDefault(x => x.Value.ID == session).Key;
            if (sessionConnection is not null)
            {
                Send(sessionConnection, message);
            }
        }

        public IEnumerable<Guid> GetUserSessions(Guid user) =>
            _sessions.Values.Where(x => x.UserGuid == user).Select(x => x.ID);

        public IEnumerable<Guid> GetSessions(Platform platform) =>
            _sessions.Values.Where(x => x.Platform == platform).Select(x => x.ID);

        #endregion
    }
}