From a2b5413e85dc6d8cbef634bbdfaf5c092189c0b2 Mon Sep 17 00:00:00 2001
From: Morb <14136326+Morb0@users.noreply.github.com>
Date: Sat, 24 Aug 2024 02:00:44 +0300
Subject: [PATCH] Add TTS rate limit (#2515)

---
 .../Corvax/TTS/TTSSystem.RateLimit.cs         | 35 +++++++++++++++++++
 Content.Server/Corvax/TTS/TTSSystem.cs        |  6 ++++
 Content.Shared/Corvax/CCCVars/CCCVars.cs      | 13 +++++++
 Resources/Locale/ru-RU/corvax/tts/tts-ui.ftl  |  1 +
 4 files changed, 55 insertions(+)
 create mode 100644 Content.Server/Corvax/TTS/TTSSystem.RateLimit.cs

diff --git a/Content.Server/Corvax/TTS/TTSSystem.RateLimit.cs b/Content.Server/Corvax/TTS/TTSSystem.RateLimit.cs
new file mode 100644
index 00000000000..261270dde08
--- /dev/null
+++ b/Content.Server/Corvax/TTS/TTSSystem.RateLimit.cs
@@ -0,0 +1,35 @@
+using Content.Server.Chat.Managers;
+using Content.Server.Players.RateLimiting;
+using Content.Shared.Corvax.CCCVars;
+using Robust.Shared.Player;
+
+namespace Content.Server.Corvax.TTS;
+
+public sealed partial class TTSSystem
+{
+    [Dependency] private readonly PlayerRateLimitManager _rateLimitManager = default!;
+    [Dependency] private readonly IChatManager _chat = default!;
+
+    private const string RateLimitKey = "TTS";
+
+    private void RegisterRateLimits()
+    {
+        _rateLimitManager.Register(RateLimitKey,
+            new RateLimitRegistration
+            {
+                CVarLimitPeriodLength = CCCVars.TTSRateLimitPeriod,
+                CVarLimitCount = CCCVars.TTSRateLimitCount,
+                PlayerLimitedAction = RateLimitPlayerLimited,
+            });
+    }
+
+    private void RateLimitPlayerLimited(ICommonSession player)
+    {
+        _chat.DispatchServerMessage(player, Loc.GetString("tts-rate-limited"), suppressLog: true);
+    }
+
+    private RateLimitStatus HandleRateLimit(ICommonSession player)
+    {
+        return _rateLimitManager.CountAction(player, RateLimitKey);
+    }
+}
diff --git a/Content.Server/Corvax/TTS/TTSSystem.cs b/Content.Server/Corvax/TTS/TTSSystem.cs
index 81caf95b099..00390fd2fc6 100644
--- a/Content.Server/Corvax/TTS/TTSSystem.cs
+++ b/Content.Server/Corvax/TTS/TTSSystem.cs
@@ -1,5 +1,6 @@
 using System.Threading.Tasks;
 using Content.Server.Chat.Systems;
+using Content.Server.Players.RateLimiting;
 using Content.Shared.Corvax.CCCVars;
 using Content.Shared.Corvax.TTS;
 using Content.Shared.GameTicking;
@@ -48,6 +49,8 @@ public override void Initialize()
         SubscribeLocalEvent<RoundRestartCleanupEvent>(OnRoundRestartCleanup);
 
         SubscribeNetworkEvent<RequestPreviewTTSEvent>(OnRequestPreviewTTS);
+
+        RegisterRateLimits();
     }
 
     private void OnRoundRestartCleanup(RoundRestartCleanupEvent ev)
@@ -61,6 +64,9 @@ private async void OnRequestPreviewTTS(RequestPreviewTTSEvent ev, EntitySessionE
             !_prototypeManager.TryIndex<TTSVoicePrototype>(ev.VoiceId, out var protoVoice))
             return;
 
+        if (HandleRateLimit(args.SenderSession) != RateLimitStatus.Allowed)
+            return;
+
         var previewText = _rng.Pick(_sampleText);
         var soundData = await GenerateTTS(previewText, protoVoice.Speaker);
         if (soundData is null)
diff --git a/Content.Shared/Corvax/CCCVars/CCCVars.cs b/Content.Shared/Corvax/CCCVars/CCCVars.cs
index 4344107d001..f0108051b9c 100644
--- a/Content.Shared/Corvax/CCCVars/CCCVars.cs
+++ b/Content.Shared/Corvax/CCCVars/CCCVars.cs
@@ -55,6 +55,19 @@ public sealed class CCCVars
     public static readonly CVarDef<int> TTSMaxCache =
         CVarDef.Create("tts.max_cache", 250, CVar.SERVERONLY | CVar.ARCHIVE);
 
+    /// <summary>
+    /// Tts rate limit values are accounted in periods of this size (seconds).
+    /// After the period has passed, the count resets.
+    /// </summary>
+    public static readonly CVarDef<int> TTSRateLimitPeriod =
+        CVarDef.Create("tts.rate_limit_period", 2, CVar.SERVERONLY);
+
+    /// <summary>
+    /// How many tts preview messages are allowed in a single rate limit period.
+    /// </summary>
+    public static readonly CVarDef<int> TTSRateLimitCount =
+        CVarDef.Create("tts.rate_limit_count", 3, CVar.SERVERONLY);
+
     /*
      * Peaceful Round End
      */
diff --git a/Resources/Locale/ru-RU/corvax/tts/tts-ui.ftl b/Resources/Locale/ru-RU/corvax/tts/tts-ui.ftl
index 0a354c3bfc7..6471cf5b91e 100644
--- a/Resources/Locale/ru-RU/corvax/tts/tts-ui.ftl
+++ b/Resources/Locale/ru-RU/corvax/tts/tts-ui.ftl
@@ -2,3 +2,4 @@ ui-options-tts-volume = Громкость TTS:
 credits-window-tts-title = Функция TTS (Text-To-Speech)
 humanoid-profile-editor-voice-label = Голос:
 humanoid-profile-editor-voice-play = ▶
+tts-rate-limited = Вы генерируете TTS слишком быстро!