From 380f576f93517d87ea4b5ebedb85202c64919b6c Mon Sep 17 00:00:00 2001 From: divyam234 Date: Thu, 24 Aug 2023 22:24:42 +0530 Subject: [PATCH] rate limit to check upload speed doesn't go too high and fail --- README.md | 1 + utils/tgclient.go | 13 +++++++++++++ 2 files changed, 14 insertions(+) diff --git a/README.md b/README.md index 32d8b9e..faa3398 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,7 @@ HTTPS=false COOKIE_SAME_SITE=true JWT_SECRET=abc DATABASE_URL=abc +RATE_LIMIT=true TG_CLIENT_DEVICE_MODEL="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36 Edg/115.0.1901.203" # Any valid browser user agent here TG_CLIENT_SYSTEM_VERSION=Win32 TG_CLIENT_APP_VERSION=2.1.9 K diff --git a/utils/tgclient.go b/utils/tgclient.go index 0a2604a..7447cef 100644 --- a/utils/tgclient.go +++ b/utils/tgclient.go @@ -15,11 +15,13 @@ import ( "github.com/gin-gonic/gin" "github.com/gotd/contrib/bg" "github.com/gotd/contrib/middleware/floodwait" + "github.com/gotd/contrib/middleware/ratelimit" tdclock "github.com/gotd/td/clock" "github.com/gotd/td/session" "github.com/gotd/td/telegram" "github.com/pkg/errors" "go.uber.org/zap" + "golang.org/x/time/rate" ) var clients map[int64]*telegram.Client @@ -54,8 +56,13 @@ func GetBotClient(clientName string) *telegram.Client { sessionStorage := &telegram.FileSessionStorage{ Path: filepath.Join("sessions", clientName+".json"), } + middlewares := []telegram.Middleware{floodwait.NewSimpleWaiter()} + if config.RateLimit { + middlewares = append(middlewares, ratelimit.New(rate.Every(time.Millisecond*100), 5)) + } + options := telegram.Options{ SessionStorage: sessionStorage, Middlewares: middlewares, @@ -88,7 +95,13 @@ func GetAuthClient(ctx context.Context, sessionStr string, userId int64) (*teleg if err := loader.Save(ctx, data); err != nil { return nil, err } + middlewares := []telegram.Middleware{floodwait.NewSimpleWaiter()} + + 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: middlewares,