package services import ( "net/http" "net/http/httptest" "strings" "testing" "time" "agent-desk/internal/pkg/enums" "github.com/gorilla/websocket" ) // The listener exists only inside the test. No app server or external service is // involved; the DB fixture is temporary and every pump is joined before cleanup. func newSecuritySocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn) { t.Helper() accepted := make(chan *websocket.Conn, 1) failed := make(chan error, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upgrader := websocket.Upgrader{} conn, err := upgrader.Upgrade(w, r, nil) if err != nil { failed <- err return } accepted <- conn })) t.Cleanup(server.Close) client, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil) if err != nil { t.Fatal(err) } var serverConn *websocket.Conn select { case serverConn = <-accepted: case err := <-failed: t.Fatal(err) case <-time.After(time.Second): t.Fatal("test websocket upgrade timed out") } t.Cleanup(func() { _ = client.Close(); _ = serverConn.Close() }) return serverConn, client } func startSecurityPumps(t *testing.T, ws *wsService, session *ClientSession, reader bool) { t.Helper() written := make(chan struct{}) go func() { defer close(written); ws.writePump(session) }() var read chan struct{} if reader { read = make(chan struct{}) go func() { defer close(read); ws.readPump(session) }() } t.Cleanup(func() { ws.closeSession(session) select { case <-written: case <-time.After(time.Second): t.Error("write pump did not stop") } if read != nil { select { case <-read: case <-time.After(time.Second): t.Error("read pump did not stop") } } }) } func requireSecuritySocketClosed(t *testing.T, client *websocket.Conn) { t.Helper() _ = client.SetReadDeadline(time.Now().Add(time.Second)) _, _, err := client.ReadMessage() if err == nil { t.Fatal("revoked connection delivered a buffered message") } if networkErr, ok := err.(interface{ Timeout() bool }); ok && networkErr.Timeout() { t.Fatal("revoked connection stayed open until test deadline") } } func TestWsWritePumpRejectsBufferedEventAfterRevocation(t *testing.T) { f := newWsAuthFixture(t, true) ws := newWsService() session := newWsAuthClient(f, "buffered", realtimeRoleAdmin) serverConn, client := newSecuritySocketPair(t) session.Conn = serverConn ws.manager.Register(session, ws.defaultTopics(session)) session.enqueueEvent(ws.newEvent(realtimeTopicAdminAll, RealtimeConversationChangedEvent{Type: enums.IMRealtimeEventConversationUpdated})) if err := f.db.Model(&f.session).Update("revoked_at", time.Now()).Error; err != nil { t.Fatal(err) } startSecurityPumps(t, ws, session, false) requireSecuritySocketClosed(t, client) if !session.Closed.Load() { t.Fatal("writer did not revoke session") } } func TestWsReadPumpRejectsIncomingFrameAfterRevocation(t *testing.T) { f := newWsAuthFixture(t, true) ws := newWsService() session := newWsAuthClient(f, "incoming", realtimeRoleAdmin) serverConn, client := newSecuritySocketPair(t) session.Conn = serverConn ws.manager.Register(session, ws.defaultTopics(session)) startSecurityPumps(t, ws, session, true) if err := f.db.Model(&f.session).Update("revoked_at", time.Now()).Error; err != nil { t.Fatal(err) } if err := client.WriteMessage(websocket.TextMessage, []byte(`{"type":"ping"}`)); err != nil { t.Fatal(err) } requireSecuritySocketClosed(t, client) if !session.Closed.Load() { t.Fatal("reader did not revoke session") } } func TestWsIdlePumpClosesRevokedConnection(t *testing.T) { f := newWsAuthFixture(t, true) ws := newWsService() // Same production timer path, shortened only on this isolated service instance. ws.authCheckPeriod = 10 * time.Millisecond session := newWsAuthClient(f, "idle", realtimeRoleAdmin) serverConn, client := newSecuritySocketPair(t) session.Conn = serverConn ws.manager.Register(session, ws.defaultTopics(session)) startSecurityPumps(t, ws, session, false) if err := f.db.Model(&f.session).Update("revoked_at", time.Now()).Error; err != nil { t.Fatal(err) } requireSecuritySocketClosed(t, client) if !session.Closed.Load() { t.Fatal("idle authorization timer did not revoke session") } }