package controller import ( "fmt" "net/http" "strconv" "github.com/gin-gonic/gin" "github.com/jinzhu/copier" "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/ctxkey" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/network" "github.com/songquanpeng/one-api/common/random" "github.com/songquanpeng/one-api/model" ) func GetRequestCost(c *gin.Context) { reqId := c.Param("request_id") if reqId == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "request_id should not be empty", }) return } docu, err := model.GetCostByRequestId(reqId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, docu) } func GetAllTokens(c *gin.Context) { userId := c.GetInt(ctxkey.Id) p, _ := strconv.Atoi(c.Query("p")) if p < 0 { p = 0 } order := c.Query("order") tokens, err := model.GetAllUserTokens(userId, p*config.ItemsPerPage, config.ItemsPerPage, order) 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": tokens, }) return } func SearchTokens(c *gin.Context) { userId := c.GetInt(ctxkey.Id) keyword := c.Query("keyword") tokens, err := model.SearchUserTokens(userId, keyword) 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": tokens, }) return } func GetToken(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt(ctxkey.Id) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } token, err := model.GetTokenByIds(id, userId) 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": token, }) return } func GetTokenStatus(c *gin.Context) { tokenId := c.GetInt(ctxkey.TokenId) userId := c.GetInt(ctxkey.Id) token, err := model.GetTokenByIds(tokenId, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) 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 validateToken(_ *gin.Context, token *model.Token) error { if len(token.Name) > 30 { return fmt.Errorf("Token name is too long") } if token.Subnet != nil && *token.Subnet != "" { err := network.IsValidSubnets(*token.Subnet) if err != nil { return fmt.Errorf("invalid network segment: %s", err.Error()) } } return nil } func AddToken(c *gin.Context) { token := new(model.Token) err := c.ShouldBindJSON(token) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } err = validateToken(c, token) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": fmt.Sprintf("invalid token: %s", err.Error()), }) return } cleanToken := model.Token{ UserId: c.GetInt(ctxkey.Id), Name: token.Name, Key: random.GenerateKey(), CreatedTime: helper.GetTimestamp(), AccessedTime: helper.GetTimestamp(), ExpiredTime: token.ExpiredTime, RemainQuota: token.RemainQuota, UnlimitedQuota: token.UnlimitedQuota, Models: token.Models, Subnet: token.Subnet, } err = cleanToken.Insert() 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": cleanToken, }) return } func DeleteToken(c *gin.Context) { id, _ := strconv.Atoi(c.Param("id")) userId := c.GetInt(ctxkey.Id) err := model.DeleteTokenById(id, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return } type consumeTokenRequest struct { // AddUsedQuota add or subtract used quota from another source AddUsedQuota uint `json:"add_used_quota" gorm:"-"` // AddReason is the reason for adding or subtracting used quota AddReason string `json:"add_reason" gorm:"-"` } // ConsumeToken consume token from another source, // let one-api to gather billing from different sources. func ConsumeToken(c *gin.Context) { tokenPatch := new(consumeTokenRequest) err := c.ShouldBindJSON(tokenPatch) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } userID := c.GetInt(ctxkey.Id) tokenID := c.GetInt(ctxkey.TokenId) cleanToken, err := model.GetTokenByIds(tokenID, userID) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if cleanToken.Status != model.TokenStatusEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "API KeysNot enabled", }) return } if cleanToken.Status == model.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The token has expired and cannot be enabled. Please modify the expiration time of the token, or set it to never expire.", }) return } if cleanToken.Status == model.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The available quota of the token has been used up and cannot be enabled. Please modify the remaining quota of the token, or set it to unlimited quota", }) return } // let admin to add or subtract used quota, // make it possible to aggregate billings from different sources. if cleanToken.RemainQuota < int64(tokenPatch.AddUsedQuota) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "Insufficient remaining quota", }) return } if err = model.DecreaseTokenQuota(cleanToken.Id, int64(tokenPatch.AddUsedQuota)); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } model.RecordConsumeLog(c.Request.Context(), &model.Log{ UserId: userID, ModelName: tokenPatch.AddReason, TokenName: cleanToken.Name, Quota: int(tokenPatch.AddUsedQuota), Content: fmt.Sprintf("External (%s) consumed %s", tokenPatch.AddReason, common.LogQuota(int64(tokenPatch.AddUsedQuota))), }) err = cleanToken.Update() 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": cleanToken, }) return } func UpdateToken(c *gin.Context) { userId := c.GetInt(ctxkey.Id) statusOnly := c.Query("status_only") tokenPatch := new(model.Token) err := c.ShouldBindJSON(tokenPatch) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } token := new(model.Token) if err = copier.Copy(token, tokenPatch); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } err = validateToken(c, token) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": fmt.Sprintf("invalid token: %s", err.Error()), }) return } cleanToken, err := model.GetTokenByIds(token.Id, userId) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } switch token.Status { case model.TokenStatusEnabled: if cleanToken.Status == model.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 && token.ExpiredTime != -1 && token.ExpiredTime < helper.GetTimestamp() { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The token has expired and cannot be enabled. Please modify the expiration time of the token, or set it to never expire.", }) return } if cleanToken.Status == model.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota && token.RemainQuota <= 0 && !token.UnlimitedQuota { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The available quota of the token has been used up and cannot be enabled. Please modify the remaining quota of the token, or set it to unlimited quota", }) return } case model.TokenStatusExhausted: if token.RemainQuota > 0 || token.UnlimitedQuota { token.Status = model.TokenStatusEnabled } case model.TokenStatusExpired: if token.ExpiredTime == -1 || token.ExpiredTime > helper.GetTimestamp() { token.Status = model.TokenStatusEnabled } } 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.UnlimitedQuota = token.UnlimitedQuota cleanToken.Models = token.Models cleanToken.Subnet = token.Subnet cleanToken.RemainQuota = token.RemainQuota cleanToken.Status = token.Status } err = cleanToken.Update() 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": cleanToken, }) return }