From 7f34de577e30aa161fe6d4f85fabe980cff4e372 Mon Sep 17 00:00:00 2001 From: Dharsan Date: Fri, 28 Jun 2024 14:43:43 +0530 Subject: [PATCH] Feature: Ability to verify webhook secret token inside the bot program (#95) * feat(webhook-token): ability to verify webhook secret token inside the bot program * test(webhook-token): added additional tests for wrong tokens --- README.md | 2 + bot.go | 9 ++-- bot_test.go | 103 ++++++++++++++++++++++++++++++++++++++++++++- options.go | 7 +++ webhook_handler.go | 5 +++ 5 files changed, 120 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 8d39710..2f623a6 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ func main() { opts := []bot.Option{ bot.WithDefaultHandler(handler), + bot.WithWebhookSecretToken(os.Getenv("EXAMPLE_TELEGRAM_WEBHOOK_SECRET_TOKEN")) } b, _ := bot.New(os.Getenv("EXAMPLE_TELEGRAM_BOT_TOKEN"), opts...) @@ -181,6 +182,7 @@ b, err := bot.New("YOUR_BOT_TOKEN_FROM_BOTFATHER", opts...) - `WithSkipGetMe()` - skip call GetMe on bot init - `WithAllowedUpdates(params AllowedUpdates)` - set [allowed_updates](https://core.telegram.org/bots/api#getupdates) for getUpdates method - `WithUpdatesChannelCap(cap int)` - set updates channel capacity, by default 1024 +- `WithWebhookSecretToken(webhookSecretToken string)` - set X-Telegram-Bot-Api-Secret-Token header sent from telegram servers to confirm validity of update ## Message.Text and CallbackQuery.Data handlers diff --git a/bot.go b/bot.go index 313d4df..b9b8247 100644 --- a/bot.go +++ b/bot.go @@ -30,10 +30,11 @@ type MatchFunc func(update *models.Update) bool // Bot represents Telegram Bot main object type Bot struct { - url string - token string - pollTimeout time.Duration - skipGetMe bool + url string + token string + pollTimeout time.Duration + skipGetMe bool + webhookSecretToken string defaultHandlerFunc HandlerFunc diff --git a/bot_test.go b/bot_test.go index ac26d47..86ff03a 100644 --- a/bot_test.go +++ b/bot_test.go @@ -164,11 +164,12 @@ func TestNew(t *testing.T) { } } -func TestBot_StartWebhook(t *testing.T) { +func TestBot_StartWebhookWithCorrectSecret(t *testing.T) { s := newServerMock("xxx") defer s.Close() - b, err := New("xxx", WithServerURL(s.URL())) + opts := []Option{WithServerURL(s.URL()), WithWebhookSecretToken("zzzz")} + b, err := New("xxx", opts...) if err != nil { t.Fatalf("unexpected error %q", err) } @@ -197,6 +198,7 @@ func TestBot_StartWebhook(t *testing.T) { t.Error(errReq) return } + req.Header.Set("X-Telegram-Bot-Api-Secret-Token", "zzzz") b.WebhookHandler().ServeHTTP(nil, req) @@ -211,6 +213,103 @@ func TestBot_StartWebhook(t *testing.T) { mx.Unlock() } +func TestBot_StartWebhookWithNoSecret(t *testing.T) { + s := newServerMock("xxx") + defer s.Close() + + opts := []Option{WithServerURL(s.URL())} + b, err := New("xxx", opts...) + if err != nil { + t.Fatalf("unexpected error %q", err) + } + + mx := sync.Mutex{} + var called bool + + b.defaultHandlerFunc = func(ctx context.Context, bot *Bot, update *models.Update) { + if update.Message.ID != 1 { + t.Errorf("unexpected message id") + } + mx.Lock() + called = true + mx.Unlock() + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go b.StartWebhook(ctx) + + time.Sleep(time.Millisecond * 100) + + req, errReq := http.NewRequest(http.MethodPost, "", strings.NewReader(`{"update_id":1,"message":{"message_id":1}}`)) + if errReq != nil { + t.Error(errReq) + return + } + + b.WebhookHandler().ServeHTTP(nil, req) + + cancel() + + time.Sleep(time.Millisecond * 100) + + mx.Lock() + if !called { + t.Errorf("not called default handler") + } + mx.Unlock() +} + +func TestBot_StartWebhookWithWrongSecret(t *testing.T) { + s := newServerMock("xxx") + defer s.Close() + + opts := []Option{WithServerURL(s.URL()), WithWebhookSecretToken("zzzz")} + b, err := New("xxx", opts...) + if err != nil { + t.Fatalf("unexpected error %q", err) + } + + mx := sync.Mutex{} + var called bool + + b.defaultHandlerFunc = func(ctx context.Context, bot *Bot, update *models.Update) { + if update.Message.ID != 1 { + t.Errorf("unexpected message id") + } + mx.Lock() + called = true + mx.Unlock() + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go b.StartWebhook(ctx) + + time.Sleep(time.Millisecond * 100) + + req, errReq := http.NewRequest(http.MethodPost, "", strings.NewReader(`{"update_id":1,"message":{"message_id":1}}`)) + if errReq != nil { + t.Error(errReq) + return + } + req.Header.Set("X-Telegram-Bot-Api-Secret-Token", "wrong_secret") + + b.WebhookHandler().ServeHTTP(nil, req) + + cancel() + + time.Sleep(time.Millisecond * 100) + + mx.Lock() + if called { + t.Errorf("not supposed to call the default handler with wrong token") + } + mx.Unlock() +} + func TestBot_Start(t *testing.T) { s := newServerMock("xxx") defer s.Close() diff --git a/options.go b/options.go index da3a0fd..667c96d 100644 --- a/options.go +++ b/options.go @@ -102,3 +102,10 @@ func WithUpdatesChannelCap(cap int) Option { b.updates = make(chan *models.Update, cap) } } + +// WithWebhookSecretToken allows setting X-Telegram-Bot-Api-Secret-Token sent from Telegram servers +func WithWebhookSecretToken(webhookSecretToken string) Option { + return func(b *Bot) { + b.webhookSecretToken = webhookSecretToken + } +} diff --git a/webhook_handler.go b/webhook_handler.go index 1734a3f..9eb23af 100644 --- a/webhook_handler.go +++ b/webhook_handler.go @@ -10,6 +10,11 @@ import ( func (b *Bot) WebhookHandler() http.HandlerFunc { return func(_ http.ResponseWriter, req *http.Request) { + if b.webhookSecretToken != "" && req.Header.Get("X-Telegram-Bot-Api-Secret-Token") != b.webhookSecretToken { + b.error("invalid webhook secret token received from update") + return + } + body, errReadBody := io.ReadAll(req.Body) if errReadBody != nil { b.error("error read request body, %w", errReadBody)