优化结构

2026-08-22 10:49:48 +08:00
parent 2b78b195c1
commit 54f76bd913
11 changed files with 501 additions and 98 deletions
@@ -26,7 +26,7 @@ func NewGinEngine(runtime *conf.Runtime, access *service.AccessControlService, h
 	return NewGinEngineWithRuntime(runtime, access, auth, security, audit, logger, version, platformmodule.NewRuntime(router.NewRoutes(handlers)), nil)
 }

-func NewGinEngineWithRuntime(runtime *conf.Runtime, access *service.AccessControlService, auth *service.AuthService, security *service.SecurityService, audit *service.AuditRecorder, logger *slog.Logger, version string, routes *platformmodule.Runtime, ws *websocket.Server) *gin.Engine {
+func NewGinEngineWithRuntime(runtime *conf.Runtime, access *service.AccessControlService, auth middleware.TokenAuthenticator, security *service.SecurityService, audit *service.AuditRecorder, logger *slog.Logger, version string, routes *platformmodule.Runtime, ws *websocket.Server) *gin.Engine {
 	gin.SetMode(gin.ReleaseMode)
 	engine := gin.New()
 	if err := engine.SetTrustedProxies(nil); err != nil && logger != nil {
@@ -51,17 +51,20 @@ func NewGinEngineWithRuntime(runtime *conf.Runtime, access *service.AccessContro
 	if routes != nil {
 		routes.RegisterRoutes(public, private, engine)
 	}
+	handleWebSocket := func(c *gin.Context) {
+		if ws == nil || !ws.Enabled() || c.Request.URL.Path != ws.Path() {
+			c.Status(http.StatusNotFound)
+			return
+		}
+		if !middleware.AuthenticateWebSocket(c, auth) {
+			return
+		}
+		if err := ws.HandleRequest(c.Writer, c.Request); err != nil && logger != nil {
+			logger.Warn("websocket request failed", "mod", "websocket", "error", err)
+		}
+	}
 	if ws != nil && ws.Enabled() {
-		path := ws.Path()
-		engine.GET(path, func(c *gin.Context) {
-			if !ws.Enabled() || c.Request.URL.Path != ws.Path() {
-				c.Status(http.StatusNotFound)
-				return
-			}
-			if err := ws.HandleRequest(c.Writer, c.Request); err != nil && logger != nil {
-				logger.Warn("websocket request failed", "mod", "websocket", "error", err)
-			}
-		})
+		engine.GET(ws.Path(), handleWebSocket)
 	}
 	registerSwagger(engine, prefix, version, logger)
 	registerLocalStorage(engine, runtime)
@@ -74,9 +77,7 @@ func NewGinEngineWithRuntime(runtime *conf.Runtime, access *service.AccessContro

 	engine.NoRoute(func(c *gin.Context) {
 		if ws != nil && ws.Enabled() && c.Request.Method == http.MethodGet && c.Request.URL.Path == ws.Path() {
-			if err := ws.HandleRequest(c.Writer, c.Request); err != nil && logger != nil {
-				logger.Warn("websocket request failed", "mod", "websocket", "error", err)
-			}
+			handleWebSocket(c)
 			return
 		}
 		if serveLocalStorage(c, runtime) {
@@ -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)
+	}
+}
@@ -0,0 +1,74 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"kra/internal/biz"
+	"kra/internal/conf"
+	"kra/internal/integration/runtimeconfig"
+	websocketintegration "kra/internal/integration/websocket"
+	"kra/internal/server/middleware"
+	platformws "kra/pkg/websocket"
+
+	gorillawebsocket "github.com/gorilla/websocket"
+)
+
+type websocketAuthStub struct {
+	claims *biz.AuthClaims
+	got    string
+}
+
+func (s *websocketAuthStub) AuthenticateToken(_ context.Context, token string) (*biz.TokenAuthentication, error) {
+	s.got = token
+	return &biz.TokenAuthentication{Claims: s.claims}, nil
+}
+
+var _ middleware.TokenAuthenticator = (*websocketAuthStub)(nil)
+
+func TestWebSocketRouteRequiresAndAcceptsLoginToken(t *testing.T) {
+	store := runtimeconfig.NewStore()
+	raw, err := json.Marshal(map[string]any{"path": "/ws"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	store.Set(runtimeconfig.Config{Kind: "websocket", Provider: websocketintegration.ProviderMelody, Enabled: true, Values: raw})
+	ws, cleanup, err := websocketintegration.New(store)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cleanup()
+
+	auth := &websocketAuthStub{claims: &biz.AuthClaims{ID: 7}}
+	engine := NewGinEngineWithRuntime(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, auth, nil, nil, nil, "test", nil, ws)
+
+	unauthorized := httptest.NewRecorder()
+	engine.ServeHTTP(unauthorized, httptest.NewRequest(http.MethodGet, "/ws", nil))
+	if unauthorized.Code != http.StatusUnauthorized {
+		t.Fatalf("unauthorized websocket status=%d body=%s", unauthorized.Code, unauthorized.Body.String())
+	}
+
+	httpServer := httptest.NewServer(engine)
+	defer httpServer.Close()
+	wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws?token=login-token"
+	connection, response, err := gorillawebsocket.DefaultDialer.Dial(wsURL, nil)
+	if response != nil && response.Body != nil {
+		_ = response.Body.Close()
+	}
+	if err != nil {
+		t.Fatalf("authorized websocket handshake failed: %v", err)
+	}
+	if connection == nil {
+		t.Fatal("authorized websocket handshake returned nil connection")
+	}
+	_ = connection.Close()
+	if auth.got != "login-token" {
+		t.Fatalf("websocket auth token=%q, want login-token", auth.got)
+	}
+}
+
+var _ platformws.Hub = (*websocketintegration.Server)(nil)