From db32473a6a6192aafbd5f7d6abeea69ed081eda9 Mon Sep 17 00:00:00 2001
From: TBXark <tbxark@outlook.com>
Date: Mon, 16 Sep 2024 16:10:59 +0800
Subject: [PATCH] feat: add middlewares for handler (#105)

* feat: each handler can add middleware.

* fix: singleFlight middleware example
---
 examples/middleware/main.go | 20 ++++++++++++++++++++
 handlers.go                 | 12 ++++++------
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/examples/middleware/main.go b/examples/middleware/main.go
index adb6209..854197a 100644
--- a/examples/middleware/main.go
+++ b/examples/middleware/main.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"sync"
 
 	"github.com/go-telegram/bot"
 	"github.com/go-telegram/bot/models"
@@ -28,6 +29,10 @@ func main() {
 		panic(err)
 	}
 
+	b.RegisterHandler(bot.HandlerTypeCallbackQueryData, "", bot.MatchTypeExact, func(ctx context.Context, b *bot.Bot, update *models.Update) {
+		log.Printf("callback query data: %s", update.CallbackQuery.Data)
+	}, singleFlight)
+
 	b.Start(ctx)
 }
 
@@ -49,6 +54,21 @@ func showMessageWithUserName(next bot.HandlerFunc) bot.HandlerFunc {
 	}
 }
 
+// singleFlight is a middleware that ensures that only one callback query is processed at a time.
+func singleFlight(next bot.HandlerFunc) bot.HandlerFunc {
+	sf := sync.Map{}
+	return func(ctx context.Context, b *bot.Bot, update *models.Update) {
+		if update.CallbackQuery != nil {
+			key := update.CallbackQuery.Message.Message.ID
+			if _, loaded := sf.LoadOrStore(key, struct{}{}); loaded {
+				return
+			}
+			defer sf.Delete(key)
+			next(ctx, b, update)
+		}
+	}
+}
+
 func handler(ctx context.Context, b *bot.Bot, update *models.Update) {
 	b.SendMessage(ctx, &bot.SendMessageParams{
 		ChatID: update.Message.Chat.ID,
diff --git a/handlers.go b/handlers.go
index bcf8e36..7dbd063 100644
--- a/handlers.go
+++ b/handlers.go
@@ -70,7 +70,7 @@ func (h handler) match(update *models.Update) bool {
 	return false
 }
 
-func (b *Bot) RegisterHandlerMatchFunc(matchFunc MatchFunc, f HandlerFunc) string {
+func (b *Bot) RegisterHandlerMatchFunc(matchFunc MatchFunc, f HandlerFunc, m ...Middleware) string {
 	b.handlersMx.Lock()
 	defer b.handlersMx.Unlock()
 
@@ -79,7 +79,7 @@ func (b *Bot) RegisterHandlerMatchFunc(matchFunc MatchFunc, f HandlerFunc) strin
 	h := handler{
 		matchType: matchTypeFunc,
 		matchFunc: matchFunc,
-		handler:   f,
+		handler:   applyMiddlewares(f, m...),
 	}
 
 	b.handlers[id] = h
@@ -87,7 +87,7 @@ func (b *Bot) RegisterHandlerMatchFunc(matchFunc MatchFunc, f HandlerFunc) strin
 	return id
 }
 
-func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp, f HandlerFunc) string {
+func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp, f HandlerFunc, m ...Middleware) string {
 	b.handlersMx.Lock()
 	defer b.handlersMx.Unlock()
 
@@ -97,7 +97,7 @@ func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp,
 		handlerType: handlerType,
 		matchType:   matchTypeRegexp,
 		re:          re,
-		handler:     f,
+		handler:     applyMiddlewares(f, m...),
 	}
 
 	b.handlers[id] = h
@@ -105,7 +105,7 @@ func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp,
 	return id
 }
 
-func (b *Bot) RegisterHandler(handlerType HandlerType, pattern string, matchType MatchType, f HandlerFunc) string {
+func (b *Bot) RegisterHandler(handlerType HandlerType, pattern string, matchType MatchType, f HandlerFunc, m ...Middleware) string {
 	b.handlersMx.Lock()
 	defer b.handlersMx.Unlock()
 
@@ -115,7 +115,7 @@ func (b *Bot) RegisterHandler(handlerType HandlerType, pattern string, matchType
 		handlerType: handlerType,
 		matchType:   matchType,
 		pattern:     pattern,
-		handler:     f,
+		handler:     applyMiddlewares(f, m...),
 	}
 
 	b.handlers[id] = h