75 lines
2.2 KiB
Go
75 lines
2.2 KiB
Go
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)
|