Просмотр исходного кода

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

JustSong 2 лет назад
Родитель
Сommit
4f8cbd643d
3 измененных файлов с 26 добавлено и 1 удалено
  1. 1 0
      README.md
  2. 12 1
      middleware/auth.go
  3. 13 0
      model/user.go

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

+ 12 - 1
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()
 	}

+ 13 - 0
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
+}