优化结构

internal/server/middleware/auth.go

@@ -1,13 +1,13 @@
 package middleware
 
 import (
+	"context"
 	"errors"
 	"net/http"
 	"strconv"
 	"strings"
 
 	"kra/internal/biz"
-	"kra/internal/service"
 
 	"github.com/gin-gonic/gin"
 	"golang.org/x/sync/singleflight"
@@ -17,48 +17,98 @@ const claimsKey = "admin_claims"
 
 var refreshTokens singleflight.Group
 
-func Auth(auth *service.AuthService) gin.HandlerFunc {
+type TokenAuthenticator interface {
+	AuthenticateToken(context.Context, string) (*biz.TokenAuthentication, error)
+}
+
+func Auth(auth TokenAuthenticator) gin.HandlerFunc {
 	return func(c *gin.Context) {
-		token := c.GetHeader("x-token")
-		if token == "" {
-			token, _ = c.Cookie("x-token")
+		if authenticate(c, auth, false) {
+			c.Next()
 		}
-		if token == "" {
-			NoAuth(c, "未登录或非法访问，请登录")
-			return
-		}
-		value, err, _ := refreshTokens.Do(token, func() (any, error) {
-			return auth.AuthenticateToken(c.Request.Context(), token)
-		})
-		if err != nil {
-			message := "无法处理此token"
-			switch {
-			case errors.Is(err, biz.ErrTokenExpired):
-				message = "登录已过期，请重新登录"
-			case errors.Is(err, biz.ErrTokenMalformed):
-				message = "这不是一个token"
-			case errors.Is(err, biz.ErrTokenSignatureInvalid):
-				message = "无效签名"
-			case errors.Is(err, biz.ErrTokenNotValidYet):
-				message = "token尚未激活"
-			case errors.Is(err, biz.ErrTokenDisabled):
-				message = "您的帐户异地登陆或令牌失效"
-			}
-			SetTokenCookie(c, "", -1)
-			NoAuth(c, message)
-			return
-		}
-		authentication := value.(*biz.TokenAuthentication)
-		if authentication.Refreshed != nil {
-			c.Header("new-token", authentication.Refreshed.Value)
-			c.Header("new-expires-at", strconv.FormatInt(authentication.Refreshed.ExpiresAt.Unix(), 10))
-			SetTokenCookie(c, authentication.Refreshed.Value, int(authentication.Refreshed.TTL.Seconds()))
-		}
-		c.Set(claimsKey, authentication.Claims)
-		c.Next()
 	}
 }
 
+// AuthenticateWebSocket validates the same login token as the HTTP auth
+// middleware. Query-string tokens are accepted only for WebSocket handshakes,
+// because browsers cannot attach a custom x-token header to WebSocket.connect.
+func AuthenticateWebSocket(c *gin.Context, auth TokenAuthenticator) bool {
+	return authenticate(c, auth, true)
+}
+
+func authenticate(c *gin.Context, auth TokenAuthenticator, allowQueryToken bool) bool {
+	token := requestToken(c, allowQueryToken)
+	if token == "" {
+		NoAuth(c, "未登录或非法访问，请登录")
+		return false
+	}
+	if auth == nil {
+		NoAuth(c, "认证服务不可用")
+		return false
+	}
+	value, err, _ := refreshTokens.Do(token, func() (any, error) {
+		return auth.AuthenticateToken(c.Request.Context(), token)
+	})
+	if err != nil {
+		SetTokenCookie(c, "", -1)
+		NoAuth(c, tokenErrorMessage(err))
+		return false
+	}
+	authentication, ok := value.(*biz.TokenAuthentication)
+	if !ok || authentication == nil || authentication.Claims == nil {
+		SetTokenCookie(c, "", -1)
+		NoAuth(c, "无法处理此token")
+		return false
+	}
+	if authentication.Refreshed != nil {
+		c.Header("new-token", authentication.Refreshed.Value)
+		c.Header("new-expires-at", strconv.FormatInt(authentication.Refreshed.ExpiresAt.Unix(), 10))
+		SetTokenCookie(c, authentication.Refreshed.Value, int(authentication.Refreshed.TTL.Seconds()))
+	}
+	c.Set(claimsKey, authentication.Claims)
+	return true
+}
+
+func requestToken(c *gin.Context, allowQueryToken bool) string {
+	token := strings.TrimSpace(c.GetHeader("x-token"))
+	if token == "" {
+		authorization := strings.TrimSpace(c.GetHeader("Authorization"))
+		if len(authorization) > len("Bearer ") && strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
+			token = strings.TrimSpace(authorization[len("Bearer "):])
+		}
+	}
+	if token == "" {
+		token, _ = c.Cookie("x-token")
+		token = strings.TrimSpace(token)
+	}
+	if token == "" && allowQueryToken {
+		for _, key := range []string{"token", "access_token"} {
+			token = strings.TrimSpace(c.Query(key))
+			if token != "" {
+				break
+			}
+		}
+	}
+	return token
+}
+
+func tokenErrorMessage(err error) string {
+	message := "无法处理此token"
+	switch {
+	case errors.Is(err, biz.ErrTokenExpired):
+		message = "登录已过期，请重新登录"
+	case errors.Is(err, biz.ErrTokenMalformed):
+		message = "这不是一个token"
+	case errors.Is(err, biz.ErrTokenSignatureInvalid):
+		message = "无效签名"
+	case errors.Is(err, biz.ErrTokenNotValidYet):
+		message = "token尚未激活"
+	case errors.Is(err, biz.ErrTokenDisabled):
+		message = "您的帐户异地登陆或令牌失效"
+	}
+	return message
+}
+
 func Claims(c *gin.Context) *biz.AuthClaims {
 	value, _ := c.Get(claimsKey)
 	claims, _ := value.(*biz.AuthClaims)