sharp-chat/SharpChat/SockChatServer.cs

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

namespace SharpChat;

public class SockChatServer {
    public const ushort DEFAULT_PORT = 6770;

    public Context Context { get; }

    public bool IsRestarting { get; set; }

    private readonly ILogger Logger;

    private readonly CachedValue<ushort> Port;

    private readonly Lock ConnectionsLock = new();
    private readonly HashSet<SockChatConnection> Connections = [];

    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(this, cancellationTokenSource)
        ]);
    }

    public async Task Listen(CancellationToken cancellationToken) {
        // TODO: protocol servers are now responsible of timing out unauthed connections by themselves

        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(WebSocketCloseCode.InternalError);
                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);
            ILogger logger = Context.LoggerFactory.CreateLogger($"sockchat:({endPoint})");

            if(cancellationToken.IsCancellationRequested) {
                logger.ZLogInformation($"{endPoint} attempted to connect after shutdown was requested. Connection will be dropped.");
                sock.Close(WebSocketCloseCode.TryAgainLater);
                return;
            }

            lock(ConnectionsLock) {
                SockChatConnection conn = new(sock, endPoint, logger);
                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);

        lock(ConnectionsLock) {
            Logger.ZLogDebug($"Disposing all clients...");
            foreach(SockChatConnection conn in Connections)
                conn.Socket.Close(IsRestarting ? WebSocketCloseCode.ServiceRestart : WebSocketCloseCode.GoingAway);
        }
    }

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

    private async Task OnError(SockChatConnection conn, Exception ex) {
        // TODO: detect timeouts and suspend the session
        conn.Logger.ZLogError($"Error: {ex.Message}");
        conn.Logger.ZLogDebug($"{ex}");
        await Context.SafeUpdate();
    }

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

        User? noMoreSessionsUser = null;

        lock(ConnectionsLock) {
            Connections.Remove(conn);

            Session? session = Context.Sessions.GetSession(conn);
            if(session is not null) {
                if(!Context.Sessions.IsSuspendedSession(session))
                    Context.Sessions.DestroySession(session);

                if(Context.Sessions.CountActiveSessions(session.UserId) < 1)
                    noMoreSessionsUser = Context.Users.GetUser(session.UserId);
            }
        }

        Context.ContextAccess.Wait();
        try {
            if(noMoreSessionsUser is not null)
                await Context.HandleDisconnect(noMoreSessionsUser);

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

    private async Task OnMessage(SockChatConnection conn, string msg) {
        conn.Logger.ZLogTrace($"Received: {msg}");

        await Context.SafeUpdate();

        Session? session = Context.Sessions.GetSession(conn);
        User? user = session is null ? null : Context.Users.GetUser(session.UserId);

        // this doesn't affect non-authed connections?????
        if(user is not null && user.Rank < Context.FloodKickExemptRank) {
            bool rusticate = false;
            string banAddr = string.Empty;
            TimeSpan banDuration = TimeSpan.MinValue;

            Context.ContextAccess.Wait();
            try {
                if(!Context.UserRateLimiters.TryGetValue(user.UserId, out RateLimiter? rateLimiter))
                    Context.UserRateLimiters.Add(user.UserId, rateLimiter = new RateLimiter());

                rateLimiter.Update();

                if(rateLimiter.IsExceeded) {
                    banDuration = TimeSpan.FromSeconds(Context.FloodKickLength);
                    rusticate = true;
                    banAddr = conn.RemoteEndPoint.Address.ToString();
                    conn.Logger.ZLogWarning($"Exceeded flood limit! Issuing ban with duration {banDuration} on {banAddr}/{user.UserId}...");
                } else if(rateLimiter.IsRisky) {
                    rusticate = true;
                    banAddr = conn.RemoteEndPoint.Address.ToString();
                    conn.Logger.ZLogWarning($"About to exceed flood limit! Issueing warning to {banAddr}/{user.UserId}...");
                }

                if(rusticate) {
                    if(banDuration == TimeSpan.MinValue) {
                        await Context.SendTo(user, new CommandResponseS2CPacket(Context.RandomSnowflake.Next(), LCR.FLOOD_WARN, false));
                    } else {
                        await Context.BanUser(user, banDuration, UserDisconnectS2CPacket.Reason.Flood);

                        if(banDuration > TimeSpan.Zero)
                            await Context.Bans.BanCreate(
                                BanKind.User,
                                banDuration,
                                conn.RemoteEndPoint.Address,
                                user.UserId,
                                "Kicked from chat for flood protection.",
                                IPAddress.IPv6Loopback
                            );

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

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

        if(handler is not null)
            await handler.Handle(context);
    }
}