341 lines
8.8 KiB
Go
341 lines
8.8 KiB
Go
package controller
|
||
|
||
import (
|
||
"fmt"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/QuantumNous/new-api/common"
|
||
"github.com/QuantumNous/new-api/i18n"
|
||
"github.com/QuantumNous/new-api/model"
|
||
"github.com/QuantumNous/new-api/service"
|
||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
func buildMaskedTokenResponse(token *model.Token) *model.Token {
|
||
if token == nil {
|
||
return nil
|
||
}
|
||
maskedToken := *token
|
||
maskedToken.Key = token.GetMaskedKey()
|
||
return &maskedToken
|
||
}
|
||
|
||
func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token {
|
||
maskedTokens := make([]*model.Token, 0, len(tokens))
|
||
for _, token := range tokens {
|
||
maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token))
|
||
}
|
||
return maskedTokens
|
||
}
|
||
|
||
func GetAllTokens(c *gin.Context) {
|
||
userId := c.GetInt("id")
|
||
pageInfo := common.GetPageQuery(c)
|
||
tokens, err := model.GetAllUserTokens(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize())
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
total, _ := model.CountUserTokens(userId)
|
||
pageInfo.SetTotal(int(total))
|
||
pageInfo.SetItems(buildMaskedTokenResponses(tokens))
|
||
common.ApiSuccess(c, pageInfo)
|
||
}
|
||
|
||
func SearchTokens(c *gin.Context) {
|
||
userId := c.GetInt("id")
|
||
keyword := c.Query("keyword")
|
||
token := c.Query("token")
|
||
|
||
pageInfo := common.GetPageQuery(c)
|
||
|
||
tokens, total, err := model.SearchUserTokens(userId, keyword, token, pageInfo.GetStartIdx(), pageInfo.GetPageSize())
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
pageInfo.SetTotal(int(total))
|
||
pageInfo.SetItems(buildMaskedTokenResponses(tokens))
|
||
common.ApiSuccess(c, pageInfo)
|
||
}
|
||
|
||
func GetToken(c *gin.Context) {
|
||
id, err := strconv.Atoi(c.Param("id"))
|
||
userId := c.GetInt("id")
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
token, err := model.GetTokenByIds(id, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
common.ApiSuccess(c, buildMaskedTokenResponse(token))
|
||
}
|
||
|
||
func GetTokenKey(c *gin.Context) {
|
||
id, err := strconv.Atoi(c.Param("id"))
|
||
userId := c.GetInt("id")
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
token, err := model.GetTokenByIds(id, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
common.ApiSuccess(c, gin.H{
|
||
"key": token.GetFullKey(),
|
||
})
|
||
}
|
||
|
||
func GetTokenStatus(c *gin.Context) {
|
||
tokenId := c.GetInt("token_id")
|
||
userId := c.GetInt("id")
|
||
token, err := model.GetTokenByIds(tokenId, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
expiredAt := token.ExpiredTime
|
||
if expiredAt == -1 {
|
||
expiredAt = 0
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "credit_summary",
|
||
"total_granted": token.RemainQuota,
|
||
"total_used": 0, // not supported currently
|
||
"total_available": token.RemainQuota,
|
||
"expires_at": expiredAt * 1000,
|
||
})
|
||
}
|
||
|
||
func GetTokenUsage(c *gin.Context) {
|
||
authHeader := c.GetHeader("Authorization")
|
||
if authHeader == "" {
|
||
c.JSON(http.StatusUnauthorized, gin.H{
|
||
"success": false,
|
||
"message": "No Authorization header",
|
||
})
|
||
return
|
||
}
|
||
|
||
parts := strings.Split(authHeader, " ")
|
||
if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
|
||
c.JSON(http.StatusUnauthorized, gin.H{
|
||
"success": false,
|
||
"message": "Invalid Bearer token",
|
||
})
|
||
return
|
||
}
|
||
tokenKey := parts[1]
|
||
|
||
token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false)
|
||
if err != nil {
|
||
common.SysError("failed to get token by key: " + err.Error())
|
||
common.ApiErrorI18n(c, i18n.MsgTokenGetInfoFailed)
|
||
return
|
||
}
|
||
|
||
expiredAt := token.ExpiredTime
|
||
if expiredAt == -1 {
|
||
expiredAt = 0
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"code": true,
|
||
"message": "ok",
|
||
"data": gin.H{
|
||
"object": "token_usage",
|
||
"name": token.Name,
|
||
"total_granted": token.RemainQuota + token.UsedQuota,
|
||
"total_used": token.UsedQuota,
|
||
"total_available": token.RemainQuota,
|
||
"unlimited_quota": token.UnlimitedQuota,
|
||
"model_limits": token.GetModelLimitsMap(),
|
||
"model_limits_enabled": token.ModelLimitsEnabled,
|
||
"expires_at": expiredAt,
|
||
},
|
||
})
|
||
}
|
||
|
||
func AddToken(c *gin.Context) {
|
||
token := model.Token{}
|
||
err := c.ShouldBindJSON(&token)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
if len(token.Name) > 50 {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
|
||
return
|
||
}
|
||
// 非无限额度时，检查额度值是否超出有效范围
|
||
if !token.UnlimitedQuota {
|
||
if token.RemainQuota < 0 {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative)
|
||
return
|
||
}
|
||
maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
|
||
if token.RemainQuota > maxQuotaValue {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue})
|
||
return
|
||
}
|
||
}
|
||
// 检查用户令牌数量是否已达上限
|
||
maxTokens := operation_setting.GetMaxUserTokens()
|
||
count, err := model.CountUserTokens(c.GetInt("id"))
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
if int(count) >= maxTokens {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": fmt.Sprintf("已达到最大令牌数量限制 (%d)", maxTokens),
|
||
})
|
||
return
|
||
}
|
||
key, err := common.GenerateKey()
|
||
if err != nil {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed)
|
||
common.SysLog("failed to generate token key: " + err.Error())
|
||
return
|
||
}
|
||
cleanToken := model.Token{
|
||
UserId: c.GetInt("id"),
|
||
Name: token.Name,
|
||
Key: key,
|
||
CreatedTime: common.GetTimestamp(),
|
||
AccessedTime: common.GetTimestamp(),
|
||
ExpiredTime: token.ExpiredTime,
|
||
RemainQuota: token.RemainQuota,
|
||
UnlimitedQuota: token.UnlimitedQuota,
|
||
ModelLimitsEnabled: token.ModelLimitsEnabled,
|
||
ModelLimits: token.ModelLimits,
|
||
AllowIps: token.AllowIps,
|
||
Group: token.Group,
|
||
CrossGroupRetry: token.CrossGroupRetry,
|
||
}
|
||
err = cleanToken.Insert()
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
service.RecordCreateOperation(c, "token", cleanToken.Id, cleanToken.Name, "创建令牌："+cleanToken.Name, "")
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
})
|
||
}
|
||
|
||
func DeleteToken(c *gin.Context) {
|
||
id, _ := strconv.Atoi(c.Param("id"))
|
||
userId := c.GetInt("id")
|
||
err := model.DeleteTokenById(id, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
service.RecordDeleteOperation(c, "token", id, "", fmt.Sprintf("删除令牌 (ID：%d)", id), "")
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
})
|
||
}
|
||
|
||
func UpdateToken(c *gin.Context) {
|
||
userId := c.GetInt("id")
|
||
statusOnly := c.Query("status_only")
|
||
token := model.Token{}
|
||
err := c.ShouldBindJSON(&token)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
if len(token.Name) > 50 {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
|
||
return
|
||
}
|
||
if !token.UnlimitedQuota {
|
||
if token.RemainQuota < 0 {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative)
|
||
return
|
||
}
|
||
maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
|
||
if token.RemainQuota > maxQuotaValue {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue})
|
||
return
|
||
}
|
||
}
|
||
cleanToken, err := model.GetTokenByIds(token.Id, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
if token.Status == common.TokenStatusEnabled {
|
||
if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenExpiredCannotEnable)
|
||
return
|
||
}
|
||
if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
|
||
common.ApiErrorI18n(c, i18n.MsgTokenExhaustedCannotEable)
|
||
return
|
||
}
|
||
}
|
||
if statusOnly != "" {
|
||
cleanToken.Status = token.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
|
||
cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
|
||
cleanToken.ModelLimits = token.ModelLimits
|
||
cleanToken.AllowIps = token.AllowIps
|
||
cleanToken.Group = token.Group
|
||
cleanToken.CrossGroupRetry = token.CrossGroupRetry
|
||
}
|
||
err = cleanToken.Update()
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
service.RecordUpdateOperation(c, "token", cleanToken.Id, cleanToken.Name, "更新令牌："+cleanToken.Name, "")
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": buildMaskedTokenResponse(cleanToken),
|
||
})
|
||
}
|
||
|
||
type TokenBatch struct {
|
||
Ids []int `json:"ids"`
|
||
}
|
||
|
||
func DeleteTokenBatch(c *gin.Context) {
|
||
tokenBatch := TokenBatch{}
|
||
if err := c.ShouldBindJSON(&tokenBatch); err != nil || len(tokenBatch.Ids) == 0 {
|
||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||
return
|
||
}
|
||
userId := c.GetInt("id")
|
||
count, err := model.BatchDeleteTokens(tokenBatch.Ids, userId)
|
||
if err != nil {
|
||
common.ApiError(c, err)
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": count,
|
||
})
|
||
}
|