sharp-chat/SharpChat/SockChatServer.cs

using Microsoft.Extensions.Logging;
using SharpChat.Bans;
using SharpChat.C2SPacketHandlers;
using SharpChat.ClientCommands;
using SharpChat.Configuration;
using SharpChat.SockChat.S2CPackets;
using System.Net;
using ZLogger;

namespace SharpChat;

public class SockChatServer {
    public const ushort DEFAULT_PORT = 6770;

    public Context Context { get; }

    private readonly ILogger Logger;

    private readonly CachedValue<ushort> Port;

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

    public SockChatServer(
        CancellationTokenSource cancellationTokenSource,
        Context ctx,
        Config config
    ) {
        Logger = ctx.LoggerFactory.CreateLogger("sockchat");
        Logger.ZLogInformation($"Initialising Sock Chat server...");

        Context = ctx;

        Logger.ZLogDebug($"Fetching configuration values...");
        Port = config.ReadCached("port", DEFAULT_PORT);

        Logger.ZLogDebug($"Registering unauthenticated handlers...");
        GuestHandlers.Add(new AuthC2SPacketHandler(
            Context.Auth,
            Context.Bans,
            Context.Channels,
            Context.RandomSnowflake,
            Context.MaxMessageLength,
            Context.MaxConnections
        ));

        Logger.ZLogDebug($"Registering authenticated handlers...");
        AuthedHandlers.AddRange([
            new PingC2SPacketHandler(Context.Auth),
            SendMessageHandler = new SendMessageC2SPacketHandler(Context.RandomSnowflake, Context.MaxMessageLength),
        ]);

        Logger.ZLogDebug($"Registering client commands...");
        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(Context.Bans),
            new PardonUserClientCommand(Context.Bans),
            new PardonAddressClientCommand(Context.Bans),
            new BanListClientCommand(Context.Bans),
            new RemoteAddressClientCommand(),
            new ShutdownRestartClientCommand(cancellationTokenSource)
        ]);
    }

    public async Task Listen(CancellationToken cancellationToken) {
        using SharpChatWebSocketServer server = new(Context.LoggerFactory.CreateLogger("sockchat:server"), $"ws://0.0.0.0:{Port}");
        server.Start(sock => {
            if(!IPAddress.TryParse(sock.ConnectionInfo.ClientIpAddress, out IPAddress? addr)) {
                Logger.ZLogError($@"A client attempted to connect with an invalid IP address: ""{sock.ConnectionInfo.ClientIpAddress}""");
                sock.Close(1011);
                return;
            }

            if(IPAddress.IsLoopback(addr) && sock.ConnectionInfo.Headers.TryGetValue("X-Real-IP", out string? addrStr)) {
                if(IPAddress.TryParse(addrStr, out IPAddress? realAddr))
                    addr = realAddr;
                else
                    Logger.ZLogWarning($@"Connection originated from loopback and supplied an X-Real-IP header, but it could not be parsed: ""{addrStr}""");
            }

            IPEndPoint endPoint = new(addr, sock.ConnectionInfo.ClientPort);

            if(cancellationToken.IsCancellationRequested) {
                Logger.ZLogInformation($"{endPoint} attepted to connect after shutdown was requested. Connection will be dropped.");
                sock.Close(1013);
                return;
            }

            ILogger logger = Context.LoggerFactory.CreateLogger($"sockchat:({endPoint})");
            Connection conn = new(logger, sock, endPoint);
            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.ZLogInformation($"Listening...");
        await Task.Delay(Timeout.Infinite, cancellationToken).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

        Logger.ZLogDebug($"Disposing all clients...");
        foreach(Connection conn in Context.Connections)
            conn.Dispose();
    }

    private async Task OnOpen(Connection conn) {
        conn.Logger.ZLogInformation($"Connection opened.");
        await Context.SafeUpdate();
    }

    private async Task OnError(Connection conn, Exception ex) {
        conn.Logger.ZLogError($"Error: {ex.Message}");
        conn.Logger.ZLogDebug($"{ex}");
        await Context.SafeUpdate();
    }

    private async Task OnClose(Connection conn) {
        conn.Logger.ZLogInformation($"Connection closed.");

        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) {
        conn.Logger.ZLogTrace($"Received: {msg}");

        await Context.SafeUpdate();

        // this doesn't affect non-authed connections?????
        if(conn.User is not null && conn.User.Rank < Context.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(Context.FloodKickLength);
                    banUser = conn.User;
                    banAddr = conn.RemoteEndPoint.Address.ToString();
                    conn.Logger.ZLogWarning($"Exceeded flood limit! Issuing ban with duration {banDuration} on {banAddr}/{banUser.UserId}...");
                } else if(rateLimiter.IsRisky) {
                    banUser = conn.User;
                    banAddr = conn.RemoteEndPoint.Address.ToString();
                    conn.Logger.ZLogWarning($"About to exceed flood limit! Issueing warning to {banAddr}/{banUser.UserId}...");
                }

                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 Context.Bans.BanCreate(
                                BanKind.User,
                                banDuration,
                                conn.RemoteEndPoint.Address,
                                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);
    }
}