优化每日限额机制

若转换/下载不成功，则不消耗次数
rroy233 · Feb 11, 2024 · a266a99 · a266a99
1 parent 8e54427
commit a266a99
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 2 deletions.
diff --git a/db/usageLimit.go b/db/usageLimit.go
@@ -1,6 +1,7 @@
 package db
 
 import (
+	"errors"
 	"fmt"
 	"github.com/go-redis/redis/v8"
 	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
@@ -9,7 +10,7 @@ import (
 	"time"
 )
 
-// CheckLimit 用户是否已达到今日限额
+// CheckLimit Determines if the user has reached today's limit
 func CheckLimit(update *tgbotapi.Update) bool {
 	UID := int64(0)
 	if update.Message != nil {
@@ -32,10 +33,31 @@ func CheckLimit(update *tgbotapi.Update) bool {
 	if limitTimes > config.Get().General.UserDailyLimit {
 		return true
 	}
-	rdb.Set(ctx, fmt.Sprintf("%s:UserLimit:%d", ServicePrefix, UID), limitTimes+1, redis.KeepTTL)
 	return false
 }
 
+// ConsumeLimit Consume the current user's daily limit
+func ConsumeLimit(update *tgbotapi.Update) error {
+	UID := int64(0)
+	if update.Message != nil {
+		UID = update.Message.Chat.ID
+	} else if update.CallbackQuery != nil {
+		UID = update.CallbackQuery.Message.Chat.ID
+	} else {
+		return errors.New("failed to get uid")
+	}
+
+	limit := rdb.Get(ctx, fmt.Sprintf("%s:UserLimit:%d", ServicePrefix, UID)).Val()
+	if limit == "" {
+		rdb.Set(ctx, fmt.Sprintf("%s:UserLimit:%d", ServicePrefix, UID), 1, 24*time.Hour)
+		return nil
+	}
+
+	limitTimes, _ := strconv.Atoi(limit)
+	rdb.Set(ctx, fmt.Sprintf("%s:UserLimit:%d", ServicePrefix, UID), limitTimes+1, redis.KeepTTL)
+	return nil
+}
+
 // 获取该用户已使用的次数
 func getUsed(UID int64) int {
 	limit := rdb.Get(ctx, fmt.Sprintf("%s:UserLimit:%d", ServicePrefix, UID)).Val()

diff --git a/handler/AnimationMessage.go b/handler/AnimationMessage.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
 	"github.com/rroy233/StickerDownloader/config"
+	"github.com/rroy233/StickerDownloader/db"
 	"github.com/rroy233/StickerDownloader/languages"
 	"github.com/rroy233/StickerDownloader/utils"
 	"gopkg.in/rroy233/logger.v2"
@@ -89,6 +90,11 @@ func AnimationMessage(update tgbotapi.Update) {
 		return
 	}
 
+	//Consume the current user's daily limit
+	if err = db.ConsumeLimit(&update); err != nil {
+		logger.Error.Println(userInfo + err.Error())
+	}
+
 	utils.EditMsgText(update.Message.Chat.ID, msg.MessageID, languages.Get(&update).BotMsg.ConvertCompleted)
 	if err != nil {
 		logger.Error.Println(userInfo+"failed to delete msg:", err)

diff --git a/handler/DownloadStickerSetQuery.go b/handler/DownloadStickerSetQuery.go
@@ -214,6 +214,11 @@ func DownloadStickerSetQuery(update tgbotapi.Update) {
 		logger.Info.Println(userInfo + "DownloadStickerSetQuery-upload(Telegram) successfully！！！")
 	}
 
+	//Consume the current user's daily limit
+	if err = db.ConsumeLimit(&update); err != nil {
+		logger.Error.Println(userInfo + "DownloadStickerSetQuery - " + err.Error())
+	}
+
 	return
 }
 

diff --git a/handler/StickersMessage.go b/handler/StickersMessage.go
@@ -125,6 +125,11 @@ func StickerMessage(update tgbotapi.Update) {
 		utils.RemoveFile(outPath)
 	}
 
+	//Consume the current user's daily limit
+	if err = db.ConsumeLimit(&update); err != nil {
+		logger.Error.Println(userInfo + err.Error())
+	}
+
 	err = utils.BotRequest(tgbotapi.NewEditMessageTextAndMarkup(update.Message.Chat.ID, msg.MessageID, languages.Get(&update).BotMsg.ConvertCompleted, tgbotapi.NewInlineKeyboardMarkup(
 		tgbotapi.NewInlineKeyboardRow(tgbotapi.NewInlineKeyboardButtonData(languages.Get(&update).BotMsg.DownloadStickerSet, DownloadStickerSetCallbackQuery)),
 	)))