mirror of
https://github.com/ChuckNorrison/LightningTipBot.git
synced 2026-08-13 12:33:14 +02:00
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>
This commit is contained in:
gohumble
GitHub
parent
25d9b8af3b
commit
aaa8a3cc19
5 changed files with 103 additions and 0 deletions
2
go.mod
2
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
|
||||
|
|
|
|||
4
go.sum
4
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=
|
||||
|
|
|
|||
87
internal/rate/limiter.go
Normal file
87
internal/rate/limiter.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue