sharp-chat/SharpChat/SockChatServer.cs

using Fleck;
using SharpChat.Auth;
using SharpChat.Bans;
using SharpChat.C2SPacketHandlers;
using SharpChat.ClientCommands;
using SharpChat.Configuration;
using SharpChat.SockChat.S2CPackets;
using System.Net;

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 Context Context { get; }

    private readonly BansClient BansClient;

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

    private readonly List<C2SPacketHandler> GuestHandlers = [];
    private readonly List<C2SPacketHandler> AuthedHandlers = [];
    private readonly SendMessageC2SPacketHandler SendMessageHandler;

    private bool IsShuttingDown = false;

    private static readonly string[] DEFAULT_CHANNELS = ["lounge"];

    private Channel DefaultChannel { get; set; }

    public SockChatServer(
        AuthClient authClient,
        BansClient bansClient,
        EventStorage.EventStorage evtStore,
        Config config
    ) {
        Logger.Write("Initialising Sock Chat server...");

        BansClient = bansClient;

        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 Context(evtStore);

        string[]? channelNames = config.ReadValue("channels", DEFAULT_CHANNELS);
        if(channelNames is not null)
            foreach(string channelName in channelNames) {
                Config channelCfg = config.ScopeTo($"channels:{channelName}");

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

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

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

        if(DefaultChannel is null)
            throw new Exception("The default channel could not be determined.");

        GuestHandlers.Add(new AuthC2SPacketHandler(authClient, bansClient, DefaultChannel, MaxMessageLength, MaxConnections));

        AuthedHandlers.AddRange([
            new PingC2SPacketHandler(authClient),
            SendMessageHandler = new SendMessageC2SPacketHandler(Context.RandomSnowflake, MaxMessageLength),
        ]);

        SendMessageHandler.AddCommands([
            new AFKClientCommand(),
            new NickClientCommand(),
            new WhisperClientCommand(),
            new ActionClientCommand(),
            new WhoClientCommand(),
            new JoinChannelClientCommand(),
            new CreateChannelClientCommand(),
            new DeleteChannelClientCommand(),
            new PasswordChannelClientCommand(),
            new RankChannelClientCommand(),
            new BroadcastClientCommand(),
            new DeleteMessageClientCommand(),
            new KickBanClientCommand(bansClient),
            new PardonUserClientCommand(bansClient),
            new PardonAddressClientCommand(bansClient),
            new BanListClientCommand(bansClient),
            new RemoteAddressClientCommand(),
        ]);

        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 ShutdownRestartClientCommand(waitHandle, () => !IsShuttingDown && (IsShuttingDown = true)));

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

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

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

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

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

    private async Task OnError(Connection conn, Exception ex) {
        Logger.Write($"[{conn.Id} {conn.RemoteAddress}] {ex}");
        await Context.SafeUpdate();
    }

    private async Task OnClose(Connection 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))
                await Context.HandleDisconnect(conn.User);

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

    private async Task OnMessage(Connection conn, string msg) {
        await Context.SafeUpdate();

        // this doesn't affect non-authed connections?????
        if(conn.User is not null && conn.User.Rank < FloodKickExemptRank) {
            User? 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(
                        User.DEFAULT_SIZE,
                        User.DEFAULT_MINIMUM_DELAY,
                        User.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) {
                        await Context.SendTo(conn.User, new CommandResponseS2CPacket(Context.RandomSnowflake.Next(), LCR.FLOOD_WARN, false));
                    } else {
                        await Context.BanUser(conn.User, banDuration, UserDisconnectS2CPacket.Reason.Flood);

                        if(banDuration > TimeSpan.Zero)
                            await BansClient.BanCreateAsync(
                                BanKind.User,
                                banDuration,
                                conn.RemoteAddress,
                                conn.User.UserId,
                                "Kicked from chat for flood protection.",
                                IPAddress.IPv6Loopback
                            );

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

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

        if(handler is not null)
            await 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(Connection conn in Context.Connections)
            conn.Dispose();

        Server?.Dispose();
    }
}