From bf349456e6422db30abe134e07c937abb4b34d0c Mon Sep 17 00:00:00 2001
From: Andrew Zhukovskiy <dir01@dir01.org>
Date: Mon, 26 Dec 2022 11:27:35 +0200
Subject: [PATCH] Added `RegisterHandlerMatchFunc` (#9)

which registers handler that can be matched based on
arbitrary `func(update *models.Update) bool`
---
 bot.go            |  1 +
 handlers.go       | 41 ++++++++++++++++++++++++++++++++++++++---
 process_update.go |  9 +++++----
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/bot.go b/bot.go
index 8abe4db..79a56a4 100644
--- a/bot.go
+++ b/bot.go
@@ -24,6 +24,7 @@ type HttpClient interface {
 type ErrorsHandler func(err error)
 type Middleware func(next HandlerFunc) HandlerFunc
 type HandlerFunc func(ctx context.Context, bot *Bot, update *models.Update)
+type MatchFunc func(update *models.Update) bool
 
 // Bot represents Telegram Bot main object
 type Bot struct {
diff --git a/handlers.go b/handlers.go
index 85dc678..f6451ba 100644
--- a/handlers.go
+++ b/handlers.go
@@ -3,6 +3,8 @@ package bot
 import (
 	"regexp"
 	"strings"
+
+	"github.com/go-telegram/bot/models"
 )
 
 type HandlerType int
@@ -20,17 +22,32 @@ const (
 	MatchTypeContains
 
 	matchTypeRegexp
+	matchTypeFunc
 )
 
 type handler struct {
 	handlerType HandlerType
 	matchType   MatchType
-	pattern     string
 	handler     HandlerFunc
-	re          *regexp.Regexp
+
+	pattern   string
+	re        *regexp.Regexp
+	matchFunc MatchFunc
 }
 
-func (h handler) match(data string) bool {
+func (h handler) match(update *models.Update) bool {
+	if h.matchType == matchTypeFunc {
+		return h.matchFunc(update)
+	}
+
+	var data string
+	switch h.handlerType {
+	case HandlerTypeMessageText:
+		data = update.Message.Text
+	case HandlerTypeCallbackQueryData:
+		data = update.CallbackQuery.Data
+	}
+
 	if h.matchType == MatchTypeExact {
 		return data == h.pattern
 	}
@@ -46,6 +63,24 @@ func (h handler) match(data string) bool {
 	return false
 }
 
+func (b *Bot) RegisterHandlerMatchFunc(handlerType HandlerType, matchFunc MatchFunc, f HandlerFunc) string {
+	b.handlersMx.Lock()
+	defer b.handlersMx.Unlock()
+
+	id := RandomString(16)
+
+	h := handler{
+		handlerType: handlerType,
+		matchType:   matchTypeFunc,
+		matchFunc:   matchFunc,
+		handler:     f,
+	}
+
+	b.handlers[id] = h
+
+	return id
+}
+
 func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp, f HandlerFunc) string {
 	b.handlersMx.Lock()
 	defer b.handlersMx.Unlock()
diff --git a/process_update.go b/process_update.go
index 72b5427..dd958fd 100644
--- a/process_update.go
+++ b/process_update.go
@@ -2,6 +2,7 @@ package bot
 
 import (
 	"context"
+
 	"github.com/go-telegram/bot/models"
 )
 
@@ -24,22 +25,22 @@ func (b *Bot) processUpdate(ctx context.Context, upd *models.Update) {
 	}()
 
 	if upd.Message != nil {
-		h = b.findHandler(HandlerTypeMessageText, upd.Message.Text)
+		h = b.findHandler(HandlerTypeMessageText, upd)
 		return
 	}
 	if upd.CallbackQuery != nil {
-		h = b.findHandler(HandlerTypeCallbackQueryData, upd.CallbackQuery.Data)
+		h = b.findHandler(HandlerTypeCallbackQueryData, upd)
 		return
 	}
 }
 
-func (b *Bot) findHandler(handlerType HandlerType, pattern string) HandlerFunc {
+func (b *Bot) findHandler(handlerType HandlerType, upd *models.Update) HandlerFunc {
 	b.handlersMx.RLock()
 	defer b.handlersMx.RUnlock()
 
 	for _, h := range b.handlers {
 		if h.handlerType == handlerType {
-			if h.match(pattern) {
+			if h.match(upd) {
 				return h.handler
 			}
 		}