优化结构

2026-08-22 10:49:48 +08:00
parent 2b78b195c1
commit 54f76bd913
11 changed files with 501 additions and 98 deletions
@@ -9,6 +9,7 @@ import (
 	"log/slog"
 	"mime"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 	"time"
@@ -131,7 +132,7 @@ func AccessLog(runtime *conf.Runtime, logger *slog.Logger, version string) gin.H
 				"request_id", stringValueFromContext(c, "request_id"), "trace_id", stringValueFromContext(c, "trace_id"),
 				"bytes_in", bytesIn, "bytes_out", bytesOut, "user_id", userID, "authority_id", authorityID,
 				"error", c.Writer.Status() >= http.StatusInternalServerError || privateErrors != "", "ua", c.Request.UserAgent(),
-				"req_query", c.Request.URL.RawQuery,
+				"req_query", redactQuery(c.Request.URL.RawQuery),
 			}
 		}
 		if !paymentCallback && config != nil && config.Zap != nil && config.Zap.AccessReqHeaders {
@@ -218,6 +219,29 @@ func redactHeaders(headers map[string][]string) map[string]string {
 	return out
 }

+func redactQuery(raw string) string {
+	if strings.TrimSpace(raw) == "" {
+		return ""
+	}
+	values, _ := url.ParseQuery(raw)
+	for key := range values {
+		if sensitiveQueryKey(key) {
+			values[key] = []string{"***"}
+		}
+	}
+	return values.Encode()
+}
+
+func sensitiveQueryKey(key string) bool {
+	normalized := strings.ToLower(strings.ReplaceAll(strings.ReplaceAll(key, "_", ""), "-", ""))
+	switch normalized {
+	case "token", "accesstoken", "refreshtoken", "authorization", "apikey", "secret":
+		return true
+	default:
+		return strings.HasSuffix(normalized, "token")
+	}
+}
+
 func stringValueFromContext(c *gin.Context, key string) string {
 	value, _ := c.Get(key)
 	return stringValue(value)
@@ -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)
@@ -0,0 +1,82 @@
+package middleware
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"kra/internal/biz"
+
+	"github.com/gin-gonic/gin"
+)
+
+type tokenAuthenticatorStub struct {
+	claims *biz.AuthClaims
+	err    error
+	got    string
+}
+
+func (s *tokenAuthenticatorStub) AuthenticateToken(_ context.Context, token string) (*biz.TokenAuthentication, error) {
+	s.got = token
+	if s.err != nil {
+		return nil, s.err
+	}
+	return &biz.TokenAuthentication{Claims: s.claims}, nil
+}
+
+func TestAuthenticateWebSocketAcceptsQueryToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	recorder := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(recorder)
+	context.Request = httptest.NewRequest(http.MethodGet, "/ws?token=query-token", nil)
+
+	if !AuthenticateWebSocket(context, stub) {
+		t.Fatalf("AuthenticateWebSocket() = false, body=%s", recorder.Body.String())
+	}
+	if stub.got != "query-token" {
+		t.Fatalf("authenticated token = %q, want query-token", stub.got)
+	}
+	if Claims(context) == nil || Claims(context).ID != 7 {
+		t.Fatalf("claims were not stored: %#v", Claims(context))
+	}
+}
+
+func TestHTTPAuthDoesNotAcceptQueryToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	engine := gin.New()
+	called := false
+	engine.GET("/protected", Auth(stub), func(c *gin.Context) { called = true })
+
+	response := httptest.NewRecorder()
+	engine.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/protected?token=query-token", nil))
+	if response.Code != http.StatusUnauthorized || called {
+		t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
+	}
+}
+
+func TestAuthenticateWebSocketSupportsBearerHeaderAndReportsInvalidToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	context, _ := gin.CreateTestContext(httptest.NewRecorder())
+	context.Request = httptest.NewRequest(http.MethodGet, "/ws", nil)
+	context.Request.Header.Set("Authorization", "Bearer header-token")
+
+	if !AuthenticateWebSocket(context, stub) || stub.got != "header-token" {
+		t.Fatalf("Bearer token was not accepted: ok=%v token=%q", Claims(context) != nil, stub.got)
+	}
+
+	invalid := &tokenAuthenticatorStub{err: errors.New("invalid token")}
+	invalidRecorder := httptest.NewRecorder()
+	invalidContext, _ := gin.CreateTestContext(invalidRecorder)
+	invalidContext.Request = httptest.NewRequest(http.MethodGet, "/ws?token=bad-token", nil)
+	if AuthenticateWebSocket(invalidContext, invalid) {
+		t.Fatal("invalid token was accepted")
+	}
+	if invalidRecorder.Code != http.StatusUnauthorized {
+		t.Fatalf("invalid token status=%d, want %d", invalidRecorder.Code, http.StatusUnauthorized)
+	}
+}