sharp-chat/SharpChat/SockChatServer.cs

using Fleck;
using SharpChat.Commands;
using SharpChat.Config;
using SharpChat.EventStorage;
using SharpChat.Misuzu;
using SharpChat.Packet;
using SharpChat.PacketHandlers;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;

namespace SharpChat {
    public class SockChatServer : IDisposable {
        public const ushort DEFAULT_PORT = 6770;
        public const int DEFAULT_MSG_LENGTH_MAX = 5000;
        public const int DEFAULT_MAX_CONNECTIONS = 5;
        public const int DEFAULT_FLOOD_KICK_LENGTH = 30;

        public bool IsDisposed { get; private set; }

        public static ChatUser Bot { get; } = new ChatUser {
            UserId = -1,
            Username = "ChatBot",
            Rank = 0,
            Colour = new ChatColour(),
        };

        public IWebSocketServer Server { get; }
        public ChatContext Context { get; }

        private readonly HttpClient HttpClient;
        private readonly MisuzuClient Misuzu;

        private readonly CachedValue<int> MaxMessageLength;
        private readonly CachedValue<int> MaxConnections;
        private readonly CachedValue<int> FloodKickLength;

        private readonly List<IChatPacketHandler> GuestHandlers = new();
        private readonly List<IChatPacketHandler> AuthedHandlers = new();
        private readonly SendMessageHandler SendMessageHandler;

        private bool IsShuttingDown = false;

        private ChatChannel DefaultChannel { get; set; }

        public SockChatServer(HttpClient httpClient, MisuzuClient msz, IEventStorage evtStore, IConfig config) {
            Logger.Write("Initialising Sock Chat server...");

            HttpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
            Misuzu = msz ?? throw new ArgumentNullException(nameof(msz));

            MaxMessageLength = config.ReadCached("msgMaxLength", DEFAULT_MSG_LENGTH_MAX);
            MaxConnections = config.ReadCached("connMaxCount", DEFAULT_MAX_CONNECTIONS);
            FloodKickLength = config.ReadCached("floodKickLength", DEFAULT_FLOOD_KICK_LENGTH);

            Context = new ChatContext(evtStore);

            string[] channelNames = config.ReadValue("channels", new[] { "lounge" });

            foreach(string channelName in channelNames) {
                ChatChannel channelInfo = new(channelName);
                IConfig channelCfg = config.ScopeTo($"channels:{channelName}");

                string tmp;
                tmp = channelCfg.SafeReadValue("name", string.Empty);
                if(!string.IsNullOrWhiteSpace(tmp))
                    channelInfo.Name = tmp;

                channelInfo.Password = channelCfg.SafeReadValue("password", string.Empty);
                channelInfo.Rank = channelCfg.SafeReadValue("minRank", 0);

                Context.Channels.Add(channelInfo);

                DefaultChannel ??= channelInfo;
            }

            GuestHandlers.Add(new AuthHandler(Misuzu, DefaultChannel, MaxMessageLength, MaxConnections));

            AuthedHandlers.AddRange(new IChatPacketHandler[] {
                new PingHandler(Misuzu),
                SendMessageHandler = new SendMessageHandler(MaxMessageLength),
            });

            SendMessageHandler.AddCommands(new IChatCommand[] {
                new AFKCommand(),
                new NickCommand(),
                new WhisperCommand(),
                new ActionCommand(),
                new WhoCommand(),
                new JoinChannelCommand(),
                new CreateChannelCommand(),
                new DeleteChannelCommand(),
                new PasswordChannelCommand(),
                new RankChannelCommand(),
                new BroadcastCommand(),
                new DeleteMessageCommand(),
                new KickBanCommand(msz),
                new PardonUserCommand(msz),
                new PardonAddressCommand(msz),
                new BanListCommand(msz),
                new SilenceApplyCommand(),
                new SilenceRevokeCommand(),
                new RemoteAddressCommand(),
            });

            ushort port = config.SafeReadValue("port", DEFAULT_PORT);
            Server = new SharpChatWebSocketServer($"ws://0.0.0.0:{port}");
        }

        public void Listen(ManualResetEvent waitHandle) {
            if(waitHandle != null)
                SendMessageHandler.AddCommand(new ShutdownRestartCommand(waitHandle, () => !IsShuttingDown && (IsShuttingDown = true)));

            Server.Start(sock => {
                if(IsShuttingDown || IsDisposed) {
                    sock.Close(1013);
                    return;
                }

                sock.OnOpen = () => OnOpen(sock);
                sock.OnClose = () => OnClose(sock);
                sock.OnError = err => OnError(sock, err);
                sock.OnMessage = msg => OnMessage(sock, msg);
            });

            Logger.Write("Listening...");
        }

        private void OnOpen(IWebSocketConnection sock) {
            Logger.Write($"Connection opened from {sock.ConnectionInfo.ClientIpAddress}:{sock.ConnectionInfo.ClientPort}");

            lock(Context.ConnectionsAccess) {
                if(!Context.Connections.Any(x => x.Socket == sock))
                    Context.Connections.Add(new ChatConnection(sock));
            }

            Context.Update();
        }

        private void OnClose(IWebSocketConnection sock) {
            Logger.Write($"Connection closed from {sock.ConnectionInfo.ClientIpAddress}:{sock.ConnectionInfo.ClientPort}");

            ChatConnection conn;
            lock(Context.ConnectionsAccess)
                conn = Context.GetConnection(sock);

            // Remove connection from user
            if(conn?.User != null) {
                // RemoveConnection sets conn.User to null so we must grab a local copy.
                ChatUser user = conn.User;

                user.RemoveConnection(conn);

                if(!user.HasConnections)
                    Context.UserLeave(null, user);
            }

            // Update context
            Context.Update();

            // Remove connection from server
            lock(Context.ConnectionsAccess)
                Context.Connections.Remove(conn);

            conn?.Dispose();
        }

        private void OnError(IWebSocketConnection sock, Exception ex) {
            string connId;
            lock(Context.ConnectionsAccess) {
                ChatConnection conn = Context.GetConnection(sock);
                connId = conn?.Id ?? new string('0', ChatConnection.ID_LENGTH);
            }

            Logger.Write($"[{connId} {sock.ConnectionInfo.ClientIpAddress}] {ex}");
            Context.Update();
        }

        private void OnMessage(IWebSocketConnection sock, string msg) {
            Context.Update();

            ChatConnection conn;
            lock(Context.ConnectionsAccess)
                conn = Context.GetConnection(sock);

            if(conn == null) {
                sock.Close();
                return;
            }

            // this doesn't affect non-authed connections?????
            if(conn.User is not null && conn.User.HasFloodProtection) {
                conn.User.RateLimiter.AddTimePoint();

                if(conn.User.RateLimiter.State == ChatRateLimitState.Kick) {
                    Task.Run(async () => {
                        TimeSpan duration = TimeSpan.FromSeconds(FloodKickLength);

                        await Misuzu.CreateBanAsync(
                            conn.User.UserId.ToString(), conn.RemoteAddress.ToString(),
                            string.Empty, "::1",
                            duration,
                            "Kicked from chat for flood protection."
                        );

                        Context.BanUser(conn.User, duration, UserDisconnectReason.Flood);
                    }).Wait();
                    return;
                } else if(conn.User.RateLimiter.State == ChatRateLimitState.Warning)
                    conn.User.Send(new FloodWarningPacket());
            }

            ChatPacketHandlerContext context = new(msg, Context, conn);
            IChatPacketHandler handler = conn.User is null
                ? GuestHandlers.FirstOrDefault(h => h.IsMatch(context))
                : AuthedHandlers.FirstOrDefault(h => h.IsMatch(context));

            handler?.Handle(context);
        }

        ~SockChatServer() {
            DoDispose();
        }

        public void Dispose() {
            DoDispose();
            GC.SuppressFinalize(this);
        }

        private void DoDispose() {
            if(IsDisposed)
                return;
            IsDisposed = true;

            lock(Context.ConnectionsAccess)
                foreach(ChatConnection conn in Context.Connections)
                    conn.Dispose();

            Server?.Dispose();
            HttpClient?.Dispose();
        }
    }
}