feat: 令牌分组

middleware/auth.go

@@ -176,6 +176,7 @@ func TokenAuth() func(c *gin.Context) {
 			c.Set("token_model_limit_enabled", false)
 		}
 		c.Set("allow_ips", token.GetIpLimitsMap())
+		c.Set("token_group", token.Group)
 		if len(parts) > 1 {
 			if model.IsAdmin(token.UserId) {
 				c.Set("specific_channel_id", parts[1])

middleware/distributor.go

@@ -39,6 +39,15 @@ func Distribute() func(c *gin.Context) {
 			return
 		}
 		userGroup, _ := model.CacheGetUserGroup(userId)
+		tokenGroup := c.GetString("token_group")
+		if tokenGroup != "" {
+			// check group in common.GroupRatio
+			if _, ok := common.GroupRatio[tokenGroup]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被禁用", tokenGroup))
+				return
+			}
+			userGroup = tokenGroup
+		}
 		c.Set("group", userGroup)
 		if ok {
 			id, err := strconv.Atoi(channelId.(string))