perler-beads-server/internal/server/websocket_test.go

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