From bc8ceca7428f5b39196883862177a4a65d2b8237 Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Fri, 29 Sep 2023 12:46:26 +0000
Subject: [PATCH] feat: api can modify token's used and remains

- Updated base image in Dockerfile from `node:16` to `node:18`
- Added `tokenPatch` struct to handle updates to the `Token` model
- Updated `UpdateToken` function to use `tokenPatch` struct for binding JSON
- Improved error handling and validation for token updates
- Refined checks for token status, expiration, and quota
- Enhanced error handling and response for token updates
- Updated JSON response to include `data` field
---
 Dockerfile          |  4 +-
 controller/token.go | 96 ++++++++++++++++++++++++++++++---------------
 2 files changed, 67 insertions(+), 33 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ffb8c21b..b378b5bc 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM node:16 as builder
+FROM node:18 as builder
 
 WORKDIR /build
 COPY web/package.json .
@@ -7,7 +7,7 @@ COPY ./web .
 COPY ./VERSION .
 RUN DISABLE_ESLINT_PLUGIN='true' REACT_APP_VERSION=$(cat VERSION) npm run build
 
-FROM golang AS builder2
+FROM golang:1.21.1 AS builder2
 
 ENV GO111MODULE=on \
     CGO_ENABLED=1 \
diff --git a/controller/token.go b/controller/token.go
index 8642122c..a452102b 100644
--- a/controller/token.go
+++ b/controller/token.go
@@ -1,11 +1,13 @@
 package controller
 
 import (
-	"github.com/gin-gonic/gin"
+	"fmt"
 	"net/http"
 	"one-api/common"
 	"one-api/model"
 	"strconv"
+
+	"github.com/gin-gonic/gin"
 )
 
 func GetAllTokens(c *gin.Context) {
@@ -159,42 +161,62 @@ func DeleteToken(c *gin.Context) {
 	return
 }
 
+type updateTokenDto struct {
+	Id             int     `json:"id"`
+	Status         int     `json:"status" gorm:"default:1"`
+	Name           *string `json:"name" gorm:"index" `
+	ExpiredTime    *int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
+	RemainQuota    *int    `json:"remain_quota" gorm:"default:0"`
+	UnlimitedQuota *bool   `json:"unlimited_quota" gorm:"default:false"`
+	// AddRemainQuota add or subtract remain quota
+	AddRemainQuota int `json:"add_remain_quota"`
+	// AddUsedQuota add or subtract used quota
+	AddUsedQuota int `json:"add_used_quota"`
+}
+
 func UpdateToken(c *gin.Context) {
 	userId := c.GetInt("id")
 	statusOnly := c.Query("status_only")
-	token := model.Token{}
-	err := c.ShouldBindJSON(&token)
+	tokenPatch := new(updateTokenDto)
+	if err := c.ShouldBindJSON(tokenPatch); err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "parse request join: " + err.Error(),
+		})
+		return
+	}
+
+	if tokenPatch.Name != nil &&
+		(len(*tokenPatch.Name) > 30 || len(*tokenPatch.Name) == 0) {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "令牌名称错误，长度应在 1-30 之间",
+		})
+		return
+	}
+
+	cleanToken, err := model.GetTokenByIds(tokenPatch.Id, userId)
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
-			"message": err.Error(),
+			"message": fmt.Sprintf("get token by id %d: %s", tokenPatch.Id, err.Error()),
 		})
 		return
 	}
-	if len(token.Name) > 30 {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": "令牌名称过长",
-		})
-		return
-	}
-	cleanToken, err := model.GetTokenByIds(token.Id, userId)
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
-		return
-	}
-	if token.Status == common.TokenStatusEnabled {
-		if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 {
+
+	if tokenPatch.Status == common.TokenStatusEnabled {
+		if cleanToken.Status == common.TokenStatusExpired &&
+			cleanToken.ExpiredTime <= common.GetTimestamp() &&
+			cleanToken.ExpiredTime != -1 {
 			c.JSON(http.StatusOK, gin.H{
 				"success": false,
 				"message": "令牌已过期，无法启用，请先修改令牌过期时间，或者设置为永不过期",
 			})
 			return
 		}
-		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
+		if cleanToken.Status == common.TokenStatusExhausted &&
+			cleanToken.RemainQuota <= 0 &&
+			!cleanToken.UnlimitedQuota {
 			c.JSON(http.StatusOK, gin.H{
 				"success": false,
 				"message": "令牌可用额度已用尽，无法启用，请先修改令牌剩余额度，或者设置为无限额度",
@@ -203,22 +225,34 @@ func UpdateToken(c *gin.Context) {
 		}
 	}
 	if statusOnly != "" {
-		cleanToken.Status = token.Status
+		cleanToken.Status = tokenPatch.Status
 	} else {
-		// If you add more fields, please also update token.Update()
-		cleanToken.Name = token.Name
-		cleanToken.ExpiredTime = token.ExpiredTime
-		cleanToken.RemainQuota = token.RemainQuota
-		cleanToken.UnlimitedQuota = token.UnlimitedQuota
+		// If you add more fields, please also update tokenPatch.Update()
+		if tokenPatch.Name != nil {
+			cleanToken.Name = *tokenPatch.Name
+		}
+		if tokenPatch.ExpiredTime != nil {
+			cleanToken.ExpiredTime = *tokenPatch.ExpiredTime
+		}
+		if tokenPatch.RemainQuota != nil {
+			cleanToken.RemainQuota = *tokenPatch.RemainQuota
+		}
+		if tokenPatch.UnlimitedQuota != nil {
+			cleanToken.UnlimitedQuota = *tokenPatch.UnlimitedQuota
+		}
 	}
-	err = cleanToken.Update()
-	if err != nil {
+
+	cleanToken.RemainQuota += tokenPatch.AddRemainQuota
+	cleanToken.UsedQuota += tokenPatch.AddUsedQuota
+
+	if err = cleanToken.Update(); err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
-			"message": err.Error(),
+			"message": "update token: " + err.Error(),
 		})
 		return
 	}
+
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",