Relay done but not working

middleware/auth.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"one-api/common"
 	"one-api/model"
+	"strings"
 )
 
 func authHelper(c *gin.Context, minRole int) {
@@ -14,34 +15,13 @@ func authHelper(c *gin.Context, minRole int) {
 	role := session.Get("role")
 	id := session.Get("id")
 	status := session.Get("status")
-	authByToken := false
 	if username == nil {
-		// Check token
-		token := c.Request.Header.Get("Authorization")
-		if token == "" {
-			c.JSON(http.StatusOK, gin.H{
-				"success": false,
-				"message": "无权进行此操作，未登录或 token 无效",
-			})
-			c.Abort()
-			return
-		}
-		user := model.ValidateUserToken(token)
-		if user != nil && user.Username != "" {
-			// Token is valid
-			username = user.Username
-			role = user.Role
-			id = user.Id
-			status = user.Status
-		} else {
-			c.JSON(http.StatusOK, gin.H{
-				"success": false,
-				"message": "无权进行此操作，token 无效",
-			})
-			c.Abort()
-			return
-		}
-		authByToken = true
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "无权进行此操作，未登录",
+		})
+		c.Abort()
+		return
 	}
 	if status.(int) == common.UserStatusDisabled {
 		c.JSON(http.StatusOK, gin.H{
@@ -62,7 +42,6 @@ func authHelper(c *gin.Context, minRole int) {
 	c.Set("username", username)
 	c.Set("role", role)
 	c.Set("id", id)
-	c.Set("authByToken", authByToken)
 	c.Next()
 }
 
@@ -84,33 +63,25 @@ func RootAuth() func(c *gin.Context) {
 	}
 }
 
-// NoTokenAuth You should always use this after normal auth middlewares.
-func NoTokenAuth() func(c *gin.Context) {
+func TokenAuth() func(c *gin.Context) {
 	return func(c *gin.Context) {
-		authByToken := c.GetBool("authByToken")
-		if authByToken {
+		key := c.Request.Header.Get("Authorization")
+		parts := strings.Split(key, "-")
+		key = parts[0]
+		token, err := model.ValidateUserToken(key)
+		if err != nil {
 			c.JSON(http.StatusOK, gin.H{
-				"success": false,
-				"message": "本接口不支持使用 token 进行验证",
+				"error": gin.H{
+					"message": err.Error(),
+					"type":    "one_api_error",
+				},
 			})
 			c.Abort()
 			return
 		}
-		c.Next()
-	}
-}
-
-// TokenOnlyAuth You should always use this after normal auth middlewares.
-func TokenOnlyAuth() func(c *gin.Context) {
-	return func(c *gin.Context) {
-		authByToken := c.GetBool("authByToken")
-		if !authByToken {
-			c.JSON(http.StatusOK, gin.H{
-				"success": false,
-				"message": "本接口仅支持使用 token 进行验证",
-			})
-			c.Abort()
-			return
+		c.Set("id", token.UserId)
+		if len(parts) > 1 {
+			c.Set("channelId", parts[1])
 		}
 		c.Next()
 	}

middleware/distributor.go

@@ -0,0 +1,68 @@
+package middleware
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"net/http"
+	"one-api/common"
+	"one-api/model"
+	"strconv"
+)
+
+func Distribute() func(c *gin.Context) {
+	return func(c *gin.Context) {
+		var channel *model.Channel
+		channelId, ok := c.Get("channelId")
+		if ok {
+			id, err := strconv.Atoi(channelId.(string))
+			if err != nil {
+				c.JSON(http.StatusOK, gin.H{
+					"error": gin.H{
+						"message": "无效的渠道 ID",
+						"type":    "one_api_error",
+					},
+				})
+				c.Abort()
+				return
+			}
+			channel, err = model.GetChannelById(id, true)
+			if err != nil {
+				c.JSON(200, gin.H{
+					"error": gin.H{
+						"message": "无效的渠道 ID",
+						"type":    "one_api_error",
+					},
+				})
+				c.Abort()
+				return
+			}
+			if channel.Status != common.ChannelStatusEnabled {
+				c.JSON(200, gin.H{
+					"error": gin.H{
+						"message": "该渠道已被禁用",
+						"type":    "one_api_error",
+					},
+				})
+				c.Abort()
+				return
+			}
+		} else {
+			// Select a channel for the user
+			var err error
+			channel, err = model.GetRandomChannel()
+			if err != nil {
+				c.JSON(200, gin.H{
+					"error": gin.H{
+						"message": "无可用渠道",
+						"type":    "one_api_error",
+					},
+				})
+				c.Abort()
+				return
+			}
+		}
+		c.Set("channel", channel.Type)
+		c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
+		c.Next()
+	}
+}