diff --git a/README.md b/README.md index f7467d6c..250439f6 100644 --- a/README.md +++ b/README.md @@ -138,6 +138,7 @@ sudo service nginx restart 之后就可以使用你的令牌访问 One API 了,使用方式与 [OpenAI API](https://platform.openai.com/docs/api-reference/introduction) 一致。 可以通过在令牌后面添加渠道 ID 的方式指定使用哪一个渠道处理本次请求,例如:`Authorization: Bearer ONE_API_KEY-CHANNEL_ID`。 +注意,需要是管理员用户创建的令牌才能指定渠道 ID。 不加的话将会使用负载均衡的方式使用多个渠道。 diff --git a/middleware/auth.go b/middleware/auth.go index 6279d008..64fdb0a1 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -83,7 +83,18 @@ func TokenAuth() func(c *gin.Context) { c.Set("token_id", token.Id) c.Set("unlimited_times", token.UnlimitedTimes) if len(parts) > 1 { - c.Set("channelId", parts[1]) + if model.IsAdmin(token.UserId) { + c.Set("channelId", parts[1]) + } else { + c.JSON(http.StatusOK, gin.H{ + "error": gin.H{ + "message": "普通用户不支持指定渠道", + "type": "one_api_error", + }, + }) + c.Abort() + return + } } c.Next() } diff --git a/model/user.go b/model/user.go index 74fc593e..4ad3b682 100644 --- a/model/user.go +++ b/model/user.go @@ -175,3 +175,16 @@ func ResetUserPasswordByEmail(email string, password string) error { err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error return err } + +func IsAdmin(userId int) bool { + if userId == 0 { + return false + } + var user User + err := DB.Where("id = ?", userId).Select("role").Find(&user).Error + if err != nil { + common.SysError("No such user " + err.Error()) + return false + } + return user.Role >= common.RoleAdminUser +}