From 8274ab09805fd07b08d5053203d7b29c2c004ef3 Mon Sep 17 00:00:00 2001
From: Sphericalkat <me@kat.bio>
Date: Mon, 30 Sep 2024 10:39:07 +0530
Subject: [PATCH] feat: add rate limiter middleware

Signed-off-by: Sphericalkat <me@kat.bio>
---
 README.md         |  2 ++
 bot.js            | 10 +++++++++-
 package-lock.json | 11 +++++++++++
 package.json      |  1 +
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index df66af6..8630358 100644
--- a/README.md
+++ b/README.md
@@ -13,4 +13,6 @@ roleplay gifs. Once you have an instance of Ugoki (and optionally
 - `export BOT_API_KEY="your-token-for-bot"`
 - `export FEEDBACK_ID="chat-id-where-feedback-is-forwarded-to"`
 - `export UGOKI_ROOT="https://root.of.ugoki.api/server/"`
+- `export RATE_TIMEFRAME=5000 # rate limit time interval in milliseconds`
+- `export RATE_LIMIT=5 # number of requests allowed in the timeframe`
 - `npm start`
diff --git a/bot.js b/bot.js
index d67a893..6682c04 100644
--- a/bot.js
+++ b/bot.js
@@ -1,12 +1,20 @@
 const { Telegraf } = require("telegraf");
-const { BOT_API_KEY, FEEDBACK_ID, UGOKI_ROOT } = process.env;
+const { BOT_API_KEY, FEEDBACK_ID, UGOKI_ROOT, RATE_TIMEFRAME, RATE_LIMIT } = process.env;
 const fs = require("fs").promises;
 const commands = require("./commands");
 const axios = require("axios");
 const roleplay = require("./static/roleplay.json");
+const { limit } = require("@grammyjs/ratelimiter");
 
 const bot = new Telegraf(BOT_API_KEY);
 bot.catch((err) => console.log(err));
+bot.use(limit({
+    // default config: 5 messages per 5 seconds
+    timeFrame: RATE_TIMEFRAME ?? 5000,
+    limit: RATE_LIMIT ?? 5,
+
+    onLimitExceeded: (ctx, next) => ctx.reply('Too many requests!'),
+}))
 
 const data = [
 	"questions",
diff --git a/package-lock.json b/package-lock.json
index f9606be..04dfc82 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -9,12 +9,18 @@
       "version": "2.0.0",
       "license": "GPL-3.0",
       "dependencies": {
+        "@grammyjs/ratelimiter": "^1.2.0",
         "axios": "^0.21.0",
         "fluent-ffmpeg": "^2.1.2",
         "form-data": "^4.0.0",
         "telegraf": "^4.15.3"
       }
     },
+    "node_modules/@grammyjs/ratelimiter": {
+      "version": "1.2.0",
+      "resolved": "https://registry.npmjs.org/@grammyjs/ratelimiter/-/ratelimiter-1.2.0.tgz",
+      "integrity": "sha512-xBkH/ATJsuv5JgVYX9yQM9DNg75Qqjw+gh82lVsBn4j+d0DkxxC+kuy6WFoB96Cb6oifQfaBJL8CTikdYG4v0A=="
+    },
     "node_modules/@telegraf/types": {
       "version": "6.9.1",
       "resolved": "https://registry.npmjs.org/@telegraf/types/-/types-6.9.1.tgz",
@@ -288,6 +294,11 @@
     }
   },
   "dependencies": {
+    "@grammyjs/ratelimiter": {
+      "version": "1.2.0",
+      "resolved": "https://registry.npmjs.org/@grammyjs/ratelimiter/-/ratelimiter-1.2.0.tgz",
+      "integrity": "sha512-xBkH/ATJsuv5JgVYX9yQM9DNg75Qqjw+gh82lVsBn4j+d0DkxxC+kuy6WFoB96Cb6oifQfaBJL8CTikdYG4v0A=="
+    },
     "@telegraf/types": {
       "version": "6.9.1",
       "resolved": "https://registry.npmjs.org/@telegraf/types/-/types-6.9.1.tgz",
diff --git a/package.json b/package.json
index dba3639..6130bb1 100644
--- a/package.json
+++ b/package.json
@@ -22,6 +22,7 @@
   },
   "homepage": "https://gitlab.com/ceda_ei/Quadnite-Bot#readme",
   "dependencies": {
+    "@grammyjs/ratelimiter": "^1.2.0",
     "axios": "^0.21.0",
     "fluent-ffmpeg": "^2.1.2",
     "form-data": "^4.0.0",