refactor: bind quota to account instead of token (close songquanpeng#64…

…, songquanpeng#31)

common/constants.go

@@ -50,10 +50,10 @@ var WeChatAccountQRCodeImageURL = ""
 var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
-var QuotaForNewUser = 100
+var QuotaForNewUser = 0
 var ChannelDisableThreshold = 5.0
 var AutomaticDisableChannelEnabled = false
-var QuotaRemindThreshold = 1000 // TODO: QuotaRemindThreshold
+var QuotaRemindThreshold = 1000
 
 var RootUserEmail = ""
 

controller/token.go

@@ -100,7 +100,6 @@ func GetTokenStatus(c *gin.Context) {
 }
 
 func AddToken(c *gin.Context) {
-	isAdmin := c.GetInt("role") >= common.RoleAdminUser
 	token := model.Token{}
 	err := c.ShouldBindJSON(&token)
 	if err != nil {
@@ -118,27 +117,14 @@ func AddToken(c *gin.Context) {
 		return
 	}
 	cleanToken := model.Token{
-		UserId:       c.GetInt("id"),
-		Name:         token.Name,
-		Key:          common.GetUUID(),
-		CreatedTime:  common.GetTimestamp(),
-		AccessedTime: common.GetTimestamp(),
-		ExpiredTime:  token.ExpiredTime,
-	}
-	if isAdmin {
-		cleanToken.RemainQuota = token.RemainQuota
-		cleanToken.UnlimitedQuota = token.UnlimitedQuota
-	} else {
-		userId := c.GetInt("id")
-		quota, err := model.GetUserQuota(userId)
-		if err != nil {
-			c.JSON(http.StatusOK, gin.H{
-				"success": false,
-				"message": err.Error(),
-			})
-			return
-		}
-		cleanToken.RemainQuota = quota
+		UserId:         c.GetInt("id"),
+		Name:           token.Name,
+		Key:            common.GetUUID(),
+		CreatedTime:    common.GetTimestamp(),
+		AccessedTime:   common.GetTimestamp(),
+		ExpiredTime:    token.ExpiredTime,
+		RemainQuota:    token.RemainQuota,
+		UnlimitedQuota: token.UnlimitedQuota,
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -148,10 +134,6 @@ func AddToken(c *gin.Context) {
 		})
 		return
 	}
-	if !isAdmin {
-		// update user quota
-		err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainQuota)
-	}
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -240,34 +222,3 @@ func UpdateToken(c *gin.Context) {
 	})
 	return
 }
-
-type topUpRequest struct {
-	Id  int    `json:"id"`
-	Key string `json:"key"`
-}
-
-func TopUp(c *gin.Context) {
-	req := topUpRequest{}
-	err := c.ShouldBindJSON(&req)
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
-		return
-	}
-	quota, err := model.Redeem(req.Key, req.Id)
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
-		return
-	}
-	c.JSON(http.StatusOK, gin.H{
-		"success": true,
-		"message": "",
-		"data":    quota,
-	})
-	return
-}

controller/user.go

@@ -654,3 +654,34 @@ func EmailBind(c *gin.Context) {
 	})
 	return
 }
+
+type topUpRequest struct {
+	Key string `json:"key"`
+}
+
+func TopUp(c *gin.Context) {
+	req := topUpRequest{}
+	err := c.ShouldBindJSON(&req)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	id := c.GetInt("id")
+	quota, err := model.Redeem(req.Key, id)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    quota,
+	})
+	return
+}

model/redemption.go

@@ -40,12 +40,12 @@ func GetRedemptionById(id int) (*Redemption, error) {
 	return &redemption, err
 }
 
-func Redeem(key string, tokenId int) (quota int, err error) {
+func Redeem(key string, userId int) (quota int, err error) {
 	if key == "" {
 		return 0, errors.New("未提供兑换码")
 	}
-	if tokenId == 0 {
-		return 0, errors.New("未提供 token id")
+	if userId == 0 {
+		return 0, errors.New("无效的 user id")
 	}
 	redemption := &Redemption{}
 	err = DB.Where("`key` = ?", key).First(redemption).Error
@@ -55,7 +55,7 @@ func Redeem(key string, tokenId int) (quota int, err error) {
 	if redemption.Status != common.RedemptionCodeStatusEnabled {
 		return 0, errors.New("该兑换码已被使用")
 	}
-	err = IncreaseTokenQuota(tokenId, redemption.Quota)
+	err = IncreaseUserQuota(userId, redemption.Quota)
 	if err != nil {
 		return 0, err
 	}

model/token.go

@@ -2,6 +2,7 @@ package model
 
 import (
 	"errors"
+	"fmt"
 	_ "gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"one-api/common"
@@ -82,6 +83,16 @@ func GetTokenByIds(id int, userId int) (*Token, error) {
 	return &token, err
 }
 
+func GetTokenById(id int) (*Token, error) {
+	if id == 0 {
+		return nil, errors.New("id 为空！")
+	}
+	token := Token{Id: id}
+	var err error = nil
+	err = DB.First(&token, "id = ?", id).Error
+	return &token, err
+}
+
 func (token *Token) Insert() error {
 	var err error
 	err = DB.Create(token).Error
@@ -116,26 +127,53 @@ func DeleteTokenById(id int, userId int) (err error) {
 	if err != nil {
 		return err
 	}
-	quota := token.RemainQuota
-	if quota != 0 {
-		if quota > 0 {
-			err = IncreaseUserQuota(userId, quota)
-		} else {
-			err = DecreaseUserQuota(userId, -quota)
-		}
+	return token.Delete()
+}
+
+func DecreaseTokenQuota(tokenId int, quota int) (err error) {
+	if quota < 0 {
+		return errors.New("quota 不能为负数！")
 	}
+	token, err := GetTokenById(tokenId)
 	if err != nil {
 		return err
 	}
-	return token.Delete()
-}
-
-func IncreaseTokenQuota(id int, quota int) (err error) {
-	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error
-	return err
-}
-
-func DecreaseTokenQuota(id int, quota int) (err error) {
-	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
+	if token.RemainQuota < quota {
+		return errors.New("令牌额度不足")
+	}
+	userQuota, err := GetUserQuota(token.UserId)
+	if err != nil {
+		return err
+	}
+	if userQuota < quota {
+		return errors.New("用户额度不足")
+	}
+	quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold
+	noMoreQuota := userQuota-quota <= 0
+	if quotaTooLow || noMoreQuota {
+		go func() {
+			email, err := GetUserEmail(token.UserId)
+			if err != nil {
+				common.SysError("获取用户邮箱失败：" + err.Error())
+			}
+			prompt := "您的额度即将用尽"
+			if noMoreQuota {
+				prompt = "您的额度已用尽"
+			}
+			if email != "" {
+				topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress)
+				err = common.SendEmail(prompt, email,
+					fmt.Sprintf("%s，剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota-quota, topUpLink, topUpLink))
+				if err != nil {
+					common.SysError("发送邮件失败：" + err.Error())
+				}
+			}
+		}()
+	}
+	err = DB.Model(&Token{}).Where("id = ?", tokenId).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
+	if err != nil {
+		return err
+	}
+	err = DecreaseUserQuota(token.UserId, quota)
 	return err
 }

model/user.go

@@ -225,12 +225,23 @@ func GetUserQuota(id int) (quota int, err error) {
 	return quota, err
 }
 
+func GetUserEmail(id int) (email string, err error) {
+	err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
+	return email, err
+}
+
 func IncreaseUserQuota(id int, quota int) (err error) {
+	if quota < 0 {
+		return errors.New("quota 不能为负数！")
+	}
 	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
 	return err
 }
 
 func DecreaseUserQuota(id int, quota int) (err error) {
+	if quota < 0 {
+		return errors.New("quota 不能为负数！")
+	}
 	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
 	return err
 }

router/api-router.go

@@ -37,6 +37,7 @@ func SetApiRouter(router *gin.Engine) {
 				selfRoute.PUT("/self", controller.UpdateSelf)
 				selfRoute.DELETE("/self", controller.DeleteSelf)
 				selfRoute.GET("/token", controller.GenerateAccessToken)
+				selfRoute.POST("/topup", controller.TopUp)
 			}
 
 			adminRoute := userRoute.Group("/")
@@ -74,7 +75,6 @@ func SetApiRouter(router *gin.Engine) {
 		{
 			tokenRoute.GET("/", controller.GetAllTokens)
 			tokenRoute.GET("/search", controller.SearchTokens)
-			tokenRoute.POST("/topup", controller.TopUp)
 			tokenRoute.GET("/:id", controller.GetToken)
 			tokenRoute.POST("/", controller.AddToken)
 			tokenRoute.PUT("/", controller.UpdateToken)

web/src/App.js

@@ -21,6 +21,7 @@ import EditToken from './pages/Token/EditToken';
 import EditChannel from './pages/Channel/EditChannel';
 import Redemption from './pages/Redemption';
 import EditRedemption from './pages/Redemption/EditRedemption';
+import TopUp from './pages/TopUp';
 
 const Home = lazy(() => import('./pages/Home'));
 const About = lazy(() => import('./pages/About'));
@@ -239,6 +240,16 @@ function App() {
           </PrivateRoute>
         }
       />
+      <Route
+        path='/topup'
+        element={
+        <PrivateRoute>
+          <Suspense fallback={<Loading></Loading>}>
+            <TopUp />
+          </Suspense>
+        </PrivateRoute>
+        }
+      />
       <Route
         path='/about'
         element={

web/src/components/Header.js

@@ -30,6 +30,12 @@ const headerButtons = [
     icon: 'dollar sign',
     admin: true,
   },
+  {
+    name: '充值',
+    to: '/topup',
+    icon: 'cart',
+    admin: true,
+  },
   {
     name: '用户',
     to: '/user',