From 4b5993bd6f344d13cd1fdd0d0ba99aa4894fe6b1 Mon Sep 17 00:00:00 2001
From: indes <wl@live.it>
Date: Sat, 26 Feb 2022 17:11:56 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=B0=86v2=E7=89=88=E6=9C=AC=E7=9A=84b?=
 =?UTF-8?q?ot=20filter=20=E8=BD=AC=E6=8D=A2=E6=88=90v3=E7=89=88=E6=9C=AC?=
 =?UTF-8?q?=E7=9A=84middleware?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/bot/bot.go                    | 28 +++-------
 internal/bot/middleware/user_filter.go | 26 +++++++++
 internal/bot/service.go                | 74 --------------------------
 3 files changed, 32 insertions(+), 96 deletions(-)
 create mode 100644 internal/bot/middleware/user_filter.go

diff --git a/internal/bot/bot.go b/internal/bot/bot.go
index 5edf56f3..3f722a73 100644
--- a/internal/bot/bot.go
+++ b/internal/bot/bot.go
@@ -3,7 +3,6 @@ package bot
 import (
 	"time"
 
-	"github.com/indes/flowerss-bot/internal/bot/fsm"
 	"github.com/indes/flowerss-bot/internal/bot/handler"
 	"github.com/indes/flowerss-bot/internal/bot/middleware"
 	"github.com/indes/flowerss-bot/internal/config"
@@ -14,9 +13,6 @@ import (
 )
 
 var (
-	// UserState 用户状态，用于标示当前用户操作所在状态
-	UserState map[int64]fsm.UserStatus = make(map[int64]fsm.UserStatus)
-
 	// B bot
 	B *tb.Bot
 )
@@ -25,38 +21,27 @@ func init() {
 	if config.RunMode == config.TestMode {
 		return
 	}
-	poller := &tb.LongPoller{Timeout: 10 * time.Second}
-	spamProtected := tb.NewMiddlewarePoller(
-		poller, func(upd *tb.Update) bool {
-			if !isUserAllowed(upd) {
-				// 检查用户是否可以使用bot
-				return false
-			}
-
-			if !CheckAdmin(upd) {
-				return false
-			}
-			return true
-		},
-	)
 	zap.S().Infow(
 		"init telegram bot",
 		"token", config.BotToken,
 		"endpoint", config.TelegramEndpoint,
 	)
-
 	// create bot
 	var err error
 	B, err = tb.NewBot(
 		tb.Settings{
 			URL:     config.TelegramEndpoint,
 			Token:   config.BotToken,
-			Poller:  spamProtected,
+			Poller:  &tb.LongPoller{Timeout: 10 * time.Second},
 			Client:  util.HttpClient,
 			Verbose: true,
 		},
 	)
-	B.Use(middleware.PreLoadMentionChat(), middleware.IsChatAdmin())
+	B.Use(
+		middleware.UserFilter(),
+		middleware.PreLoadMentionChat(),
+		middleware.IsChatAdmin(),
+	)
 	if err != nil {
 		zap.S().Fatal(err)
 		return
@@ -124,5 +109,4 @@ func setCommands() {
 	if err := B.SetCommands(commands); err != nil {
 		zap.S().Errorw("set bot commands failed", "error", err.Error())
 	}
-
 }
diff --git a/internal/bot/middleware/user_filter.go b/internal/bot/middleware/user_filter.go
new file mode 100644
index 00000000..4204f9ba
--- /dev/null
+++ b/internal/bot/middleware/user_filter.go
@@ -0,0 +1,26 @@
+package middleware
+
+import (
+	"fmt"
+
+	tb "gopkg.in/telebot.v3"
+
+	"github.com/indes/flowerss-bot/internal/config"
+)
+
+func UserFilter() tb.MiddlewareFunc {
+	return func(next tb.HandlerFunc) tb.HandlerFunc {
+		return func(c tb.Context) error {
+			if len(config.AllowUsers) == 0 {
+				return next(c)
+			}
+			userID := c.Sender().ID
+			for _, allowUserID := range config.AllowUsers {
+				if allowUserID == userID {
+					return next(c)
+				}
+			}
+			return fmt.Errorf("deny user %d", userID)
+		}
+	}
+}
diff --git a/internal/bot/service.go b/internal/bot/service.go
index f9c6a632..8ef9bc35 100644
--- a/internal/bot/service.go
+++ b/internal/bot/service.go
@@ -102,80 +102,6 @@ func BroadcastSourceError(source *model.Source) {
 	}
 }
 
-// CheckAdmin check user is admin of group/channel
-func CheckAdmin(upd *tb.Update) bool {
-
-	if upd.Message != nil {
-		if HasAdminType(upd.Message.Chat.Type) {
-			adminList, _ := B.AdminsOf(upd.Message.Chat)
-			for _, admin := range adminList {
-				if admin.User.ID == upd.Message.Sender.ID {
-					return true
-				}
-			}
-
-			return false
-		}
-
-		return true
-	} else if upd.Callback != nil {
-		if HasAdminType(upd.Callback.Message.Chat.Type) {
-			adminList, _ := B.AdminsOf(upd.Callback.Message.Chat)
-			for _, admin := range adminList {
-				if admin.User.ID == upd.Callback.Sender.ID {
-					return true
-				}
-			}
-
-			return false
-		}
-
-		return true
-	}
-	return false
-}
-
-// IsUserAllowed check user is allowed to use bot
-func isUserAllowed(upd *tb.Update) bool {
-	if upd == nil {
-		return false
-	}
-
-	var userID int64
-
-	if upd.Message != nil {
-		userID = int64(upd.Message.Sender.ID)
-	} else if upd.Callback != nil {
-		userID = int64(upd.Callback.Sender.ID)
-	} else {
-		return false
-	}
-
-	if len(config.AllowUsers) == 0 {
-		return true
-	}
-
-	for _, allowUserID := range config.AllowUsers {
-		if allowUserID == userID {
-			return true
-		}
-	}
-
-	zap.S().Infow("user not allowed", "userID", userID)
-	return false
-}
-
-// HasAdminType check if the message is sent in the group/channel environment
-func HasAdminType(t tb.ChatType) bool {
-	hasAdmin := []tb.ChatType{tb.ChatGroup, tb.ChatSuperGroup, tb.ChatChannel, tb.ChatChannelPrivate}
-	for _, n := range hasAdmin {
-		if t == n {
-			return true
-		}
-	}
-	return false
-}
-
 // GetMentionFromMessage get message mention
 func GetMentionFromMessage(m *tb.Message) (mention string) {
 	if m.Text != "" {