rate limitier for the bot (#8)

Cryptographic-API-Services · Aug 8, 2024 · e3ba45a · e3ba45a
1 parent 7bab4ba
commit e3ba45a
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 9 deletions.
diff --git a/handlers/message-create.ts b/handlers/message-create.ts
@@ -17,14 +17,17 @@ import { ChatMessageRepository } from "../repositories/chat-message-repository.t
 import { ReloadDocumentationCommand } from "./reload-documentation.command.ts";
 import { PGVectorStore } from "npm:@langchain/community/vectorstores/pgvector";
 import { DocumentationVectorStoreRepository } from "../repositories/documentation-vector-store-repository.ts";
+import { MessageRateLimiter } from "./message-rate-limiter.ts";
 
 export default class MessageCreateHandler {
   bot: Bot;
   message: Message;
+  rateLimiter: MessageRateLimiter;
 
-  constructor(bot: Bot, message: Message) {
+  constructor(bot: Bot, message: Message, rateLimiter: MessageRateLimiter) {
     this.bot = bot;
     this.message = message;
+    this.rateLimiter = rateLimiter;
   }
 
   public async handleMessage(): Promise<void> {
@@ -34,13 +37,17 @@ export default class MessageCreateHandler {
         this.message.channelId,
         this.message.id,
       );
-      const sliceMessage = gotMessage.content.slice(0, 20).toLowerCase();
-      if (sliceMessage.includes("hey bot:")) {
-        await this.performHeyBotQuery(gotMessage);
-      } else if (sliceMessage.includes("documentation:")) {
-        await this.performDocumentationQuery(gotMessage);
-      } else if (sliceMessage.includes("replace docs:")) {
-        this.replaceDocumentation(gotMessage);
+      if (!gotMessage.isFromBot && this.rateLimiter.checkRateLimitForUser(gotMessage.authorId)) {
+        await sendMessage(this.bot, gotMessage.channelId, {content: `<@${gotMessage.authorId}> ` + `You have messaged the bot to many times in the last thirty seconds, please wait.`});
+      } else {
+        const sliceMessage = gotMessage.content.slice(0, 20).toLowerCase();
+        if (sliceMessage.includes("hey bot:")) {
+          await this.performHeyBotQuery(gotMessage);
+        } else if (sliceMessage.includes("documentation:")) {
+          await this.performDocumentationQuery(gotMessage);
+        } else if (sliceMessage.includes("replace docs:")) {
+          this.replaceDocumentation(gotMessage);
+        }
       }
     } catch (error) {
       console.error(error);

diff --git a/handlers/message-rate-limiter.ts b/handlers/message-rate-limiter.ts
@@ -0,0 +1,29 @@
+export class MessageRateLimiter {
+    private rateLimitMap: Map<bigint, number>;
+    private interval: number;
+
+    constructor() {
+        this.rateLimitMap = new Map();
+        this.interval = setInterval(() => {
+            for (const key of this.rateLimitMap.keys()) {
+                this.rateLimitMap.delete(key);
+            }
+        }, 30000);
+    }
+
+
+    public checkRateLimitForUser(authorId: bigint): boolean {
+        let hasRateLimitBeenMeet = false;
+        if (!this.rateLimitMap.has(authorId)) {
+            this.rateLimitMap.set(authorId, 1);
+        } else {
+            let limit = this.rateLimitMap.get(authorId);
+            if (limit && limit >= 5) {
+                hasRateLimitBeenMeet = true;
+            }
+            limit = limit + 1;
+            this.rateLimitMap.set(authorId, limit);
+        }
+        return hasRateLimitBeenMeet;
+    }
+}
diff --git a/index.ts b/index.ts
@@ -6,6 +6,7 @@ import {
   startBot
 } from "https://deno.land/x/[email protected]/mod.ts";
 import MessageCreateHandler from "./handlers/message-create.ts";
+import { MessageRateLimiter } from "./handlers/message-rate-limiter.ts";
 
 const bot = createBot({
   token: Deno.env.get("DISCORD_TOKEN"),
@@ -17,9 +18,11 @@ const bot = createBot({
   },
 });
 
+const rateLimiter = new MessageRateLimiter();
+
 bot.events.messageCreate = async function (b: Bot, message: Message) {
   try {
-    const handler = new MessageCreateHandler(b, message);
+    const handler = new MessageCreateHandler(b, message, rateLimiter);
     await handler.handleMessage();
   } catch(error) {
     console.error(error);