From aaa8a3cc197c8af800cd93686adf5f34e2da40ca Mon Sep 17 00:00:00 2001 From: gohumble <55599638+gohumble@users.noreply.github.com> Date: Mon, 6 Dec 2021 22:22:02 +0100 Subject: [PATCH] WIP: rate limiter per chat (#86) * add rate limiter per chat * differentiate between global and chat rate limiting * remove diffs * resolve conflicts * resolve conflict * refactor limiter * rename * rename * logs * remove logs * reset buckets * new rate limiter * fix comments * separate id and global limiter * rename stuff * reset bucket size Co-authored-by: callebtc <93376500+callebtc@users.noreply.github.com> --- go.mod | 2 + go.sum | 4 ++ internal/rate/limiter.go | 87 +++++++++++++++++++++++++++++++++++ internal/telegram/bot.go | 4 ++ internal/telegram/telegram.go | 6 +++ 5 files changed, 103 insertions(+) create mode 100644 internal/rate/limiter.go diff --git a/go.mod b/go.mod index d5a69f7..de57f04 100644 --- a/go.mod +++ b/go.mod @@ -13,11 +13,13 @@ require ( github.com/makiuchi-d/gozxing v0.0.2 github.com/nicksnyder/go-i18n/v2 v2.1.2 github.com/patrickmn/go-cache v2.1.0+incompatible + github.com/sethvargo/go-limiter v0.7.2 github.com/sirupsen/logrus v1.6.0 github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e github.com/tidwall/buntdb v1.2.7 github.com/tidwall/gjson v1.10.2 golang.org/x/text v0.3.5 + golang.org/x/time v0.0.0-20191024005414-555d28b269f0 gopkg.in/lightningtipbot/telebot.v2 v2.4.2-0.20211201074627-babf9f2cc955 gorm.io/driver/sqlite v1.1.4 gorm.io/gorm v1.21.12 diff --git a/go.sum b/go.sum index 0929553..6bf2307 100644 --- a/go.sum +++ b/go.sum @@ -453,6 +453,8 @@ github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQD github.com/ryanuber/columnize v0.0.0-20160712163229-9b3edd62028f/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts= github.com/samuel/go-zookeeper v0.0.0-20190923202752-2cc03de413da/go.mod h1:gi+0XIa01GRL2eRQVjQkKGqKF3SF9vZR/HnPullcV2E= github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc= +github.com/sethvargo/go-limiter v0.7.2 h1:FgC4N7RMpV5gMrUdda15FaFTkQ/L4fEqM7seXMs4oO8= +github.com/sethvargo/go-limiter v0.7.2/go.mod h1:C0kbSFbiriE5k2FFOe18M1YZbAR2Fiwf72uGu0CXCcU= github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc= github.com/sirupsen/logrus v1.2.0 h1:juTguoYk5qI21pwyTXY3B3Y5cOTH3ZUyZCg1v/mihuo= github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= @@ -649,7 +651,9 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.5 h1:i6eZZ+zk0SOf0xgBpEpPD18qWcJda6q1sxt3S0kzyUQ= golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2 h1:+DCIGbF/swA92ohVg0//6X2IVY3KZs6p9mix0ziNYJM= golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180828015842-6cd1fcedba52/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/internal/rate/limiter.go b/internal/rate/limiter.go new file mode 100644 index 0000000..4184778 --- /dev/null +++ b/internal/rate/limiter.go @@ -0,0 +1,87 @@ +package rate + +import ( + "context" + "golang.org/x/time/rate" + tb "gopkg.in/lightningtipbot/telebot.v2" + "strconv" + "sync" +) + +// Limiter +type Limiter struct { + keys map[string]*rate.Limiter + mu *sync.RWMutex + r rate.Limit + b int +} + +var idLimiter *Limiter +var globalLimiter *rate.Limiter + +// NewLimiter creates both chat and global rate limiters. +func Start() { + idLimiter = newIdRateLimiter(rate.Limit(0.3), 20) + globalLimiter = rate.NewLimiter(rate.Limit(30), 30) +} + +// NewRateLimiter . +func newIdRateLimiter(r rate.Limit, b int) *Limiter { + i := &Limiter{ + keys: make(map[string]*rate.Limiter), + mu: &sync.RWMutex{}, + r: r, + b: b, + } + + return i +} + +func CheckLimit(to interface{}) { + globalLimiter.Wait(context.Background()) + var id string + switch to.(type) { + case *tb.Chat: + id = strconv.FormatInt(to.(*tb.Chat).ID, 10) + case *tb.User: + id = strconv.FormatInt(to.(*tb.User).ID, 10) + case tb.Recipient: + id = to.(tb.Recipient).Recipient() + case *tb.Message: + if to.(*tb.Message).Chat != nil { + id = strconv.FormatInt(to.(*tb.Message).Chat.ID, 10) + } + } + if len(id) > 0 { + idLimiter.GetLimiter(id).Wait(context.Background()) + } +} + +// Add creates a new rate limiter and adds it to the keys map, +// using the key +func (i *Limiter) Add(key string) *rate.Limiter { + i.mu.Lock() + defer i.mu.Unlock() + + limiter := rate.NewLimiter(i.r, i.b) + + i.keys[key] = limiter + + return limiter +} + +// GetLimiter returns the rate limiter for the provided key if it exists. +// Otherwise, calls Add to add key address to the map +func (i *Limiter) GetLimiter(key string) *rate.Limiter { + i.mu.Lock() + limiter, exists := i.keys[key] + + if !exists { + i.mu.Unlock() + return i.Add(key) + } + + i.mu.Unlock() + + return limiter +} diff --git a/internal/telegram/bot.go b/internal/telegram/bot.go index efff987..a5b9876 100644 --- a/internal/telegram/bot.go +++ b/internal/telegram/bot.go @@ -5,6 +5,8 @@ import ( "sync" "time" + limiter "github.com/LightningTipBot/LightningTipBot/internal/rate" + "github.com/eko/gocache/store" "github.com/LightningTipBot/LightningTipBot/internal" @@ -23,6 +25,7 @@ type TipBot struct { logger *gorm.DB Telegram *telebot.Bot Client *lnbits.Client + limiter map[string]limiter.Limiter Cache } type Cache struct { @@ -40,6 +43,7 @@ func NewBot() TipBot { gocacheStore := store.NewGoCache(gocacheClient, nil) // create sqlite databases db, txLogger := AutoMigration() + limiter.Start() return TipBot{ Database: db, Client: lnbits.NewClient(internal.Configuration.Lnbits.AdminKey, internal.Configuration.Lnbits.Url), diff --git a/internal/telegram/telegram.go b/internal/telegram/telegram.go index 210b81c..6a8ecdf 100644 --- a/internal/telegram/telegram.go +++ b/internal/telegram/telegram.go @@ -2,6 +2,7 @@ package telegram import ( "fmt" + "github.com/LightningTipBot/LightningTipBot/internal/rate" "github.com/eko/gocache/store" log "github.com/sirupsen/logrus" tb "gopkg.in/lightningtipbot/telebot.v2" @@ -9,6 +10,7 @@ import ( ) func (bot TipBot) tryForwardMessage(to tb.Recipient, what tb.Editable, options ...interface{}) (msg *tb.Message) { + rate.CheckLimit(to) msg, err := bot.Telegram.Forward(to, what, options...) if err != nil { log.Warnln(err.Error()) @@ -16,6 +18,7 @@ func (bot TipBot) tryForwardMessage(to tb.Recipient, what tb.Editable, options . return } func (bot TipBot) trySendMessage(to tb.Recipient, what interface{}, options ...interface{}) (msg *tb.Message) { + rate.CheckLimit(to) msg, err := bot.Telegram.Send(to, what, options...) if err != nil { log.Warnln(err.Error()) @@ -24,6 +27,7 @@ func (bot TipBot) trySendMessage(to tb.Recipient, what interface{}, options ...i } func (bot TipBot) tryReplyMessage(to *tb.Message, what interface{}, options ...interface{}) (msg *tb.Message) { + rate.CheckLimit(to) msg, err := bot.Telegram.Reply(to, what, options...) if err != nil { log.Warnln(err.Error()) @@ -35,6 +39,7 @@ func (bot TipBot) tryEditMessage(to tb.Editable, what interface{}, options ...in if !allowedToPerformAction(bot, to, isAdminAndCanEdit) { return } + rate.CheckLimit(to) var err error msg, err = bot.Telegram.Edit(to, what, options...) if err != nil { @@ -48,6 +53,7 @@ func (bot TipBot) tryDeleteMessage(msg tb.Editable) { if !allowedToPerformAction(bot, msg, isAdminAndCanDelete) { return } + rate.CheckLimit(msg) err := bot.Telegram.Delete(msg) if err != nil { log.Warnln(err.Error())