From 54f76bd913d87878bfad5bc9d0d5f69537facdee Mon Sep 17 00:00:00 2001 From: Yvan <8574526@qq,com> Date: Sat, 22 Aug 2026 10:49:48 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/wire_gen.go | 15 ++- internal/app/catalog.go | 8 +- internal/integration/mq/emqx.go | 55 +++++++--- internal/integration/mq/emqx_test.go | 67 +++++++++++- internal/server/gin.go | 29 +++--- internal/server/middleware/access_log.go | 26 ++++- internal/server/middleware/auth.go | 126 ++++++++++++++++------- internal/server/middleware/auth_test.go | 82 +++++++++++++++ internal/server/websocket_auth_test.go | 74 +++++++++++++ pkg/module/module.go | 22 +++- pkg/mq/mqtt.go | 95 ++++++++++++++--- 11 files changed, 501 insertions(+), 98 deletions(-) create mode 100644 internal/server/middleware/auth_test.go create mode 100644 internal/server/websocket_auth_test.go diff --git a/cmd/wire_gen.go b/cmd/wire_gen.go index 40b26f6..41901ee 100644 --- a/cmd/wire_gen.go +++ b/cmd/wire_gen.go @@ -156,20 +156,25 @@ func wireApp(confServer *conf.Server, runtime *conf.Runtime, logger *slog.Logger v := handler.NewSet(authority, menu, api, permission, organization, announcement, handlerEmail, handlerPayment, task, media, audit, export, version, dictionary, parameter, apiToken, systemConfig, public, user, navigation, session, integrationConfig) routes := router.NewRoutes(v) taskMethods := worker.NewTaskMethods(taskUsecase, mediaUsecase, runtime) - moduleRuntime := app.Runtime(routes, taskMethods, registry) - websocketServer, cleanup2, err := websocket.New(store) + mqReloadable, cleanup2, err := mq.New(store, logger) if err != nil { cleanup() return nil, nil, err } - engine := server.NewGinEngineWithRuntime(runtime, accessControlService, authService, securityService, auditRecorder, logger, string2, moduleRuntime, websocketServer) - httpServer := server.NewGinServer(confServer, engine) - mqReloadable, cleanup3, err := mq.New(store, logger) + moduleRuntime, err := app.Runtime(routes, taskMethods, registry, catalog, mqReloadable) if err != nil { cleanup2() cleanup() return nil, nil, err } + websocketServer, cleanup3, err := websocket.New(store) + if err != nil { + cleanup2() + cleanup() + return nil, nil, err + } + engine := server.NewGinEngineWithRuntime(runtime, accessControlService, authService, securityService, auditRecorder, logger, string2, moduleRuntime, websocketServer) + httpServer := server.NewGinServer(confServer, engine) kratosApp := newApp(logger, httpServer, taskScheduler, auditRecorder, reloadableLogger, mqReloadable) return kratosApp, func() { cleanup3() diff --git a/internal/app/catalog.go b/internal/app/catalog.go index 594305e..634f1eb 100644 --- a/internal/app/catalog.go +++ b/internal/app/catalog.go @@ -7,6 +7,7 @@ import ( systemrouter "kra/internal/server/router" systemworker "kra/internal/worker" "kra/pkg/module" + "kra/pkg/mq" platformtask "kra/pkg/task" ) @@ -27,7 +28,10 @@ func TaskRegistry(catalog module.Catalog) *platformtask.Registry { } // Runtime composes HTTP route contributors from the enabled modules. -func Runtime(systemRoutes *systemrouter.Routes, systemTasks *systemworker.TaskMethods, registry *platformtask.Registry) *module.Runtime { +func Runtime(systemRoutes *systemrouter.Routes, systemTasks *systemworker.TaskMethods, registry *platformtask.Registry, catalog module.Catalog, subscriptions mq.SubscriptionRegistrar) (*module.Runtime, error) { platformtask.Apply(registry, systemTasks) - return module.NewRuntime(systemRoutes) + if err := mq.ApplySubscriptions(subscriptions, catalog.SubscriptionContributors()...); err != nil { + return nil, err + } + return module.NewRuntime(systemRoutes), nil } diff --git a/internal/integration/mq/emqx.go b/internal/integration/mq/emqx.go index 2a46be7..38abf32 100644 --- a/internal/integration/mq/emqx.go +++ b/internal/integration/mq/emqx.go @@ -402,7 +402,7 @@ func (r *Reloadable) clientLocked(provider string) platformmq.Client { return r.clients[provider] } -func (r *Reloadable) reconcileProviderLocked(provider string) error { +func (r *Reloadable) reconcileProviderLocked(provider string) (err error) { client := r.clientLocked(provider) desired := r.desiredSubscriptions(provider) if client == nil || !client.Connected() { @@ -419,6 +419,10 @@ func (r *Reloadable) reconcileProviderLocked(provider string) error { current[topic] = qos } r.mu.RUnlock() + // Broker operations can partially succeed. Persist every successful + // change even when a later operation fails, otherwise the next retry will + // repeat stale unbinds and may never reach the remaining subscriptions. + defer func() { r.setBindings(provider, current) }() for topic := range current { if _, exists := desired[topic]; exists { continue @@ -436,20 +440,31 @@ func (r *Reloadable) reconcileProviderLocked(provider string) error { if err := client.Unsubscribe(context.Background(), topic); err != nil { return err } + delete(current, topic) } if err := client.Subscribe(context.Background(), topic, qos, r.dispatcher(provider, topic)); err != nil { return err } current[topic] = qos } - r.mu.Lock() - r.bindings[provider] = current - r.mu.Unlock() delete(r.pending, provider) delete(r.nextRetry, provider) return nil } +func (r *Reloadable) setBindings(provider string, bindings map[string]byte) { + copyOfBindings := make(map[string]byte, len(bindings)) + for topic, qos := range bindings { + copyOfBindings[topic] = qos + } + r.mu.Lock() + if r.bindings == nil { + r.bindings = make(map[string]map[string]byte) + } + r.bindings[provider] = copyOfBindings + r.mu.Unlock() +} + func (r *Reloadable) Register(set platformmq.SubscriptionSet) error { normalized, err := platformmq.NormalizeSubscriptionSet(set) if err != nil { @@ -526,6 +541,9 @@ func (r *Reloadable) Publish(ctx context.Context, topic string, payload []byte, } func (r *Reloadable) PublishTo(ctx context.Context, provider, topic string, payload []byte, qos byte, retain bool) error { + if ctx == nil { + ctx = context.Background() + } provider = strings.ToLower(strings.TrimSpace(provider)) client := r.clientLocked(provider) if client == nil { @@ -539,6 +557,9 @@ func (r *Reloadable) Subscribe(ctx context.Context, topic string, qos byte, hand } func (r *Reloadable) SubscribeTo(ctx context.Context, provider, topic string, qos byte, handler platformmq.Handler) error { + if ctx == nil { + ctx = context.Background() + } if ctx != nil { select { case <-ctx.Done(): @@ -555,6 +576,9 @@ func (r *Reloadable) SubscribeTo(ctx context.Context, provider, topic string, qo return err } if !r.ConnectedTo(provider) { + // The legacy API reports unavailable when no live broker exists. Do not + // leave a durable declaration behind when that call failed. + _ = r.Unregister(owner) return platformmq.ErrUnavailable } return nil @@ -565,6 +589,9 @@ func (r *Reloadable) Unsubscribe(ctx context.Context, topics ...string) error { } func (r *Reloadable) UnsubscribeFrom(ctx context.Context, provider string, topics ...string) error { + if ctx == nil { + ctx = context.Background() + } if ctx != nil { select { case <-ctx.Done(): @@ -660,22 +687,24 @@ func (r *Reloadable) Close() error { } func (r *Reloadable) ensureStateLocked() { - if r.clients == nil { - r.clients = make(map[string]platformmq.Client) - } if r.configs == nil { r.configs = make(map[string]runtimeconfig.Config) } - if r.subscriptions == nil { - r.subscriptions = make(map[string]map[string]map[string]subscription) - } - if r.bindings == nil { - r.bindings = make(map[string]map[string]byte) - } if r.pending == nil { r.pending = make(map[string]bool) } if r.nextRetry == nil { r.nextRetry = make(map[string]time.Time) } + r.mu.Lock() + if r.clients == nil { + r.clients = make(map[string]platformmq.Client) + } + if r.subscriptions == nil { + r.subscriptions = make(map[string]map[string]map[string]subscription) + } + if r.bindings == nil { + r.bindings = make(map[string]map[string]byte) + } + r.mu.Unlock() } diff --git a/internal/integration/mq/emqx_test.go b/internal/integration/mq/emqx_test.go index f333c82..5d6bfc6 100644 --- a/internal/integration/mq/emqx_test.go +++ b/internal/integration/mq/emqx_test.go @@ -2,6 +2,7 @@ package mq import ( "context" + "errors" "log/slog" "testing" "time" @@ -11,14 +12,19 @@ import ( ) type fakeClient struct { - subscribed []string - unsubscribed []string - handlers map[string]platformmq.Handler + subscribed []string + unsubscribed []string + handlers map[string]platformmq.Handler + subscribeErrors map[string]error + unsubscribeErrors map[string]error } func (*fakeClient) Publish(context.Context, string, []byte, byte, bool) error { return nil } func (f *fakeClient) Subscribe(_ context.Context, topic string, _ byte, handler platformmq.Handler) error { f.subscribed = append(f.subscribed, topic) + if err := f.subscribeErrors[topic]; err != nil { + return err + } if f.handlers == nil { f.handlers = make(map[string]platformmq.Handler) } @@ -27,6 +33,11 @@ func (f *fakeClient) Subscribe(_ context.Context, topic string, _ byte, handler } func (f *fakeClient) Unsubscribe(_ context.Context, topics ...string) error { f.unsubscribed = append(f.unsubscribed, topics...) + for _, topic := range topics { + if err := f.unsubscribeErrors[topic]; err != nil { + return err + } + } return nil } func (*fakeClient) Connected() bool { return true } @@ -148,3 +159,53 @@ func TestReloadableDispatchesSameTopicToMultipleOwners(t *testing.T) { t.Fatal("audit handler was not called") } } + +func TestReloadablePersistsPartialBindingChangesWhenReconcileFails(t *testing.T) { + client := &fakeClient{subscribeErrors: map[string]error{"events.new": errors.New("subscribe failed")}} + r := &Reloadable{ + clients: map[string]platformmq.Client{ProviderEMQX: client}, + subscriptions: map[string]map[string]map[string]subscription{ + ProviderEMQX: { + "events.keep": {"owner": {qos: platformmq.AtMostOnce, handler: func(context.Context, platformmq.Message) {}}}, + "events.new": {"owner": {qos: platformmq.AtMostOnce, handler: func(context.Context, platformmq.Message) {}}}, + }, + }, + bindings: map[string]map[string]byte{ProviderEMQX: { + "events.old": platformmq.AtMostOnce, + "events.keep": platformmq.AtMostOnce, + }}, + pending: make(map[string]bool), + nextRetry: make(map[string]time.Time), + } + + if err := r.reconcileProviderLocked(ProviderEMQX); err == nil { + t.Fatal("reconcileProviderLocked() error = nil, want subscribe failure") + } + if _, exists := r.bindings[ProviderEMQX]["events.old"]; exists { + t.Fatal("successfully removed binding remained after failed reconcile") + } + if _, exists := r.bindings[ProviderEMQX]["events.keep"]; !exists { + t.Fatal("unchanged binding was lost after failed reconcile") + } + if _, exists := r.bindings[ProviderEMQX]["events.new"]; exists { + t.Fatal("failed subscription was recorded as bound") + } +} + +func TestReloadableLegacySubscribeDoesNotLeaveOfflineDeclaration(t *testing.T) { + r := &Reloadable{ + clients: make(map[string]platformmq.Client), + configs: map[string]runtimeconfig.Config{ProviderEMQX: {Kind: "mq", Provider: ProviderEMQX, Enabled: true}}, + subscriptions: make(map[string]map[string]map[string]subscription), + bindings: make(map[string]map[string]byte), + pending: make(map[string]bool), + nextRetry: make(map[string]time.Time), + } + err := r.Subscribe(context.Background(), "events.offline", platformmq.AtLeastOnce, func(context.Context, platformmq.Message) {}) + if !errors.Is(err, platformmq.ErrUnavailable) { + t.Fatalf("Subscribe() error = %v, want ErrUnavailable", err) + } + if len(r.subscriptions[ProviderEMQX]) != 0 { + t.Fatalf("offline legacy declaration remained: %#v", r.subscriptions[ProviderEMQX]) + } +} diff --git a/internal/server/gin.go b/internal/server/gin.go index 27e089d..1eef31a 100644 --- a/internal/server/gin.go +++ b/internal/server/gin.go @@ -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) { diff --git a/internal/server/middleware/access_log.go b/internal/server/middleware/access_log.go index afb8f12..b32b733 100644 --- a/internal/server/middleware/access_log.go +++ b/internal/server/middleware/access_log.go @@ -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) diff --git a/internal/server/middleware/auth.go b/internal/server/middleware/auth.go index 2ca661d..1084b38 100644 --- a/internal/server/middleware/auth.go +++ b/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) diff --git a/internal/server/middleware/auth_test.go b/internal/server/middleware/auth_test.go new file mode 100644 index 0000000..b9a51b1 --- /dev/null +++ b/internal/server/middleware/auth_test.go @@ -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) + } +} diff --git a/internal/server/websocket_auth_test.go b/internal/server/websocket_auth_test.go new file mode 100644 index 0000000..68ebc53 --- /dev/null +++ b/internal/server/websocket_auth_test.go @@ -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) diff --git a/pkg/module/module.go b/pkg/module/module.go index 432943e..ebf730a 100644 --- a/pkg/module/module.go +++ b/pkg/module/module.go @@ -6,6 +6,7 @@ package module import ( "github.com/gin-gonic/gin" "kra/pkg/database/migration" + "kra/pkg/mq" "kra/pkg/task" ) @@ -29,11 +30,12 @@ type TimedTask struct { } type Definition struct { - Name string - Migrations []migration.Step - Surface Surface - TimedTasks []TimedTask - Tasks []task.Method + Name string + Migrations []migration.Step + Surface Surface + TimedTasks []TimedTask + Tasks []task.Method + SubscriptionContributors []mq.SubscriptionContributor } type Catalog struct { @@ -73,6 +75,16 @@ func (c Catalog) TaskMethods() []task.Method { return methods } +// SubscriptionContributors returns the module declarations that are applied +// to the process-wide MQ runtime during application startup. +func (c Catalog) SubscriptionContributors() []mq.SubscriptionContributor { + var contributors []mq.SubscriptionContributor + for _, item := range c.Definitions { + contributors = append(contributors, item.SubscriptionContributors...) + } + return contributors +} + type RouteRegistrar interface { RegisterRoutes(public, private *gin.RouterGroup, engine *gin.Engine) } diff --git a/pkg/mq/mqtt.go b/pkg/mq/mqtt.go index e9c090e..08ffe75 100644 --- a/pkg/mq/mqtt.go +++ b/pkg/mq/mqtt.go @@ -4,15 +4,17 @@ import ( "context" "fmt" "sync" + "sync/atomic" "time" paho "github.com/eclipse/paho.mqtt.golang" ) type MQTT struct { - client paho.Client - mu sync.RWMutex - closed bool + client paho.Client + mu sync.RWMutex + closed bool + reconnecting atomic.Bool } func NewMQTT(cfg Config) (*MQTT, error) { @@ -31,19 +33,29 @@ func NewMQTT(cfg Config) (*MQTT, error) { if cfg.ReconnectInterval <= 0 { cfg.ReconnectInterval = 5 * time.Second } + c := &MQTT{} opts := paho.NewClientOptions().AddBroker(cfg.Broker).SetClientID(cfg.ClientID).SetUsername(cfg.Username).SetPassword(cfg.Password) opts.SetKeepAlive(cfg.KeepAlive).SetCleanSession(cfg.CleanSession).SetConnectTimeout(cfg.ConnectTimeout).SetAutoReconnect(true) opts.SetMaxReconnectInterval(cfg.ReconnectInterval).SetResumeSubs(true).SetOrderMatters(false) - c := &MQTT{} - token := paho.NewClient(opts) - connect := token.Connect() + opts.SetConnectionNotificationHandler(func(_ paho.Client, notification paho.ConnectionNotification) { + switch notification.Type() { + case paho.ConnectionNotificationTypeConnecting, paho.ConnectionNotificationTypeLost: + c.reconnecting.Store(true) + case paho.ConnectionNotificationTypeConnected: + c.reconnecting.Store(false) + } + }) + client := paho.NewClient(opts) + connect := client.Connect() if !connect.WaitTimeout(cfg.ConnectTimeout) { + client.Disconnect(0) return nil, fmt.Errorf("connect mqtt: %w", ErrUnavailable) } if err := connect.Error(); err != nil { + client.Disconnect(0) return nil, err } - c.client = token + c.client = client return c, nil } @@ -54,7 +66,9 @@ func (c *MQTT) Publish(ctx context.Context, topic string, payload []byte, qos by if qos > ExactlyOnce { return fmt.Errorf("invalid mqtt qos %d", qos) } - if c == nil || c.client == nil || !c.client.IsConnected() { + ctx = nonNilContext(ctx) + client, err := c.activeClient() + if err != nil || client == nil || !client.IsConnected() { return ErrUnavailable } select { @@ -62,7 +76,7 @@ func (c *MQTT) Publish(ctx context.Context, topic string, payload []byte, qos by return ctx.Err() default: } - return waitToken(ctx, c.client.Publish(topic, qos, retain, payload)) + return waitToken(ctx, client.Publish(topic, qos, retain, payload)) } func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error { @@ -72,7 +86,9 @@ func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Ha if qos > ExactlyOnce { return fmt.Errorf("invalid mqtt qos %d", qos) } - if c == nil || c.client == nil || !c.client.IsConnected() { + ctx = nonNilContext(ctx) + client, err := c.activeClient() + if err != nil || client == nil || !client.IsConnected() { return ErrUnavailable } select { @@ -83,7 +99,7 @@ func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Ha if handler == nil { return fmt.Errorf("mqtt handler is nil") } - return waitToken(ctx, c.client.Subscribe(topic, qos, func(_ paho.Client, msg paho.Message) { + return waitToken(ctx, client.Subscribe(topic, qos, func(_ paho.Client, msg paho.Message) { handler(context.Background(), Message{Topic: msg.Topic(), Payload: append([]byte(nil), msg.Payload()...), QoS: msg.Qos(), Retain: msg.Retained()}) })) } @@ -92,7 +108,9 @@ func (c *MQTT) Unsubscribe(ctx context.Context, topics ...string) error { if len(topics) == 0 { return fmt.Errorf("mqtt topics are empty") } - if c == nil || c.client == nil || !c.client.IsConnected() { + ctx = nonNilContext(ctx) + client, err := c.activeClient() + if err != nil || client == nil || !client.IsConnected() { return ErrUnavailable } select { @@ -100,10 +118,11 @@ func (c *MQTT) Unsubscribe(ctx context.Context, topics ...string) error { return ctx.Err() default: } - return waitToken(ctx, c.client.Unsubscribe(topics...)) + return waitToken(ctx, client.Unsubscribe(topics...)) } func waitToken(ctx context.Context, token paho.Token) error { + ctx = nonNilContext(ctx) if token == nil { return ErrUnavailable } @@ -114,16 +133,58 @@ func waitToken(ctx context.Context, token paho.Token) error { return token.Error() } } -func (c *MQTT) Connected() bool { return c != nil && c.client != nil && c.client.IsConnected() } + +func nonNilContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return ctx +} + +func (c *MQTT) activeClient() (paho.Client, error) { + if c == nil { + return nil, ErrUnavailable + } + c.mu.RLock() + client := c.client + closed := c.closed + c.mu.RUnlock() + if closed || client == nil { + return nil, ErrUnavailable + } + return client, nil +} + +func (c *MQTT) Connected() bool { + client, err := c.activeClient() + return err == nil && client.IsConnected() +} + +// Reconnecting reports whether the MQTT driver is recovering its connection. +// The reloadable integration waits for this state instead of replacing a +// client while driver-level recovery is still in progress. +func (c *MQTT) Reconnecting() bool { + if c == nil { + return false + } + c.mu.RLock() + closed := c.closed + c.mu.RUnlock() + return !closed && c.reconnecting.Load() +} + func (c *MQTT) Close() error { c.mu.Lock() - defer c.mu.Unlock() if c.closed { + c.mu.Unlock() return nil } c.closed = true - if c.client != nil && c.client.IsConnected() { - c.client.Disconnect(250) + client := c.client + c.client = nil + c.mu.Unlock() + if client != nil { + client.Disconnect(250) } return nil }