yvan
/
kra-new
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化结构

This commit is contained in:
Yvan 2026-08-22 10:49:48 +08:00
parent 2b78b195c1
commit 54f76bd913
11 changed files with 501 additions and 98 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
cmd/wire_gen.go generated
View File

@ -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()

8
internal/app/catalog.go
View File

@ -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
}

55
internal/integration/mq/emqx.go
View File

@ -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()
}

61
internal/integration/mq/emqx_test.go
View File

@ -2,6 +2,7 @@ package mq
import (
"context"
"errors"
"log/slog"
"testing"
"time"
@ -14,11 +15,16 @@ type fakeClient struct {
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])
}
}

19
internal/server/gin.go
View File

@ -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)
}
if ws != nil && ws.Enabled() {
path := ws.Path()
engine.GET(path, func(c *gin.Context) {
if !ws.Enabled() || c.Request.URL.Path != ws.Path() {
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() {
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) {

26
internal/server/middleware/access_log.go
View File

@ -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)

88
internal/server/middleware/auth.go
View File

@ -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,20 +17,82 @@ 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()
}
}
}
// 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
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):
@ -44,19 +106,7 @@ func Auth(auth *service.AuthService) gin.HandlerFunc {
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()
}
return message
}
func Claims(c *gin.Context) *biz.AuthClaims {

82
internal/server/middleware/auth_test.go Normal file
View File

@ -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)
}
}

74
internal/server/websocket_auth_test.go Normal file
View File

@ -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)

12
pkg/module/module.go
View File

@ -6,6 +6,7 @@ package module
import (
"github.com/gin-gonic/gin"
"kra/pkg/database/migration"
"kra/pkg/mq"
"kra/pkg/task"
)
@ -34,6 +35,7 @@ type Definition struct {
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)
}

89
pkg/mq/mqtt.go
View File

@ -4,6 +4,7 @@ import (
"context"
"fmt"
"sync"
"sync/atomic"
"time"
paho "github.com/eclipse/paho.mqtt.golang"
@ -13,6 +14,7 @@ type MQTT struct {
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
}