sharp-chat/SharpChat.SockChat/SockChatServer.cs

using Fleck;
using SharpChat.Config;
using SharpChat.EventStorage;
using SharpChat.Misuzu;
using SharpChat.SockChat.Commands;
using SharpChat.SockChat.PacketsC2S;
using SharpChat.SockChat.PacketsS2C;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Threading;

namespace SharpChat.SockChat {
    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 IWebSocketServer Server { get; }
        public SockChatContext 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<IC2SPacketHandler> GuestHandlers = new();
        private readonly List<IC2SPacketHandler> AuthedHandlers = new();
        private readonly SendMessageC2SPacketHandler SendMessageHandler;

        private bool IsShuttingDown = false;
        private bool IsRestarting = false;

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

            DateTimeOffset started = DateTimeOffset.UtcNow;
            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);

            Context = new SockChatContext(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;

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

                    Context.Channels.Add(channelInfo);
                }

            if(Context.Channels.PublicCount < 1)
                Context.Channels.Add(new ChannelInfo("Default"));

            GuestHandlers.Add(new AuthC2SPacketHandler(
                started,
                Misuzu,
                Context.Channels.MainChannel,
                MaxMessageLength,
                MaxConnections
            ));

            AuthedHandlers.AddRange(new IC2SPacketHandler[] {
                new PingC2SPacketHandler(Misuzu),
                SendMessageHandler = new SendMessageC2SPacketHandler(MaxMessageLength),
            });

            SendMessageHandler.AddCommands(new ISockChatClientCommand[] {
                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 SockChatWebSocketServer($"ws://0.0.0.0:{port}");
        }

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

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

                SockChatConnectionInfo 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(SockChatConnectionInfo conn) {
            Logger.Write($"Connection opened from {conn.RemoteEndPoint}");
            Context.SafeUpdate();
        }

        private void OnError(SockChatConnectionInfo conn, Exception ex) {
            Logger.Write($"<{conn.RemoteEndPoint}> {ex}");
            Context.SafeUpdate();
        }

        private void OnClose(SockChatConnectionInfo conn) {
            Logger.Write($"Connection closed from {conn.RemoteEndPoint}");

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

                if(!Context.Connections.HasUser(conn.UserId)) {
                    UserInfo? userInfo = Context.Users.Get(conn.UserId);
                    if(userInfo != null)
                        Context.HandleDisconnect(userInfo);
                }

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

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

            // this doesn't affect non-authed connections?????
            if(conn.UserId > 0) {
                long banUserId = 0;
                string banAddr = string.Empty;
                TimeSpan banDuration = TimeSpan.MinValue;

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

                    rateLimiter.Update();

                    if(rateLimiter.IsExceeded) {
                        banDuration = TimeSpan.FromSeconds(FloodKickLength);
                        banUserId = conn.UserId;
                        banAddr = conn.RemoteAddress;
                    } else if(rateLimiter.IsRisky) {
                        banUserId = conn.UserId;
                    }

                    if(banUserId > 0) {
                        UserInfo? userInfo = Context.Users.Get(banUserId);
                        if(userInfo != null) {
                            if(banDuration == TimeSpan.MinValue) {
                                Context.SendTo(userInfo, new FloodWarningS2CPacket());
                            } else {
                                Context.BanUser(userInfo, banDuration, UserDisconnectReason.Flood);

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

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

            C2SPacketHandlerContext context = new(msg, Context, conn);
            IC2SPacketHandler? handler = conn.UserId > 0
                ? AuthedHandlers.FirstOrDefault(h => h.IsMatch(context))
                : GuestHandlers.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;

            Context.Connections.WithAll(conn => conn.Close(IsRestarting ? 1012 : 1001));

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