agent-desk-ng/internal/services/ws_pump_security_test.go

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