132 lines
4.0 KiB
Go
132 lines
4.0 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"kra/internal/biz/system"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"kra/internal/config"
|
|
"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 *system.AuthClaims
|
|
got string
|
|
}
|
|
|
|
func (s *websocketAuthStub) AuthenticateToken(_ context.Context, token string) (*system.TokenAuthentication, error) {
|
|
s.got = token
|
|
return &system.TokenAuthentication{Claims: s.claims}, nil
|
|
}
|
|
|
|
var _ middleware.TokenAuthenticator = (*websocketAuthStub)(nil)
|
|
|
|
func TestWebSocketRouteAcceptsAnonymousAndAuthenticatedViewers(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: &system.AuthClaims{ID: 7}}
|
|
engine := NewGinEngineWithRuntime(config.NewStore(&config.Config{Admin: &config.Admin{}}), nil, auth, nil, nil, nil, "test", nil, ws)
|
|
|
|
httpServer := httptest.NewServer(engine)
|
|
defer httpServer.Close()
|
|
baseURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws"
|
|
anonymous, response, err := gorillawebsocket.DefaultDialer.Dial(baseURL, nil)
|
|
if response != nil && response.Body != nil {
|
|
_ = response.Body.Close()
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("anonymous websocket handshake failed: %v", err)
|
|
}
|
|
if anonymous == nil {
|
|
t.Fatal("anonymous websocket handshake returned nil connection")
|
|
}
|
|
_ = anonymous.Close()
|
|
if auth.got != "" {
|
|
t.Fatalf("anonymous websocket unexpectedly authenticated token=%q", auth.got)
|
|
}
|
|
|
|
authenticated, response, err := gorillawebsocket.DefaultDialer.Dial(baseURL+"?token=login-token", nil)
|
|
if response != nil && response.Body != nil {
|
|
_ = response.Body.Close()
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("authenticated websocket handshake failed: %v", err)
|
|
}
|
|
if authenticated == nil {
|
|
t.Fatal("authenticated websocket handshake returned nil connection")
|
|
}
|
|
_ = authenticated.Close()
|
|
if auth.got != "login-token" {
|
|
t.Fatalf("websocket auth token=%q, want login-token", auth.got)
|
|
}
|
|
}
|
|
|
|
func TestWebSocketRouteSupportsRouterPrefix(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: &system.AuthClaims{ID: 7}}
|
|
runtime := config.NewStore(&config.Config{Admin: &config.Admin{RouterPrefix: "/admin"}})
|
|
engine := NewGinEngineWithRuntime(runtime, nil, auth, nil, nil, nil, "test", nil, ws)
|
|
httpServer := httptest.NewServer(engine)
|
|
defer httpServer.Close()
|
|
|
|
wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/admin/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("prefixed websocket handshake failed: %v", err)
|
|
}
|
|
if connection == nil {
|
|
t.Fatal("prefixed websocket handshake returned nil connection")
|
|
}
|
|
_ = connection.Close()
|
|
for _, route := range engine.Routes() {
|
|
if route.Method == http.MethodGet && route.Path == "/ws" {
|
|
t.Fatal("unprefixed websocket route should not be registered")
|
|
}
|
|
}
|
|
found := false
|
|
for _, route := range engine.Routes() {
|
|
if route.Method == http.MethodGet && route.Path == "/admin/ws" {
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
t.Fatal("prefixed websocket route was not registered")
|
|
}
|
|
}
|
|
|
|
var _ platformws.Hub = (*websocketintegration.Server)(nil)
|