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;

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 const int DEFAULT_FLOOD_KICK_EXEMPT_RANK = 9;

        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 CachedValue<int> FloodKickExemptRank;

        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;
            Misuzu = msz;

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

            Context = new ChatContext(evtStore);

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

            if(channelNames != null)
                foreach(string channelName in channelNames) {
                    IConfig channelCfg = config.ScopeTo($"channels:{channelName}");

                    string? name = channelCfg.SafeReadValue("name", string.Empty);
                    if(string.IsNullOrWhiteSpace(name))
                        name = channelName;

                    ChatChannel channelInfo = new(
                        name,
                        channelCfg.SafeReadValue("password", string.Empty),
                        rank: channelCfg.SafeReadValue("minRank", 0)
                    );

                    Context.Channels.Add(channelInfo.Name, channelInfo);
                    DefaultChannel ??= channelInfo;
                }

            DefaultChannel ??= new ChatChannel("Default");
            if(Context.Channels.Count < 1)
                Context.Channels.Add(DefaultChannel.Name, DefaultChannel);

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

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

            SendMessageHandler.AddCommands(new IChatCommand[] {
                new UserAFKCommand(),
                new UserNickCommand(),
                new MessageWhisperCommand(),
                new MessageActionCommand(),
                new WhoCommand(),
                new ChannelJoinCommand(),
                new ChannelCreateCommand(),
                new ChannelDeleteCommand(),
                new ChannelPasswordCommand(),
                new ChannelRankCommand(),
                new MessageBroadcastCommand(),
                new MessageDeleteCommand(),
                new KickBanCommand(msz),
                new PardonUserCommand(msz),
                new PardonAddressCommand(msz),
                new BanListCommand(msz),
                new WhoisCommand(),
            });

            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) {
                    sock.Close(1013);
                    return;
                }

                ChatConnection conn = new(sock);
                Context.Connections.Add(conn);

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

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

        private void OnOpen(ChatConnection conn) {
            Logger.Write($"Connection opened from {conn.RemoteAddress}:{conn.RemotePort}");
            Context.SafeUpdate();
        }

        private void OnError(ChatConnection conn, Exception ex) {
            Logger.Write($"[{conn.Id} {conn.RemoteAddress}] {ex}");
            Context.SafeUpdate();
        }

        private void OnClose(ChatConnection conn) {
            Logger.Write($"Connection closed from {conn.RemoteAddress}:{conn.RemotePort}");

            Context.ContextAccess.Wait();
            try {
                Context.Connections.Remove(conn);

                if(conn.User != null && !Context.Connections.Any(c => c.User == conn.User))
                    Context.HandleDisconnect(conn.User);

                Context.Update();
            } finally {
                Context.ContextAccess.Release();
            }
        }

        private void OnMessage(ChatConnection conn, string msg) {
            Context.SafeUpdate();

            // this doesn't affect non-authed connections?????
            if(conn.User is not null && conn.User.Rank < FloodKickExemptRank) {
                ChatUser? banUser = null;
                string banAddr = string.Empty;
                TimeSpan banDuration = TimeSpan.MinValue;

                Context.ContextAccess.Wait();
                try {
                    if(!Context.UserRateLimiters.TryGetValue(conn.User.UserId, out RateLimiter? rateLimiter))
                        Context.UserRateLimiters.Add(conn.User.UserId, rateLimiter = new RateLimiter(
                            ChatUser.DEFAULT_SIZE,
                            ChatUser.DEFAULT_MINIMUM_DELAY,
                            ChatUser.DEFAULT_RISKY_OFFSET
                        ));

                    rateLimiter.Update();

                    if(rateLimiter.IsExceeded) {
                        banDuration = TimeSpan.FromSeconds(FloodKickLength);
                        banUser = conn.User;
                        banAddr = conn.RemoteAddress.ToString();
                    } else if(rateLimiter.IsRisky) {
                        banUser = conn.User;
                    }

                    if(banUser is not null) {
                        if(banDuration == TimeSpan.MinValue) {
                            Context.SendTo(conn.User, new FloodWarningPacket());
                        } else {
                            Context.BanUser(conn.User, banDuration, ChatUserDisconnectReason.Flood);

                            if(banDuration > TimeSpan.Zero)
                                Misuzu.CreateBanAsync(
                                    conn.User.UserId.ToString(), conn.RemoteAddress.ToString(),
                                    string.Empty, "::1",
                                    banDuration,
                                    "Kicked from chat for flood protection."
                                ).Wait();

                            return;
                        }
                    }
                } finally {
                    Context.ContextAccess.Release();
                }
            }

            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);
        }

        private bool IsDisposed;

        ~SockChatServer() {
            DoDispose();
        }

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

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

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

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