From ea7afd951e50b2cc4975553c5c84100a53dcc576 Mon Sep 17 00:00:00 2001
From: balki <3070606-balki@users.noreply.gitlab.com>
Date: Mon, 2 May 2022 13:24:29 -0400
Subject: [PATCH] Add rate limiter

---
 go.mod               |  1 +
 go.sum               |  2 ++
 telegram/telegram.go | 25 ++++++++++++++++++++++---
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 14e436c..ec397c6 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/stretchr/testify v1.7.1
 	go.uber.org/zap v1.21.0
 	golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4
+	golang.org/x/time v0.0.0-20220411224347-583f2d630306
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
 )
 
diff --git a/go.sum b/go.sum
index ebee1af..48b3113 100644
--- a/go.sum
+++ b/go.sum
@@ -47,6 +47,8 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/time v0.0.0-20220411224347-583f2d630306 h1:+gHMid33q6pen7kv9xvT+JRinntgeXO2AeZVd0AWD3w=
+golang.org/x/time v0.0.0-20220411224347-583f2d630306/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
diff --git a/telegram/telegram.go b/telegram/telegram.go
index 42b7184..7d3f4aa 100644
--- a/telegram/telegram.go
+++ b/telegram/telegram.go
@@ -2,14 +2,17 @@ package telegram
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
+	"time"
 
 	"go.balki.me/tss/log"
+	"golang.org/x/time/rate"
 )
 
 func main() {
@@ -21,8 +24,10 @@ type TelegramSender interface {
 }
 
 type telegramSender struct {
-	client    *http.Client
-	authToken string
+	client            *http.Client
+	authToken         string
+	rateLimiterPerMin *rate.Limiter
+	rateLimiterPerSec *rate.Limiter
 }
 
 func (ts *telegramSender) SendLink(link string, channel string, rhash string) error {
@@ -42,6 +47,17 @@ func (ts *telegramSender) SendLink(link string, channel string, rhash string) er
 	}
 
 	apiUrl := fmt.Sprintf("https://api.telegram.org/bot%s/sendMessage", ts.authToken)
+
+	err = ts.rateLimiterPerMin.Wait(context.TODO())
+	if err != nil {
+		return err
+	}
+
+	err = ts.rateLimiterPerSec.Wait(context.TODO())
+	if err != nil {
+		return err
+	}
+
 	res, err := ts.client.Post(apiUrl, "application/json", bytes.NewReader(data))
 	if err != nil {
 		return err
@@ -60,10 +76,13 @@ func (ts *telegramSender) SendLink(link string, channel string, rhash string) er
 }
 
 func NewTelegramSender(transport http.RoundTripper, authToken string) TelegramSender {
-	//TODO: Rate limit
 	return &telegramSender{
 		client:    &http.Client{Transport: transport},
 		authToken: authToken,
+		// 20 requests per min with some buffer
+		rateLimiterPerMin: rate.NewLimiter(rate.Every(65*time.Second), 20),
+		// 1 msg per sec with some buffer
+		rateLimiterPerSec: rate.NewLimiter(rate.Every(1050*time.Millisecond), 1),
 	}
 }