fix: prevent common user from specifying channel id (#12)

This commit is contained in:
JustSong 2023-04-26 14:49:27 +08:00
parent 1dd92a3f92
commit 4f8cbd643d
3 changed files with 26 additions and 1 deletions

View File

@ -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。
不加的话将会使用负载均衡的方式使用多个渠道。

View File

@ -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()
}

View File

@ -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
}