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)), )))