From 097860af6b6557563791724114a5babcfe59c7b4 Mon Sep 17 00:00:00 2001 From: zenix Date: Tue, 20 Sep 2022 17:02:02 +0900 Subject: [PATCH 1/3] fix: add rate limit on telegram api and split messages by unicode with size limitation --- pkg/interact/telegram.go | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/pkg/interact/telegram.go b/pkg/interact/telegram.go index 9d4cfc7c6..68f341f5b 100644 --- a/pkg/interact/telegram.go +++ b/pkg/interact/telegram.go @@ -5,8 +5,10 @@ import ( "fmt" "strings" "time" + "unicode/utf8" log "github.com/sirupsen/logrus" + "golang.org/x/time/rate" "gopkg.in/tucnak/telebot.v2" ) @@ -15,6 +17,10 @@ func init() { _ = Reply(&TelegramReply{}) } +var sendLimiter = rate.NewLimiter(10, 2) + +const maxMessageSize int = 3000 + type TelegramSessionMap map[int64]*TelegramSession type TelegramSession struct { @@ -62,7 +68,14 @@ type TelegramReply struct { } func (r *TelegramReply) Send(message string) { - checkSendErr(r.bot.Send(r.session.Chat, message)) + var left, right int + for left, right = 0, maxMessageSize; right < len(message); left, right = right, right+maxMessageSize { + for !utf8.RuneStart(message[right]) { + right-- + } + checkSendErr(r.bot.Send(r.session.Chat, message[left:right])) + } + checkSendErr(r.bot.Send(r.session.Chat, message[left:])) } func (r *TelegramReply) Message(message string) { @@ -132,7 +145,7 @@ func (tm *Telegram) SetTextMessageResponder(responder Responder) { tm.textMessageResponder = responder } -func (tm *Telegram) Start(context.Context) { +func (tm *Telegram) Start(ctx context.Context) { tm.Bot.Handle(telebot.OnCallback, func(c *telebot.Callback) { log.Infof("[telegram] onCallback: %+v", c) }) @@ -158,7 +171,24 @@ func (tm *Telegram) Start(context.Context) { if reply.set { reply.build() if len(reply.message) > 0 || reply.menu != nil { - checkSendErr(tm.Bot.Send(m.Chat, reply.message, reply.menu)) + var left, right int + for left, right = 0, maxMessageSize; right < len(reply.message); left, right = right, right+maxMessageSize { + for !utf8.RuneStart(reply.message[right]) { + right-- + } + if err := sendLimiter.Wait(ctx); err != nil { + log.WithError(err).Errorf("telegram send limit exceeded") + } + checkSendErr(tm.Bot.Send(m.Chat, reply.message[left:right])) + + } + if left < len(reply.message) { + if err := sendLimiter.Wait(ctx); err != nil { + log.WithError(err).Errorf("telegram send limit exceeded") + } + // only set menu on the last message + checkSendErr(tm.Bot.Send(m.Chat, reply.message[left:], reply.menu)) + } } } }) From 9cce165aa516f87aab15b17710a8e6e658e31965 Mon Sep 17 00:00:00 2001 From: zenix Date: Wed, 21 Sep 2022 15:14:33 +0900 Subject: [PATCH 2/3] fix: extract split string by length as a function --- pkg/interact/telegram.go | 33 +++++++++++++-------------------- pkg/util/string.go | 19 ++++++++++++++++++- pkg/util/string_test.go | 13 ++++++++++++- 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/pkg/interact/telegram.go b/pkg/interact/telegram.go index 68f341f5b..deba77053 100644 --- a/pkg/interact/telegram.go +++ b/pkg/interact/telegram.go @@ -5,8 +5,8 @@ import ( "fmt" "strings" "time" - "unicode/utf8" + "github.com/c9s/bbgo/pkg/util" log "github.com/sirupsen/logrus" "golang.org/x/time/rate" "gopkg.in/tucnak/telebot.v2" @@ -68,14 +68,13 @@ type TelegramReply struct { } func (r *TelegramReply) Send(message string) { - var left, right int - for left, right = 0, maxMessageSize; right < len(message); left, right = right, right+maxMessageSize { - for !utf8.RuneStart(message[right]) { - right-- + splits := util.StringSplitByLength(message, maxMessageSize) + for _, split := range splits { + if err := sendLimiter.Wait(ctx); err != nil { + log.WithError(err).Errorf("telegram send limit exceeded") } - checkSendErr(r.bot.Send(r.session.Chat, message[left:right])) + checkSendErr(r.bot.Send(r.session.Chat, split)) } - checkSendErr(r.bot.Send(r.session.Chat, message[left:])) } func (r *TelegramReply) Message(message string) { @@ -171,23 +170,17 @@ func (tm *Telegram) Start(ctx context.Context) { if reply.set { reply.build() if len(reply.message) > 0 || reply.menu != nil { - var left, right int - for left, right = 0, maxMessageSize; right < len(reply.message); left, right = right, right+maxMessageSize { - for !utf8.RuneStart(reply.message[right]) { - right-- - } + splits := util.StringSplitByLength(reply.message, maxMessageSize) + for i, split := range splits { if err := sendLimiter.Wait(ctx); err != nil { log.WithError(err).Errorf("telegram send limit exceeded") } - checkSendErr(tm.Bot.Send(m.Chat, reply.message[left:right])) - - } - if left < len(reply.message) { - if err := sendLimiter.Wait(ctx); err != nil { - log.WithError(err).Errorf("telegram send limit exceeded") + if i == len(splits)-1 { + // only set menu on the last message + checkSendErr(tm.Bot.Send(m.Chat, split, reply.menu)) + } else { + checkSendErr(tm.Bot.Send(m.Chat, split)) } - // only set menu on the last message - checkSendErr(tm.Bot.Send(m.Chat, reply.message[left:], reply.menu)) } } } diff --git a/pkg/util/string.go b/pkg/util/string.go index 268233acb..af3b25bf4 100644 --- a/pkg/util/string.go +++ b/pkg/util/string.go @@ -1,6 +1,9 @@ package util -import "strings" +import ( + "strings" + "unicode/utf8" +) func StringSliceContains(slice []string, needle string) bool { for _, s := range slice { @@ -27,3 +30,17 @@ func MaskKey(key string) string { maskKey += key[len(key)-h:] return maskKey } + +func StringSplitByLength(s string, length int) (result []string) { + var left, right int + for left, right = 0, length; right < len(s); left, right = right, right+length { + for !utf8.RuneStart(s[right]) { + right-- + } + result = append(result, s[left:right]) + } + if len(s)-left > 0 { + result = append(result, s[left:]) + } + return result +} diff --git a/pkg/util/string_test.go b/pkg/util/string_test.go index 7d55f425a..4f09a3a26 100644 --- a/pkg/util/string_test.go +++ b/pkg/util/string_test.go @@ -1,6 +1,10 @@ package util -import "testing" +import ( + "testing" + + "github.com/stretchr/testify/assert" +) func TestMaskKey(t *testing.T) { type args struct { @@ -40,3 +44,10 @@ func TestMaskKey(t *testing.T) { }) } } + +func TestStringSplitByLength(t *testing.T) { + result := StringSplitByLength("1234567890", 3) + assert.Equal(t, result, []string{"123", "456", "789", "0"}) + result = StringSplitByLength("123許456", 4) + assert.Equal(t, result, []string{"123", "許4", "56"}) +} From d8dea22e10c5b4462b502b9836fe4371dd21919f Mon Sep 17 00:00:00 2001 From: zenix Date: Wed, 21 Sep 2022 15:32:55 +0900 Subject: [PATCH 3/3] fix: set ctx --- pkg/interact/telegram.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/interact/telegram.go b/pkg/interact/telegram.go index deba77053..c2a45a002 100644 --- a/pkg/interact/telegram.go +++ b/pkg/interact/telegram.go @@ -68,6 +68,7 @@ type TelegramReply struct { } func (r *TelegramReply) Send(message string) { + ctx := context.Background() splits := util.StringSplitByLength(message, maxMessageSize) for _, split := range splits { if err := sendLimiter.Wait(ctx); err != nil {