From eb881a6fe562eb5ddfda4d2894324cbc46beeeae Mon Sep 17 00:00:00 2001
From: Nikita <nikitadiscore@gmail.com>
Date: Mon, 27 Nov 2023 20:56:33 +0300
Subject: [PATCH] rate limiter: add (#487)

Co-authored-by: demget <30910794+demget@users.noreply.github.com>
---
 README.md | 16 ++++++++++++++++
 api.go    | 12 ++++++++++++
 bot.go    | 40 +++++++++++++++++++++++++++++++++++++++-
 go.mod    |  5 +++++
 4 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 060a70d..90e6ff9 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,7 @@ go get -u gopkg.in/telebot.v3
 	- [Editable](#editable)
 	- [Keyboards](#keyboards)
 	- [Inline mode](#inline-mode)
+	- [Rate limiter](#rate-limiter)
 * [Contributing](#contributing)
 * [Donate](#donate)
 * [License](#license)
@@ -460,6 +461,21 @@ b.Handle(tele.OnQuery, func(c tele.Context) error {
 })
 ```
 
+## Rate limiter
+In order not to catch the anti-flood, an implementation has been made that allows you to limit the number of outgoing requests by n times per second.
+```go
+pref := tele.Settings{
+	Token:      "123456:token",
+	// PerSeconds limits the number of requests per second executed
+	// by the client for Raw function. Default per in seconds is -1.
+	// To enable the queue, set the value greater than zero.
+	PerSeconds: 30,
+	// PerBufferSize sets the size of the queue that is waiting for the signal to be sent.
+	// By default, the buffer is infinite and has a value equal to zero.
+	PerBufferSize: 0,
+}
+```
+
 There's not much to talk about really. It also supports some form of authentication
 through deep-linking. For that, use fields `SwitchPMText` and `SwitchPMParameter`
 of `QueryResponse`.
diff --git a/api.go b/api.go
index 6ed4eb1..e482974 100644
--- a/api.go
+++ b/api.go
@@ -13,6 +13,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -20,6 +21,17 @@ import (
 // It also handles API errors, so you only need to unwrap
 // result field from json data.
 func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
+	if b.ratelimit {
+		c := sync.NewCond(&sync.Mutex{})
+		c.L.Lock()
+		b.raws <- c
+		c.Wait()
+		c.L.Unlock()
+	}
+	return b.raw(method, payload)
+}
+
+func (b *Bot) raw(method string, payload interface{}) ([]byte, error) {
 	url := b.URL + "/bot" + b.Token + "/" + method
 
 	var buf bytes.Buffer
diff --git a/bot.go b/bot.go
index e3a9a41..50e92a7 100644
--- a/bot.go
+++ b/bot.go
@@ -9,8 +9,11 @@ import (
 	"os"
 	"regexp"
 	"strconv"
-	"strings"
 	"time"
+	"strings"
+	"sync"
+
+	"go.uber.org/ratelimit"
 )
 
 // NewBot does try to build a Bot with token `token`, which
@@ -19,6 +22,9 @@ func NewBot(pref Settings) (*Bot, error) {
 	if pref.Updates == 0 {
 		pref.Updates = 100
 	}
+	if pref.PerSeconds == 0 {
+		pref.PerSeconds = -1
+	}
 
 	client := pref.Client
 	if client == nil {
@@ -49,8 +55,29 @@ func NewBot(pref Settings) (*Bot, error) {
 		verbose:     pref.Verbose,
 		parseMode:   pref.ParseMode,
 		client:      client,
+
+		ratelimit: false,
 	}
 
+	wait := make(chan struct{})
+	if pref.PerSeconds != -1 {
+		bot.ratelimit = true
+		go func(b *Bot) {
+			rl := ratelimit.New(pref.PerSeconds)
+			bot.raws = make(chan *sync.Cond, pref.PerBufferSize)
+			wait <- struct{}{}
+			for raw := range bot.raws {
+				rl.Take()
+				raw.Broadcast()
+			}
+		}(bot)
+	} else {
+		go func() {
+			wait <- struct{}{}
+		}()
+	}
+	<-wait
+
 	if pref.Offline {
 		bot.Me = &User{}
 	} else {
@@ -82,6 +109,8 @@ type Bot struct {
 	stop        chan chan struct{}
 	client      *http.Client
 	stopClient  chan struct{}
+	raws        chan *sync.Cond
+	ratelimit   bool
 }
 
 // Settings represents a utility struct for passing certain
@@ -119,6 +148,15 @@ type Settings struct {
 
 	// Offline allows to create a bot without network for testing purposes.
 	Offline bool
+
+	// PerSeconds limits the number of requests per second executed
+	// by the client for Raw function. Default per in seconds is -1.
+	// To enable the queue, set the value greater than zero.
+	PerSeconds int
+
+	// PerBufferSize sets the size of the queue that is waiting for the signal to be sent.
+	// By default, the buffer is infinite and has a value equal to zero.
+	PerBufferSize int
 }
 
 var defaultOnError = func(err error, c Context) {
diff --git a/go.mod b/go.mod
index 10ccd31..d31d2d7 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,11 @@ go 1.13
 
 require (
 	github.com/goccy/go-yaml v1.9.5
+	github.com/mattn/go-colorable v0.1.12 // indirect
+	github.com/spf13/cast v1.3.1
+	go.uber.org/ratelimit v0.2.0 // indirect
+	golang.org/x/sys v0.0.0-20220114195835-da31bd327af9 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 	github.com/spf13/viper v1.13.0
 	github.com/stretchr/testify v1.8.0
 )