From b1847509a4e5120f55fb308851acb65f823ca8c3 Mon Sep 17 00:00:00 2001
From: "1808837298@qq.com" <1808837298@qq.com>
Date: Wed, 19 Feb 2025 15:12:26 +0800
Subject: [PATCH] refactor: Optimize user caching and token retrieval methods

---
 controller/pricing.go  |  2 +-
 controller/user.go     |  2 +-
 model/token_cache.go   |  2 +-
 model/user.go          |  4 ++--
 model/user_cache.go    | 15 ++++++++++++++-
 service/quota.go       |  2 +-
 service/user_notify.go |  2 +-
 7 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/controller/pricing.go b/controller/pricing.go
index 36caff9d1..d7af5a4c8 100644
--- a/controller/pricing.go
+++ b/controller/pricing.go
@@ -17,7 +17,7 @@ func GetPricing(c *gin.Context) {
 	}
 	var group string
 	if exists {
-		user, err := model.GetUserById(userId.(int), false)
+		user, err := model.GetUserCache(userId.(int))
 		if err == nil {
 			group = user.Group
 		}
diff --git a/controller/user.go b/controller/user.go
index ac2cc8394..3002a6131 100644
--- a/controller/user.go
+++ b/controller/user.go
@@ -472,7 +472,7 @@ func GetUserModels(c *gin.Context) {
 	if err != nil {
 		id = c.GetInt("id")
 	}
-	user, err := model.GetUserById(id, true)
+	user, err := model.GetUserCache(id)
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
diff --git a/model/token_cache.go b/model/token_cache.go
index 99b762f51..0fe02fea5 100644
--- a/model/token_cache.go
+++ b/model/token_cache.go
@@ -52,7 +52,7 @@ func cacheSetTokenField(key string, field string, value string) error {
 func cacheGetTokenByKey(key string) (*Token, error) {
 	hmacKey := common.GenerateHMAC(key)
 	if !common.RedisEnabled {
-		return nil, nil
+		return nil, fmt.Errorf("redis is not enabled")
 	}
 	var token Token
 	err := common.RedisHGetObj(fmt.Sprintf("token:%s", hmacKey), &token)
diff --git a/model/user.go b/model/user.go
index 5aa0bdd31..427b0625f 100644
--- a/model/user.go
+++ b/model/user.go
@@ -42,8 +42,8 @@ type User struct {
 	Setting          string         `json:"setting" gorm:"type:text;column:setting"`
 }
 
-func (user *User) ToBaseUser() UserBase {
-	cache := UserBase{
+func (user *User) ToBaseUser() *UserBase {
+	cache := &UserBase{
 		Id:       user.Id,
 		Group:    user.Group,
 		Quota:    user.Quota,
diff --git a/model/user_cache.go b/model/user_cache.go
index 38ae03975..cc08288d6 100644
--- a/model/user_cache.go
+++ b/model/user_cache.go
@@ -79,7 +79,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) {
 	}()
 
 	// Try getting from Redis first
-	err = common.RedisHGetObj(getUserCacheKey(userId), &userCache)
+	userCache, err = cacheGetUserBase(userId)
 	if err == nil {
 		return userCache, nil
 	}
@@ -105,6 +105,19 @@ func GetUserCache(userId int) (userCache *UserBase, err error) {
 	return userCache, nil
 }
 
+func cacheGetUserBase(userId int) (*UserBase, error) {
+	if !common.RedisEnabled {
+		return nil, fmt.Errorf("redis is not enabled")
+	}
+	var userCache UserBase
+	// Try getting from Redis first
+	err := common.RedisHGetObj(getUserCacheKey(userId), &userCache)
+	if err != nil {
+		return nil, err
+	}
+	return &userCache, nil
+}
+
 // Add atomic quota operations using hash fields
 func cacheIncrUserQuota(userId int, delta int64) error {
 	if !common.RedisEnabled {
diff --git a/service/quota.go b/service/quota.go
index 13ce97631..2ec04fe0a 100644
--- a/service/quota.go
+++ b/service/quota.go
@@ -252,7 +252,7 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
 	//if relayInfo.TokenUnlimited {
 	//	return nil
 	//}
-	token, err := model.GetTokenById(relayInfo.TokenId)
+	token, err := model.GetTokenByKey(relayInfo.TokenKey, false)
 	if err != nil {
 		return err
 	}
diff --git a/service/user_notify.go b/service/user_notify.go
index d8b3c939e..d51bbcec0 100644
--- a/service/user_notify.go
+++ b/service/user_notify.go
@@ -11,7 +11,7 @@ import (
 
 func NotifyRootUser(t string, subject string, content string) {
 	user := model.GetRootUser().ToBaseUser()
-	_ = NotifyUser(&user, dto.NewNotify(t, subject, content, nil))
+	_ = NotifyUser(user, dto.NewNotify(t, subject, content, nil))
 }
 
 func NotifyUser(user *model.UserBase, data dto.Notify) error {