Files
OwnCord/Server/ws/ws_integration_test.go
T
J3vb 9116a880a3 feat(phase-bc): pass 5 — pgdbgen, postgres EventStore/PluginStore, plugin hub wiring, OTel stack, reconnect DB tier
- Generate Server/db/dbgen/{events,plugins}.sql.go and full Server/db/pgdbgen/ (//go:build postgres gated)
- Implement PostgresStore EventStore and PluginStore methods in store/postgres.go
- Wire plugin host_events.go EventSink into hub broadcast path (SetPluginEventSink)
- Wire host_commands.go slash-command dispatcher: chat_command V1 handler + hub.SetPluginRegistry
- Add handlers_command.go + handlers_command_test.go for plugin slash-command dispatch
- Add reconnect_db_test.go: TestReconnect_BufferMiss_FallsBackToDBTier (cold-tier DB replay)
- Add otel-up/otel-down Makefile targets; docker-compose.otel.yml + prometheus.dev.yml
- Update PHASE_BC_LOCAL_TODO.md: mark in-session items complete; document remaining network-blocked steps
- Minor fixes: channel_handler access-control, router plugin handler wiring, service span instrumentation
2026-04-06 22:48:59 +02:00

1301 lines
39 KiB
Go

package ws_test
// ws_integration_test.go covers ServeWS, authenticateConn, writePump, and
// readPump by spinning up a real httptest server and dialing it with the
// nhooyr.io/websocket client.
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"nhooyr.io/websocket"
"github.com/owncord/server/auth"
"github.com/owncord/server/service"
"github.com/owncord/server/store"
"github.com/owncord/server/ws"
)
// ─── ServeWS / authenticateConn happy path ────────────────────────────────────
// TestServeWS_InvalidUpgrade verifies that a plain HTTP GET (non-WS) returns
// a non-101 status without panicking.
func TestServeWS_InvalidUpgrade_ReturnsError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
// Plain GET without WebSocket upgrade headers should fail gracefully.
resp, err := http.Get(srv.URL)
if err != nil {
t.Fatalf("http.Get: %v", err)
}
defer func() { _ = resp.Body.Close() }()
// nhooyr.io/websocket returns 400 or 426 when upgrade is absent.
if resp.StatusCode == 200 {
t.Errorf("expected non-200 for plain HTTP, got %d", resp.StatusCode)
}
}
// ─── authenticateConn — error paths ──────────────────────────────────────────
// TestAuthenticateConn_NoAuthMessage verifies that a connection that closes
// immediately (without sending auth) causes the server to close it gracefully.
func TestAuthenticateConn_NoAuthMessage_ServerClosesConn(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, dialResp, err := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
defer dialResp.Body.Close() //nolint:errcheck // test cleanup
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
// Close without sending auth — the server's authDeadline (10s) will fire,
// but closing immediately should cause a read error on the server side.
_ = conn.Close(websocket.StatusNormalClosure, "no auth")
// Give the server a moment to react.
time.Sleep(50 * time.Millisecond)
// Hub should have no clients registered.
if hub.ClientCount() != 0 {
t.Errorf("ClientCount = %d after unauthenticated connection, want 0", hub.ClientCount())
}
}
// TestAuthenticateConn_InvalidJSON verifies that sending invalid JSON as the
// first message causes the server to send an auth_error and close.
func TestAuthenticateConn_InvalidJSON_ReceivesAuthError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, dialResp2, err := websocket.Dial(ctx, wsURL, nil)
if dialResp2 != nil && dialResp2.Body != nil {
defer dialResp2.Body.Close() //nolint:errcheck // test cleanup
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
// Send invalid JSON as first message.
if err := conn.Write(ctx, websocket.MessageText, []byte("NOT JSON")); err != nil {
t.Fatalf("write: %v", err)
}
// Server should respond with auth_error.
_, raw, readErr := conn.Read(ctx)
if readErr != nil {
// Server may close connection — also acceptable.
return
}
var msg map[string]any
if err := json.Unmarshal(raw, &msg); err == nil {
if msg["type"] == "auth_error" {
return // expected
}
t.Errorf("expected auth_error, got type=%q", msg["type"])
}
}
// TestAuthenticateConn_WrongMessageType verifies that sending a non-auth
// first message causes the server to send an auth_error.
func TestAuthenticateConn_WrongMessageType_ReceivesAuthError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
// Send a chat_send instead of auth.
wrongMsg := map[string]any{
"type": "chat_send",
"payload": map[string]string{"content": "hello"},
}
raw, _ := json.Marshal(wrongMsg)
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write: %v", err)
}
_, respRaw, readErr := conn.Read(ctx)
if readErr != nil {
return // server closed — acceptable
}
var msg map[string]any
if err := json.Unmarshal(respRaw, &msg); err == nil {
if msg["type"] == "auth_error" {
return // expected
}
t.Errorf("expected auth_error, got type=%q", msg["type"])
}
}
// TestAuthenticateConn_MissingToken verifies that an auth message without
// a token field receives an auth_error.
func TestAuthenticateConn_MissingToken_ReceivesAuthError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{}, // no token field
}
raw, _ := json.Marshal(authMsg)
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write: %v", err)
}
_, respRaw, readErr := conn.Read(ctx)
if readErr != nil {
return
}
var msg map[string]any
if err := json.Unmarshal(respRaw, &msg); err == nil {
if msg["type"] == "auth_error" {
return
}
t.Errorf("expected auth_error, got type=%q", msg["type"])
}
}
// TestAuthenticateConn_InvalidToken verifies that an auth message with a
// non-existent token receives an auth_error.
func TestAuthenticateConn_InvalidToken_ReceivesAuthError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": "totally-invalid-token-xyz"},
}
raw, _ := json.Marshal(authMsg)
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write: %v", err)
}
_, respRaw, readErr := conn.Read(ctx)
if readErr != nil {
return
}
var msg map[string]any
if err := json.Unmarshal(respRaw, &msg); err == nil {
if msg["type"] == "auth_error" {
return
}
t.Errorf("expected auth_error, got type=%q", msg["type"])
}
}
// TestServeWS_ValidAuth_FullHandshake verifies the complete happy path:
// valid token → auth_ok + ready received, client counted in hub.
func TestServeWS_ValidAuth_FullHandshake(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
// Seed user and session.
userID, err := database.CreateUser("ws-handshake-user", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
// Send auth.
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
}
raw, _ := json.Marshal(authMsg)
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write auth: %v", err)
}
// Expect auth_ok.
_, respRaw, err := conn.Read(ctx)
if err != nil {
t.Fatalf("read auth_ok: %v", err)
}
var authOK map[string]any
if err := json.Unmarshal(respRaw, &authOK); err != nil {
t.Fatalf("unmarshal auth_ok: %v", err)
}
if authOK["type"] != "auth_ok" {
t.Errorf("first response type = %q, want auth_ok", authOK["type"])
}
// Expect ready.
_, respRaw2, err := conn.Read(ctx)
if err != nil {
t.Fatalf("read ready: %v", err)
}
var readyMsg map[string]any
if err := json.Unmarshal(respRaw2, &readyMsg); err != nil {
t.Fatalf("unmarshal ready: %v", err)
}
if readyMsg["type"] != "ready" {
t.Errorf("second response type = %q, want ready", readyMsg["type"])
}
// Give hub a moment to register the client.
time.Sleep(30 * time.Millisecond)
if hub.ClientCount() != 1 {
t.Errorf("ClientCount = %d after successful auth, want 1", hub.ClientCount())
}
}
// TestServeWS_ImmediateDisconnect_DoesNotLeaveGhostClient verifies that a
// client that drops immediately after sending auth does not remain registered
// or stuck online.
func TestServeWS_ImmediateDisconnect_DoesNotLeaveGhostClient(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
userID, err := database.CreateUser("abruptclose", "hash", 4)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token := "abrupt-close-token"
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer conn.CloseNow() //nolint:errcheck
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
}
raw, err := json.Marshal(authMsg)
if err != nil {
t.Fatalf("marshal auth: %v", err)
}
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write auth: %v", err)
}
_ = conn.CloseNow()
deadline := time.Now().Add(2 * time.Second)
cleanedUp := false
for time.Now().Before(deadline) {
user, getErr := database.GetUserByID(userID)
if getErr != nil {
t.Fatalf("GetUserByID: %v", getErr)
}
if hub.ClientCount() == 0 && user.Status == "offline" {
cleanedUp = true
break
}
time.Sleep(20 * time.Millisecond)
}
if !cleanedUp {
user, getErr := database.GetUserByID(userID)
if getErr != nil {
t.Fatalf("GetUserByID final: %v", getErr)
}
t.Fatalf("immediate disconnect left stale state: client_count=%d user_status=%q", hub.ClientCount(), user.Status)
}
}
// TestServeWS_DuplicateLogin_KeepsUserOnline verifies that replacing an
// existing connection does not broadcast or persist a false offline state for
// the still-connected replacement session.
func TestServeWS_DuplicateLogin_KeepsUserOnline(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
userID, err := database.CreateUser("ws-reconnect-user", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
dialAndAuth := func() *websocket.Conn {
conn, dialResp, dialErr := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
dialResp.Body.Close()
}
if dialErr != nil {
t.Fatalf("websocket.Dial: %v", dialErr)
}
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
}
raw, marshalErr := json.Marshal(authMsg)
if marshalErr != nil {
t.Fatalf("marshal auth: %v", marshalErr)
}
if writeErr := conn.Write(ctx, websocket.MessageText, raw); writeErr != nil {
t.Fatalf("write auth: %v", writeErr)
}
for i := 0; i < 2; i++ {
if _, _, readErr := conn.Read(ctx); readErr != nil {
t.Fatalf("read handshake message %d: %v", i, readErr)
}
}
return conn
}
conn1 := dialAndAuth()
defer func() { _ = conn1.Close(websocket.StatusNormalClosure, "") }()
conn2 := dialAndAuth()
defer func() { _ = conn2.Close(websocket.StatusNormalClosure, "") }()
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
user, getErr := database.GetUserByID(userID)
if getErr != nil {
t.Fatalf("GetUserByID: %v", getErr)
}
if hub.ClientCount() == 1 && user.Status == "online" {
return
}
time.Sleep(20 * time.Millisecond)
}
user, getErr := database.GetUserByID(userID)
if getErr != nil {
t.Fatalf("GetUserByID final: %v", getErr)
}
t.Fatalf("duplicate login left wrong state: client_count=%d user_status=%q", hub.ClientCount(), user.Status)
}
// TestServeWS_Reconnect_PreservesVoiceState verifies that replacing a
// connection via network reconnect (last_seq > 0) preserves voice state:
// no voice_leave broadcast, voiceChID transferred, DB row intact.
func TestServeWS_Reconnect_PreservesVoiceState(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
userID, err := database.CreateUser("ws-voice-reconnect", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
// Create a voice channel.
chID, err := database.CreateChannel("voice-reconnect", "voice", "", "", 0)
if err != nil {
t.Fatalf("CreateChannel: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
dialAndAuth := func(lastSeq uint64) *websocket.Conn {
t.Helper()
conn, dialResp, dialErr := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
dialResp.Body.Close()
}
if dialErr != nil {
t.Fatalf("websocket.Dial: %v", dialErr)
}
authMsg := map[string]any{
"type": "auth",
"payload": map[string]any{"token": token, "last_seq": lastSeq},
}
raw, marshalErr := json.Marshal(authMsg)
if marshalErr != nil {
t.Fatalf("marshal auth: %v", marshalErr)
}
if writeErr := conn.Write(ctx, websocket.MessageText, raw); writeErr != nil {
t.Fatalf("write auth: %v", writeErr)
}
// Read auth_ok + ready
for i := 0; i < 2; i++ {
if _, _, readErr := conn.Read(ctx); readErr != nil {
t.Fatalf("read handshake message %d: %v", i, readErr)
}
}
return conn
}
// First connection: fresh
conn1 := dialAndAuth(0)
defer func() { _ = conn1.Close(websocket.StatusNormalClosure, "") }()
var originalClient *ws.Client
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
originalClient = hub.GetClient(userID)
if originalClient != nil {
break
}
time.Sleep(20 * time.Millisecond)
}
if originalClient == nil {
t.Fatal("expected first client to be registered")
}
// Simulate voice join AFTER conn1 is established — both in-memory and DB.
// (Setting it before conn1 would cause serve.go's fresh-connect cleanup
// to delete the DB row during conn1's handshake.)
if err := database.JoinVoiceChannel(userID, chID); err != nil {
t.Fatalf("JoinVoiceChannel: %v", err)
}
vsBeforeReconnect, err := database.GetVoiceState(userID)
if err != nil {
t.Fatalf("GetVoiceState(before reconnect): %v", err)
}
if vsBeforeReconnect == nil {
t.Fatal("expected voice state row after JoinVoiceChannel")
}
ws.SetClientVoiceStateForTest(originalClient, chID, vsBeforeReconnect.JoinedAt)
// Second connection: reconnect (lastSeq > oldestSeq) — voice state should transfer.
// Use lastSeq=2 because conn1's join produces at least 2 broadcasts
// (member_join seq=1, presence seq=2), and afterSeq must be > oldestSeq
// for EventsSince to return a replay instead of nil (BUG-085).
conn2 := dialAndAuth(2)
defer func() { _ = conn2.Close(websocket.StatusNormalClosure, "") }()
var replacementClient *ws.Client
deadline = time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
replacementClient = hub.GetClient(userID)
if replacementClient != nil && ws.GetClientVoiceChIDForTest(replacementClient) == chID {
break
}
time.Sleep(20 * time.Millisecond)
}
if replacementClient == nil {
t.Fatal("expected replacement client to be registered")
}
// Assert: voiceChID transferred
if got := ws.GetClientVoiceChIDForTest(replacementClient); got != chID {
t.Fatalf("replacement client voiceChID = %d, want %d", got, chID)
}
// Assert: DB row still intact
vs, vsErr := database.GetVoiceState(userID)
if vsErr != nil {
t.Fatalf("GetVoiceState: %v", vsErr)
}
if vs == nil {
t.Fatal("reconnect: DB voice_state row was deleted, expected it to be preserved")
}
if vs.ChannelID != chID {
t.Fatalf("reconnect: DB voice_state channel_id = %d, want %d", vs.ChannelID, chID)
}
// Assert: no voice_leave broadcast
readDeadline := time.Now().Add(400 * time.Millisecond)
for time.Now().Before(readDeadline) {
readCtx, readCancel := context.WithTimeout(ctx, 100*time.Millisecond)
_, raw, readErr := conn2.Read(readCtx)
readCancel()
if readErr != nil {
break
}
var msg map[string]any
if err := json.Unmarshal(raw, &msg); err != nil {
continue
}
if msg["type"] == "voice_leave" {
t.Fatalf("reconnect must not broadcast voice_leave: %s", string(raw))
}
}
}
// TestServeWS_FreshReconnect_CleansStaleVoiceState verifies that when a user
// presses F5 (fresh connection, lastSeq = 0) while in voice, the server:
// 1. cleans the DB voice_state row before building ready
// 2. does NOT include the user in ready.payload.voice_states
// 3. sets replacement client voiceChID = 0
// 4. broadcasts exactly one voice_leave visible to an observer client
func TestServeWS_FreshReconnect_CleansStaleVoiceState(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
// Create two users: the voice user who F5-reloads, and an observer.
userID, err := database.CreateUser("ws-voice-f5", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
observerID, err := database.CreateUser("ws-observer", "hash", 1)
if err != nil {
t.Fatalf("CreateUser (observer): %v", err)
}
obsToken, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken (observer): %v", err)
}
obsTokenHash := auth.HashToken(obsToken)
if _, err := database.CreateSession(observerID, obsTokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession (observer): %v", err)
}
// Create a voice channel for the user to be "in".
chID, err := database.CreateChannel("voice-test", "voice", "", "", 0)
if err != nil {
t.Fatalf("CreateChannel: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
dialAndAuthFresh := func(tok string) *websocket.Conn {
t.Helper()
conn, dialResp, dialErr := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
dialResp.Body.Close()
}
if dialErr != nil {
t.Fatalf("websocket.Dial: %v", dialErr)
}
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": tok},
}
raw, marshalErr := json.Marshal(authMsg)
if marshalErr != nil {
t.Fatalf("marshal auth: %v", marshalErr)
}
if writeErr := conn.Write(ctx, websocket.MessageText, raw); writeErr != nil {
t.Fatalf("write auth: %v", writeErr)
}
// Read auth_ok + ready
for i := 0; i < 2; i++ {
if _, _, readErr := conn.Read(ctx); readErr != nil {
t.Fatalf("read handshake message %d: %v", i, readErr)
}
}
return conn
}
// dialAndReadReady dials, authenticates with lastSeq=0, and returns the
// conn plus the parsed ready payload so the caller can inspect voice_states.
dialAndReadReady := func(tok string) (*websocket.Conn, map[string]any) {
t.Helper()
conn, dialResp, dialErr := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
dialResp.Body.Close()
}
if dialErr != nil {
t.Fatalf("websocket.Dial: %v", dialErr)
}
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": tok},
}
raw, marshalErr := json.Marshal(authMsg)
if marshalErr != nil {
t.Fatalf("marshal auth: %v", marshalErr)
}
if writeErr := conn.Write(ctx, websocket.MessageText, raw); writeErr != nil {
t.Fatalf("write auth: %v", writeErr)
}
// Read auth_ok (skip it)
if _, _, readErr := conn.Read(ctx); readErr != nil {
t.Fatalf("read auth_ok: %v", readErr)
}
// Read ready — parse it
_, readyRaw, readErr := conn.Read(ctx)
if readErr != nil {
t.Fatalf("read ready: %v", readErr)
}
var readyMsg map[string]any
if err := json.Unmarshal(readyRaw, &readyMsg); err != nil {
t.Fatalf("unmarshal ready: %v", err)
}
return conn, readyMsg
}
// First connection: user joins voice
conn1 := dialAndAuthFresh(token)
defer func() { _ = conn1.Close(websocket.StatusNormalClosure, "") }()
var originalClient *ws.Client
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
originalClient = hub.GetClient(userID)
if originalClient != nil {
break
}
time.Sleep(20 * time.Millisecond)
}
if originalClient == nil {
t.Fatal("expected first client to be registered")
}
// Simulate voice join — both in-memory and DB
if err := database.JoinVoiceChannel(userID, chID); err != nil {
t.Fatalf("JoinVoiceChannel: %v", err)
}
vsBeforeReload, err := database.GetVoiceState(userID)
if err != nil {
t.Fatalf("GetVoiceState(before reload): %v", err)
}
if vsBeforeReload == nil {
t.Fatal("expected voice state row after JoinVoiceChannel")
}
ws.SetClientVoiceStateForTest(originalClient, chID, vsBeforeReload.JoinedAt)
// Connect the observer (will receive broadcasts).
// Do NOT drain the observer — nhooyr.io/websocket closes the conn
// when a Read context expires. Instead, collect all messages below
// and filter for voice_leave in the assertion.
obsConn := dialAndAuthFresh(obsToken)
defer func() { _ = obsConn.Close(websocket.StatusNormalClosure, "") }()
// F5 reload: fresh connection (lastSeq = 0)
conn2, readyMsg := dialAndReadReady(token)
defer func() { _ = conn2.Close(websocket.StatusNormalClosure, "") }()
// Wait for replacement to register fully and broadcasts to propagate.
time.Sleep(500 * time.Millisecond)
// Assert 1: replacement client voiceChID == 0
replacementClient := hub.GetClient(userID)
if replacementClient == nil {
t.Fatal("expected replacement client to be registered")
}
if got := ws.GetClientVoiceChIDForTest(replacementClient); got != 0 {
t.Fatalf("fresh reconnect: replacement client voiceChID = %d, want 0", got)
}
// Assert 2: DB voice row is gone
vs, vsErr := database.GetVoiceState(userID)
if vsErr != nil {
t.Fatalf("GetVoiceState: %v", vsErr)
}
if vs != nil {
t.Fatalf("fresh reconnect: stale voice state still in DB: channel_id=%d", vs.ChannelID)
}
// Assert 3: ready.payload.voice_states does not include the reconnecting user
payload, _ := readyMsg["payload"].(map[string]any)
voiceStates, _ := payload["voice_states"].([]any)
for _, vsRaw := range voiceStates {
vsMap, _ := vsRaw.(map[string]any)
vsUserID, _ := vsMap["user_id"].(float64)
if int64(vsUserID) == userID {
t.Fatalf("ready payload must not include stale voice state for user %d: %+v", userID, vsMap)
}
}
// Assert 4: observer saw exactly one voice_leave for our user+channel.
// Read all pending messages — the observer may have received
// member_join/presence/voice_leave since connecting.
voiceLeaveCount := 0
for {
readCtx, readCancel := context.WithTimeout(ctx, 500*time.Millisecond)
_, raw, readErr := obsConn.Read(readCtx)
readCancel()
if readErr != nil {
break
}
var msg map[string]any
if err := json.Unmarshal(raw, &msg); err != nil {
continue
}
msgType, _ := msg["type"].(string)
if msgType == "voice_leave" {
msgPayload, _ := msg["payload"].(map[string]any)
msgUserID, _ := msgPayload["user_id"].(float64)
msgChID, _ := msgPayload["channel_id"].(float64)
if int64(msgUserID) == userID && int64(msgChID) == chID {
voiceLeaveCount++
}
}
}
if voiceLeaveCount == 0 {
t.Fatal("fresh reconnect: observer never received voice_leave for the ghost user")
}
if voiceLeaveCount > 1 {
t.Fatalf("fresh reconnect: observer received %d voice_leave messages, want exactly 1", voiceLeaveCount)
}
}
// TestServeWS_writePump_MessageDelivered verifies that messages queued on the
// hub are written through writePump to the connected client.
func TestServeWS_writePump_MessageDelivered(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
// Seed user and session.
userID, err := database.CreateUser("ws-pump-user", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
// Authenticate.
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
}
raw, _ := json.Marshal(authMsg)
_ = conn.Write(ctx, websocket.MessageText, raw)
// Drain auth_ok and ready.
for i := 0; i < 2; i++ {
_, _, err := conn.Read(ctx)
if err != nil {
t.Fatalf("drain initial messages: %v", err)
}
}
// Wait for client to be registered and then broadcast a server_restart.
time.Sleep(50 * time.Millisecond)
hub.BroadcastServerRestart("test", 0)
// The client should receive the broadcast via writePump.
readCtx, readCancel := context.WithTimeout(ctx, 2*time.Second)
defer readCancel()
_, broadcastRaw, err := conn.Read(readCtx)
if err != nil {
t.Fatalf("read broadcast: %v", err)
}
var bcast map[string]any
if err := json.Unmarshal(broadcastRaw, &bcast); err != nil {
t.Fatalf("unmarshal broadcast: %v", err)
}
// May receive member_join or presence first; drain until server_restart found.
found := bcast["type"] == "server_restart"
if !found {
// Drain a few more messages.
for i := 0; i < 5 && !found; i++ {
rCtx, rCancel := context.WithTimeout(ctx, 500*time.Millisecond)
_, raw2, err2 := conn.Read(rCtx)
rCancel()
if err2 != nil {
break
}
var m map[string]any
if json.Unmarshal(raw2, &m) == nil && m["type"] == "server_restart" {
found = true
}
}
}
if !found {
t.Error("did not receive server_restart broadcast via writePump")
}
}
// TestIntegration_MessageRoundTrip verifies that two clients can exchange messages
// through the real WebSocket upgrade path: Client A sends chat_send, Client B
// receives chat_message via the hub broadcast.
func TestIntegration_MessageRoundTrip(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
st := store.NewSQLiteStore(database)
svc := service.New(st, limiter)
hub := ws.NewHub(database, limiter, svc)
go hub.Run()
defer hub.Stop()
// Seed two users with sessions.
userIDA, err := database.CreateUser("roundtrip-a", "hash", 1)
if err != nil {
t.Fatalf("CreateUser A: %v", err)
}
tokenA, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken A: %v", err)
}
if _, err := database.CreateSession(userIDA, auth.HashToken(tokenA), "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession A: %v", err)
}
userIDB, err := database.CreateUser("roundtrip-b", "hash", 1)
if err != nil {
t.Fatalf("CreateUser B: %v", err)
}
tokenB, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken B: %v", err)
}
if _, err := database.CreateSession(userIDB, auth.HashToken(tokenB), "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession B: %v", err)
}
// Create a text channel for the chat.
chID, err := database.CreateChannel("integration-chat", "text", "", "", 0)
if err != nil {
t.Fatalf("CreateChannel: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
// --- Helper: connect and authenticate a WebSocket client ---
connectAndAuth := func(label, token string) *websocket.Conn {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
conn, dialResp, dialErr := websocket.Dial(ctx, wsURL, nil)
if dialResp != nil && dialResp.Body != nil {
dialResp.Body.Close()
}
if dialErr != nil {
t.Fatalf("%s dial: %v", label, dialErr)
}
authMsg, _ := json.Marshal(map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
})
if writeErr := conn.Write(ctx, websocket.MessageText, authMsg); writeErr != nil {
t.Fatalf("%s write auth: %v", label, writeErr)
}
// Drain auth_ok + ready.
for i := 0; i < 2; i++ {
if _, _, readErr := conn.Read(ctx); readErr != nil {
t.Fatalf("%s drain initial msg %d: %v", label, i, readErr)
}
}
return conn
}
connA := connectAndAuth("clientA", tokenA)
defer func() { _ = connA.Close(websocket.StatusNormalClosure, "") }()
connB := connectAndAuth("clientB", tokenB)
defer func() { _ = connB.Close(websocket.StatusNormalClosure, "") }()
// Wait for both clients to be registered in the hub.
time.Sleep(50 * time.Millisecond)
// Client B focuses on the channel so it receives channel-scoped broadcasts.
focusMsg, _ := json.Marshal(map[string]any{
"type": "channel_focus",
"payload": map[string]any{"channel_id": chID},
})
ctxB, cancelB := context.WithTimeout(context.Background(), 5*time.Second)
defer cancelB()
if err := connB.Write(ctxB, websocket.MessageText, focusMsg); err != nil {
t.Fatalf("clientB write channel_focus: %v", err)
}
time.Sleep(30 * time.Millisecond)
// Client A sends a chat message.
chatSend, _ := json.Marshal(map[string]any{
"type": "chat_send",
"id": "req-1",
"payload": map[string]any{
"channel_id": chID,
"content": "hello from A",
},
})
ctxA, cancelA := context.WithTimeout(context.Background(), 5*time.Second)
defer cancelA()
if err := connA.Write(ctxA, websocket.MessageText, chatSend); err != nil {
t.Fatalf("clientA write chat_send: %v", err)
}
// Client B should receive a chat_message broadcast.
// Drain a few messages (member_join, presence, etc.) until we find chat_message.
found := false
readCtx, readCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer readCancel()
for i := 0; i < 15 && !found; i++ {
_, raw, readErr := connB.Read(readCtx)
if readErr != nil {
t.Fatalf("clientB read: %v", readErr)
}
var env map[string]any
if json.Unmarshal(raw, &env) != nil {
continue
}
if env["type"] == "chat_message" {
payload, _ := env["payload"].(map[string]any)
if payload == nil {
t.Fatal("chat_message has nil payload")
}
if payload["content"] != "hello from A" {
t.Errorf("content = %q, want 'hello from A'", payload["content"])
}
user, _ := payload["user"].(map[string]any)
if user == nil {
t.Fatal("chat_message missing user")
}
if user["username"] != "roundtrip-a" {
t.Errorf("username = %q, want 'roundtrip-a'", user["username"])
}
found = true
}
}
if !found {
t.Error("clientB never received chat_message from clientA")
}
}
// TestIntegration_SequenceNumbers verifies that broadcast messages delivered via
// the real WebSocket path carry a monotonically increasing `seq` field.
func TestIntegration_SequenceNumbers(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
userID, err := database.CreateUser("seq-user", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
if _, err := database.CreateSession(userID, auth.HashToken(token), "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
// Authenticate.
authMsg, _ := json.Marshal(map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
})
if err := conn.Write(ctx, websocket.MessageText, authMsg); err != nil {
t.Fatalf("write auth: %v", err)
}
// Drain auth_ok and ready (these are direct writes, not broadcasts).
for i := 0; i < 2; i++ {
if _, _, err := conn.Read(ctx); err != nil {
t.Fatalf("drain msg %d: %v", i, err)
}
}
// Wait for registration.
time.Sleep(50 * time.Millisecond)
// Trigger two broadcasts.
hub.BroadcastServerRestart("test-seq-1", 10)
hub.BroadcastServerRestart("test-seq-2", 20)
// Collect broadcast messages — they must carry monotonically increasing seq.
var seqs []float64
readCtx, readCancel := context.WithTimeout(ctx, 3*time.Second)
defer readCancel()
for i := 0; i < 10; i++ {
_, raw, readErr := conn.Read(readCtx)
if readErr != nil {
break
}
var env map[string]any
if json.Unmarshal(raw, &env) != nil {
continue
}
// Broadcasts go through deliverBroadcast which stamps seq.
if seq, ok := env["seq"].(float64); ok {
seqs = append(seqs, seq)
}
// Stop once we've collected at least 2 seq-bearing messages.
if len(seqs) >= 2 {
break
}
}
if len(seqs) < 2 {
t.Fatalf("expected at least 2 messages with seq field, got %d", len(seqs))
}
for i := 1; i < len(seqs); i++ {
if seqs[i] <= seqs[i-1] {
t.Errorf("seq not monotonically increasing: seq[%d]=%.0f seq[%d]=%.0f", i-1, seqs[i-1], i, seqs[i])
}
}
}
// TestServeWS_BannedUser_ReceivesError verifies that a banned user cannot connect.
func TestServeWS_BannedUser_ReceivesError(t *testing.T) {
database := openServeTestDB(t)
limiter := auth.NewRateLimiter()
hub := ws.NewHub(database, limiter, nil)
go hub.Run()
defer hub.Stop()
// Seed user, then ban them.
userID, err := database.CreateUser("ws-banned-user", "hash", 1)
if err != nil {
t.Fatalf("CreateUser: %v", err)
}
token, err := auth.GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
tokenHash := auth.HashToken(token)
if _, err := database.CreateSession(userID, tokenHash, "test", "127.0.0.1"); err != nil {
t.Fatalf("CreateSession: %v", err)
}
// Ban the user permanently.
if err := database.BanUser(userID, "test ban", nil); err != nil {
t.Fatalf("BanUser: %v", err)
}
handler := ws.ServeWS(hub, database, []string{"*"})
srv := httptest.NewServer(handler)
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, resp, err := websocket.Dial(ctx, wsURL, nil)
if resp != nil && resp.Body != nil {
defer resp.Body.Close()
}
if err != nil {
t.Fatalf("websocket.Dial: %v", err)
}
defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()
authMsg := map[string]any{
"type": "auth",
"payload": map[string]string{"token": token},
}
raw, _ := json.Marshal(authMsg)
if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatalf("write auth: %v", err)
}
_, respRaw, readErr := conn.Read(ctx)
if readErr != nil {
return // server closed connection — acceptable
}
var msg map[string]any
if err := json.Unmarshal(respRaw, &msg); err == nil {
msgType, _ := msg["type"].(string)
if msgType == "auth_ok" {
t.Error("banned user should not receive auth_ok")
}
}
}