优化结构

2026-08-22 10:49:48 +08:00
parent 2b78b195c1
commit 54f76bd913
11 changed files with 501 additions and 98 deletions
@@ -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)