From a266a99a47e881acaf0368085a4dc6fa60420daa Mon Sep 17 00:00:00 2001
From: Roy233 <oisroy233@gmail.com>
Date: Mon, 12 Feb 2024 00:11:54 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=AF=8F=E6=97=A5=E9=99=90?=
 =?UTF-8?q?=E9=A2=9D=E6=9C=BA=E5=88=B6=20=E8=8B=A5=E8=BD=AC=E6=8D=A2/?=
 =?UTF-8?q?=E4=B8=8B=E8=BD=BD=E4=B8=8D=E6=88=90=E5=8A=9F=EF=BC=8C=E5=88=99?=
 =?UTF-8?q?=E4=B8=8D=E6=B6=88=E8=80=97=E6=AC=A1=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 db/usageLimit.go                   | 26 ++++++++++++++++++++++++--
 handler/AnimationMessage.go        |  6 ++++++
 handler/DownloadStickerSetQuery.go |  5 +++++
 handler/StickersMessage.go         |  5 +++++
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/db/usageLimit.go b/db/usageLimit.go
index 2b39656..66cefff 100644
--- 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
index 478d68e..95cab7d 100644
--- 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
index a013fc7..f113730 100644
--- 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
index 21cccd1..8c4cd33 100644
--- 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)),
 	)))