From 5edebb43b3d7f81c489e0bc645d069c55e82869b Mon Sep 17 00:00:00 2001 From: J3vb Date: Thu, 2 Apr 2026 12:12:56 +0200 Subject: [PATCH] fix: add 30s session sweep to kick revoked WS connections (BUG-109) Idle WebSocket connections only revalidated sessions every 10 sent messages, allowing revoked tokens to stay connected indefinitely. Add sweepRevokedSessions() on a 30s ticker that checks all connected clients against the DB and kicks any with deleted/expired sessions or banned users. --- Server/ws/export_test.go | 5 +++ Server/ws/hub.go | 40 +++++++++++++++++ Server/ws/hub_test.go | 93 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+) diff --git a/Server/ws/export_test.go b/Server/ws/export_test.go index 2f33055b..744feb69 100644 --- a/Server/ws/export_test.go +++ b/Server/ws/export_test.go @@ -25,6 +25,11 @@ func (h *Hub) SweepStaleVoiceStatesForTest() { h.sweepStaleVoiceStates() } +// SweepRevokedSessionsForTest exposes sweepRevokedSessions for external tests. +func (h *Hub) SweepRevokedSessionsForTest() { + h.sweepRevokedSessions() +} + // SetClientLastActivityForTest overwrites a client's lastActivity timestamp. func SetClientLastActivityForTest(c *Client, t time.Time) { c.mu.Lock() diff --git a/Server/ws/hub.go b/Server/ws/hub.go index 708d2a6a..a453a84c 100644 --- a/Server/ws/hub.go +++ b/Server/ws/hub.go @@ -148,6 +148,8 @@ func (h *Hub) Run() { func() { staleTicker := time.NewTicker(30 * time.Second) defer staleTicker.Stop() + sessionSweepTicker := time.NewTicker(30 * time.Second) + defer sessionSweepTicker.Stop() voiceSweepTicker := time.NewTicker(60 * time.Second) defer voiceSweepTicker.Stop() @@ -187,6 +189,8 @@ func (h *Hub) Run() { h.deliverBroadcast(bm) case <-staleTicker.C: h.sweepStaleClients() + case <-sessionSweepTicker.C: + h.sweepRevokedSessions() case <-voiceSweepTicker.C: h.sweepStaleVoiceStates() } @@ -491,6 +495,42 @@ func (h *Hub) sweepStaleClients() { } } +// sweepRevokedSessions iterates all connected clients and kicks any whose +// session has been deleted, expired, or whose user has been banned. This +// provides time-based session enforcement for idle WebSocket connections +// that never trigger the message-count-based check (BUG-109). +func (h *Hub) sweepRevokedSessions() { + if h.db == nil { + return + } + + h.mu.RLock() + snapshot := make([]*Client, 0, len(h.clients)) + for _, c := range h.clients { + if c.tokenHash != "" { + snapshot = append(snapshot, c) + } + } + h.mu.RUnlock() + + for _, c := range snapshot { + result, err := h.db.GetSessionWithBanStatus(c.tokenHash) + if err != nil || result == nil || auth.IsSessionExpired(result.ExpiresAt) { + slog.Info("session sweep: revoked/expired session, disconnecting", + "user_id", c.userID) + h.kickClient(c) + continue + } + tempUser := &db.User{Banned: result.Banned, BanExpires: result.BanExpires} + if auth.IsEffectivelyBanned(tempUser) { + slog.Info("session sweep: banned user, disconnecting", + "user_id", c.userID) + c.sendMsg(buildErrorMsg(ErrCodeBanned, "you are banned")) + h.kickClient(c) + } + } +} + // sweepStaleVoiceStates queries all voice_states rows and removes any that // don't match a connected client's voiceChID. This catches ghost users that // slip through the primary cleanup paths (registerNow, readPump defer, diff --git a/Server/ws/hub_test.go b/Server/ws/hub_test.go index 08fbbcd8..7eb169e4 100644 --- a/Server/ws/hub_test.go +++ b/Server/ws/hub_test.go @@ -699,6 +699,99 @@ func TestHub_SweepStaleClients_AllFresh(t *testing.T) { } } +// ─── Session sweep (BUG-109) ────────────────────────────────────────────── + +// TestHub_SweepRevokedSessions_KicksRevokedClient verifies that the periodic +// session sweep disconnects clients whose sessions have been deleted from the +// database (e.g. after logout on another device). +func TestHub_SweepRevokedSessions_KicksRevokedClient(t *testing.T) { + hub, database := newTestHub(t) + go hub.Run() + defer hub.Stop() + + // Create two users with sessions. + uid1, err := database.CreateUser("alice-revoke", "hash", 3) + if err != nil { + t.Fatalf("CreateUser: %v", err) + } + uid2, err := database.CreateUser("bob-valid", "hash", 3) + if err != nil { + t.Fatalf("CreateUser: %v", err) + } + + u1, _ := database.GetUserByID(uid1) + u2, _ := database.GetUserByID(uid2) + + token1 := "revoke-token-1" + token2 := "valid-token-2" + hash1 := auth.HashToken(token1) + hash2 := auth.HashToken(token2) + + if _, err := database.CreateSession(uid1, hash1, "test", "127.0.0.1"); err != nil { + t.Fatalf("CreateSession 1: %v", err) + } + if _, err := database.CreateSession(uid2, hash2, "test", "127.0.0.1"); err != nil { + t.Fatalf("CreateSession 2: %v", err) + } + + s1 := make(chan []byte, 4) + s2 := make(chan []byte, 4) + c1 := ws.NewTestClientWithTokenHash(hub, u1, hash1, 0, s1) + c2 := ws.NewTestClientWithTokenHash(hub, u2, hash2, 0, s2) + + hub.Register(c1) + hub.Register(c2) + time.Sleep(20 * time.Millisecond) + + // Delete alice's session (simulating logout from another device). + if err := database.DeleteSession(hash1); err != nil { + t.Fatalf("DeleteSession: %v", err) + } + + // Run the session sweep. + hub.SweepRevokedSessionsForTest() + time.Sleep(20 * time.Millisecond) + + // Alice should be kicked, Bob should remain. + if hub.GetClient(uid1) != nil { + t.Error("revoked client alice should have been kicked") + } + if hub.GetClient(uid2) == nil { + t.Error("valid client bob should still be connected") + } + if hub.ClientCount() != 1 { + t.Errorf("ClientCount = %d, want 1", hub.ClientCount()) + } +} + +// TestHub_SweepRevokedSessions_NoDBNoPanic verifies the sweep is a no-op +// when the hub has no database (nil-safe). +func TestHub_SweepRevokedSessions_NoDBNoPanic(t *testing.T) { + hub := ws.NewHubForTest() + hub.SweepRevokedSessionsForTest() // should not panic +} + +// TestHub_SweepRevokedSessions_EmptyTokenHashSkipped verifies that clients +// without a token hash (e.g. test clients) are not kicked by the sweep. +func TestHub_SweepRevokedSessions_EmptyTokenHashSkipped(t *testing.T) { + hub, database := newTestHub(t) + go hub.Run() + defer hub.Stop() + + uid := seedTestUser(t, database, "no-hash-user") + s := make(chan []byte, 4) + c := ws.NewTestClient(hub, uid, s) + hub.Register(c) + time.Sleep(20 * time.Millisecond) + + hub.SweepRevokedSessionsForTest() + time.Sleep(20 * time.Millisecond) + + if hub.ClientCount() != 1 { + t.Errorf("ClientCount = %d, want 1 (client without token hash should survive)", hub.ClientCount()) + } +} + // ─── LiveKitHealthCheck ───────────────────────────────────────────────────── func TestHub_LiveKitHealthCheck_NilReturnsError(t *testing.T) {