package server import ( "context" "crypto/sha256" "encoding/hex" "encoding/json" "github.com/google/uuid" "github.com/gorilla/websocket" "github.com/redis/go-redis/v9" "net/http" "net/http/httptest" "os" "perler-beads-server/internal/domain" "strings" "testing" "time" ) func wsTestServer(t *testing.T) (*Server, *httptest.Server) { t.Helper() dsn := os.Getenv("TEST_REDIS_URL") if dsn == "" { t.Skip("TEST_REDIS_URL absent; real Redis websocket integration skipped") } options, e := redis.ParseURL(dsn) if e != nil { t.Fatal(e) } client := redis.NewClient(options) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() if e = client.Ping(ctx).Err(); e != nil { t.Fatal(e) } s := &Server{Redis: client, Origins: []string{"https://staff.example.test"}} httpServer := httptest.NewServer(s.Router()) t.Cleanup(func() { httpServer.Close(); client.Close() }) return s, httpServer } func wsToken(t *testing.T, s *Server, role, store string) string { t.Helper() token := uuid.NewString() h := sha256.Sum256([]byte(token)) key := "session:" + hex.EncodeToString(h[:]) b, _ := json.Marshal(domain.User{ID: uuid.NewString(), Role: role, StoreID: store}) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() if e := s.Redis.Set(ctx, key, b, time.Minute).Err(); e != nil { t.Fatal(e) } t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() s.Redis.Del(ctx, key) }) return token } func wsConnect(t *testing.T, h *httptest.Server, origin string) *websocket.Conn { t.Helper() headers := http.Header{} if origin != "" { headers.Set("Origin", origin) } c, resp, e := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(h.URL, "http")+"/api/ws", headers) if e != nil { if resp != nil { t.Fatalf("websocket upgrade: %v (status %d)", e, resp.StatusCode) } t.Fatal(e) } t.Cleanup(func() { c.Close() }) return c } func wsAuth(t *testing.T, c *websocket.Conn, token string) { t.Helper() c.SetReadDeadline(time.Now().Add(5 * time.Second)) if e := c.WriteJSON(map[string]string{"type": "auth", "token": token}); e != nil { t.Fatal(e) } var v map[string]any if e := c.ReadJSON(&v); e != nil { t.Fatal(e) } if v["type"] != "ready" || len(v) != 1 { t.Fatalf("bad ready: %+v", v) } } func TestWebSocketAuthenticationAndOrigin(t *testing.T) { s, h := wsTestServer(t) token := wsToken(t, s, "admin", "s1") wsAuth(t, wsConnect(t, h, h.URL), token) wsAuth(t, wsConnect(t, h, "https://staff.example.test"), token) for _, origin := range []string{"https://evil.example", "null", h.URL + "/path", ""} { headers := http.Header{} if origin != "" { headers.Set("Origin", origin) } conn, resp, e := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(h.URL, "http")+"/api/ws", headers) if conn != nil { conn.Close() } if e == nil || resp == nil || resp.StatusCode != http.StatusForbidden { t.Fatalf("origin %q should be forbidden: %v %+v", origin, e, resp) } } } func TestWebSocketGuestInvalidAuthAndTimeout(t *testing.T) { t.Parallel() s, h := wsTestServer(t) for _, token := range []string{"missing", wsToken(t, s, "guest", "s1"), wsToken(t, s, "admin", "")} { c := wsConnect(t, h, h.URL) c.WriteJSON(map[string]string{"type": "auth", "token": token}) c.SetReadDeadline(time.Now().Add(6 * time.Second)) _, _, e := c.ReadMessage() if !websocket.IsCloseError(e, 4403) { t.Fatalf("auth must close4403: %v", e) } } c := wsConnect(t, h, h.URL) c.SetReadDeadline(time.Now().Add(7 * time.Second)) _, _, e := c.ReadMessage() if !websocket.IsCloseError(e, 4403) { t.Fatalf("auth timeout must close4403: %v", e) } } func TestWebSocketLivesPastHTTPDeadline(t *testing.T) { t.Parallel() s, h := wsTestServer(t) c := wsConnect(t, h, h.URL) wsAuth(t, c, wsToken(t, s, "staff", "s1")) c.SetReadDeadline(time.Now().Add(24 * time.Second)) var event map[string]any if e := c.ReadJSON(&event); e != nil { t.Fatal(e) } if event["type"] != "heartbeat" { t.Fatalf("expected heartbeat beyond15s, got%+v", event) } } func TestWebSocketRevokedSessionCloses(t *testing.T) { t.Parallel() s, h := wsTestServer(t) token := wsToken(t, s, "admin", "s1") c := wsConnect(t, h, h.URL) wsAuth(t, c, token) hash := sha256.Sum256([]byte(token)) s.Redis.Del(context.Background(), "session:"+hex.EncodeToString(hash[:])) c.SetReadDeadline(time.Now().Add(25 * time.Second)) for { _, _, e := c.ReadMessage() if e != nil { if !websocket.IsCloseError(e, 4403) { t.Fatalf("revocation should close4403: %v", e) } break } } } func TestWebSocketOversizeAndProtocol(t *testing.T) { s, h := wsTestServer(t) for _, payload := range []string{strings.Repeat("x", 4097), `{"type":"other","token":"missing"}`} { c := wsConnect(t, h, h.URL) c.WriteMessage(websocket.TextMessage, []byte(payload)) c.SetReadDeadline(time.Now().Add(6 * time.Second)) _, _, e := c.ReadMessage() if e == nil { t.Fatal("invalid client message accepted") } } c := wsConnect(t, h, h.URL) wsAuth(t, c, wsToken(t, s, "admin", "s1")) c.WriteJSON(map[string]string{"type": "auth", "token": "anything"}) c.SetReadDeadline(time.Now().Add(6 * time.Second)) _, _, e := c.ReadMessage() if !websocket.IsCloseError(e, 4400) { t.Fatalf("reauth must close4400: %v", e) } }