From aeb04343561f6f98ca88aea7de5e4abba7215cd2 Mon Sep 17 00:00:00 2001
From: divyam234 <divyam@mnnit.ac.in>
Date: Sun, 20 Aug 2023 02:48:04 +0530
Subject: [PATCH] optional rate limit

---
 README.md         |  3 ++-
 utils/config.go   |  1 +
 utils/tgclient.go | 38 +++++++++++++++++++++-----------------
 3 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 9b5309a..c810444 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,7 @@ An example of `.env` file:
 APP_ID=1234
 APP_HASH=abc
 CHANNEL_ID=1234
+RATE_LIMIT=true
 HTTPS=false
 COOKIE_SAME_SITE=true
 JWT_SECRET=abc
@@ -125,7 +126,7 @@ Before running the bot, you will need to set up the following mandatory variable
 
 ### Optional Vars
 In addition to the mandatory variables, you can also set the following optional variables:
-
+- `RATE_LIMIT` : Rate Limit Calls to prevent flood errors
 - `HTTPS` : Only needed when frontend is deployed on vercel.
 
 - `COOKIE_SAME_SITE` : Only needed when frontend is deployed on vercel.
diff --git a/utils/config.go b/utils/config.go
index 7ac88d2..13b5d5b 100644
--- a/utils/config.go
+++ b/utils/config.go
@@ -15,6 +15,7 @@ type Config struct {
 	Https                  bool   `envconfig:"HTTPS" default:"false"`
 	CookieSameSite         bool   `envconfig:"COOKIE_SAME_SITE" default:"true"`
 	DatabaseUrl            string `envconfig:"DATABASE_URL" required:"true"`
+	RateLimit              bool   `envconfig:"RATE_LIMIT" default:"true"`
 	TgClientDeviceModel    string `envconfig:"TG_CLIENT_DEVICE_MODEL" required:"true"`
 	TgClientSystemVersion  string `envconfig:"TG_CLIENT_SYSTEM_VERSION" default:"Win32"`
 	TgClientAppVersion     string `envconfig:"TG_CLIENT_APP_VERSION" default:"2.1.9 K"`
diff --git a/utils/tgclient.go b/utils/tgclient.go
index 28bfdf5..d1f505b 100644
--- a/utils/tgclient.go
+++ b/utils/tgclient.go
@@ -43,13 +43,15 @@ func getBotClient(appID int, appHash, clientName, sessionDir string) *telegram.C
 	sessionStorage := &telegram.FileSessionStorage{
 		Path: filepath.Join(sessionDir, clientName+".json"),
 	}
+	middlewares := []telegram.Middleware{}
+	if config.RateLimit {
+		middlewares = append(middlewares, ratelimit.New(rate.Every(time.Millisecond*100), 5))
+	}
 	options := telegram.Options{
 		SessionStorage: sessionStorage,
-		Middlewares: []telegram.Middleware{
-			ratelimit.New(rate.Every(time.Millisecond*100), 5),
-		},
-		Device:    getDeviceConfig(),
-		NoUpdates: true,
+		Middlewares:    middlewares,
+		Device:         getDeviceConfig(),
+		NoUpdates:      true,
 	}
 
 	client := telegram.NewClient(appID, appHash, options)
@@ -139,14 +141,15 @@ func GetAuthClient(sessionStr string, userId int) (*Client, bg.StopFunc, error)
 	if err := loader.Save(ctx, data); err != nil {
 		return nil, nil, err
 	}
-
+	middlewares := []telegram.Middleware{}
+	if config.RateLimit {
+		middlewares = append(middlewares, ratelimit.New(rate.Every(time.Millisecond*100), 5))
+	}
 	client := telegram.NewClient(config.AppId, config.AppHash, telegram.Options{
 		SessionStorage: storage,
-		Middlewares: []telegram.Middleware{
-			ratelimit.New(rate.Every(time.Millisecond*100), 5),
-		},
-		Device:    getDeviceConfig(),
-		NoUpdates: true,
+		Middlewares:    middlewares,
+		Device:         getDeviceConfig(),
+		NoUpdates:      true,
 	})
 
 	stop, err := bg.Connect(client)
@@ -181,14 +184,15 @@ func GetBotClient() *Client {
 }
 
 func GetNonAuthClient(handler telegram.UpdateHandler, storage telegram.SessionStorage) (*telegram.Client, bg.StopFunc, error) {
-
+	middlewares := []telegram.Middleware{}
+	if config.RateLimit {
+		middlewares = append(middlewares, ratelimit.New(rate.Every(time.Millisecond*100), 5))
+	}
 	client := telegram.NewClient(config.AppId, config.AppHash, telegram.Options{
 		SessionStorage: storage,
-		Middlewares: []telegram.Middleware{
-			ratelimit.New(rate.Every(time.Millisecond*100), 5),
-		},
-		Device:        getDeviceConfig(),
-		UpdateHandler: handler,
+		Middlewares:    middlewares,
+		Device:         getDeviceConfig(),
+		UpdateHandler:  handler,
 	})
 
 	stop, err := bg.Connect(client)