DoSProtection.cs

#nullable enable
 
using System;
using System.Collections.Generic;
using System.Diagnostics;
using Barotrauma.Networking;
 
namespace Barotrauma
{
    internal sealed class DoSProtection
    {
        public readonly ref struct DoSAction
        {
            private readonly Client sender;
            private readonly Action<Client> end;
 
            public DoSAction(Client sender, Action<Client> start, Action<Client> end)
            {
                this.sender = sender;
                this.end = end;
                start(sender);
            }
 
            public void Dispose()
            {
                end(sender);
            }
        }
 
        private sealed class OffenseData
        {
            public readonly Stopwatch Stopwatch = new();
 
            public int Strikes;
 
            public int PacketCount;
 
            public void ResetStrikes()
            {
                Strikes = 0;
                PacketCount = 0;
            }
 
            public void ResetTimer() => Stopwatch.Reset();
        }
 
        private readonly Dictionary<Client, OffenseData> clients = new();
 
        private float stopwatchResetTimer,
                      strikesResetTimer;
 
        private const int StopwatchResetInterval = 1,
                          StrikesResetInterval = 60,
                          StrikeThreshold = 6;
 
        private const int MinPacketLimitMultipler = 1;
 
        private static int GetMaxPacketLimit(ServerSettings settings)
            => (int)MathF.Ceiling(
                settings.MaxPacketAmount *
                MathF.Max(
                    settings.TickRate / (float)ServerSettings.DefaultTickRate,
                    MinPacketLimitMultipler)); // Prevent the rate limit multiplier from being less than 1.
 
        public DoSAction Start(Client client) => new DoSAction(client, StartFor, EndFor);
 
        public DoSAction Pause(Client client) => new DoSAction(client, PauseFor, ResumeFor);
 
        private void StartFor(Client client)
        {
            clients.TryAdd(client, new OffenseData());
            clients[client].Stopwatch.Start();
        }
 
        private void EndFor(Client client)
        {
            if (GetData(client) is not { } data) { return; }
 
            data.PacketCount++;
            data.Stopwatch.Stop();
            UpdateOffense(client, data);
        }
 
        // stops the clock but doesn't update offenses
        private void PauseFor(Client client) => GetData(client)?.Stopwatch.Stop();
 
        private void ResumeFor(Client client) => GetData(client)?.Stopwatch.Start();
 
        private void UpdateOffense(Client client, OffenseData data)
        {
            if (GameMain.Server?.ServerSettings is not { } settings) { return; }
 
            // client is sending too many packets, kick them
            if (data.PacketCount > GetMaxPacketLimit(settings) && settings.MaxPacketAmount > ServerSettings.PacketLimitMin)
            {
                AttemptKickClient(client, TextManager.Get("PacketLimitKicked"));
                clients.Remove(client);
                return;
            }
 
            // if the stopwatch has been running for an entire second without the Update() method resetting it (which it does every second) then something is wrong
            if (data.Stopwatch.ElapsedMilliseconds < 100) { return; }
 
            data.Strikes++;
            data.ResetTimer();
 
            GameServer.Log($"{NetworkMember.ClientLogName(client)} is causing the server to slow down.", ServerLog.MessageType.DoSProtection);
 
            // too many strikes, get them out of here
            if (data.Strikes < StrikeThreshold) { return; }
 
            if (settings.EnableDoSProtection)
            {
                AttemptKickClient(client, TextManager.Get("DoSProtectionKicked"));
            }
 
            clients.Remove(client);
 
            static void AttemptKickClient(Client client, LocalizedString reason)
            {
                // ReSharper disable once ConvertToConstant.Local
                bool doesRateLimitAffectClient =
#if DEBUG
                    true; // for testing
#else
                    !RateLimiter.IsExempt(client);
#endif
 
                if (!doesRateLimitAffectClient)
                {
                    return;
                }
 
                GameMain.Server?.KickClient(client, reason.Value);
            }
        }
 
        public void Update(float deltaTime)
        {
            stopwatchResetTimer += deltaTime;
            strikesResetTimer += deltaTime;
 
            // reset the stopwatch every second
            if (stopwatchResetTimer > StopwatchResetInterval)
            {
                stopwatchResetTimer = 0;
                foreach (OffenseData data in clients.Values)
                {
                    data.ResetTimer();
                }
            }
 
            // reset the strikes every minute
            if (strikesResetTimer > StrikesResetInterval)
            {
                strikesResetTimer = 0;
                foreach (var (client, data) in clients)
                {
                    if (GameMain.Server?.ServerSettings is { MaxPacketAmount: > ServerSettings.PacketLimitMin } settings)
                    {
                        if (data.PacketCount > GetMaxPacketLimit(settings) * 0.9f)
                        {
                            GameServer.Log($"{NetworkMember.ClientLogName(client)} is sending a lot of packets and almost got kicked! ({data.PacketCount}).", ServerLog.MessageType.DoSProtection);
                        }
                    }
 
                    data.ResetStrikes();
                }
            }
        }
 
        private OffenseData? GetData(Client client) => clients.TryGetValue(client, out OffenseData? data) ? data : null;
    }
}