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

package services

import (
	"slices"
	"testing"
	"time"

	"agent-desk/internal/models"
	"agent-desk/internal/pkg/constants"
	"agent-desk/internal/pkg/dto"
	"agent-desk/internal/pkg/enums"

	"github.com/mlogclub/simple/sqls"
	"gorm.io/gorm"
)

func newWsAuthFixture(t *testing.T, hasConversationView bool) *authSessionFixture {
	t.Helper()
	f := newAuthSessionFixture(t)
	f.permission.Code = constants.PermissionConversationView.Code
	if err := f.db.Model(&f.permission).Update("code", f.permission.Code).Error; err != nil {
		t.Fatal(err)
	}
	permissions := []string{f.permission.Code}
	if !hasConversationView {
		if err := f.db.Where("role_id = ?", f.role.ID).Delete(&models.RolePermission{}).Error; err != nil {
			t.Fatal(err)
		}
		permissions = nil
	}
	f.principal = buildAuthPrincipal(&f.user, &f.session, []string{f.role.Code}, permissions)
	previousDB := sqls.DB()
	sqls.SetDB(f.db)
	t.Cleanup(func() { sqls.SetDB(previousDB) })
	return f
}

func newWsAuthClient(f *authSessionFixture, id, role string) *ClientSession {
	return &ClientSession{
		ID:        id,
		Role:      role,
		Principal: f.principal,
		Topics:    make(map[string]struct{}),
		Send:      make(chan []byte, 16),
	}
}

func TestWsStaffWithoutViewPermissionHasNoConversationDefaults(t *testing.T) {
	f := newWsAuthFixture(t, false)
	ws := newWsService()
	session := newWsAuthClient(f, "no-view", realtimeRoleAdmin)
	if !ws.revalidateSession(session) {
		t.Fatal("expected a valid authenticated staff session")
	}
	if topics := ws.defaultTopics(session); len(topics) != 0 {
		t.Fatalf("staff without view permission received default topics: %v", topics)
	}
	if topics := ws.filterAllowedTopics(session, []string{realtimeTopicAdminAll, ws.adminTopic(f.user.ID), "conversation:100"}); len(topics) != 0 {
		t.Fatalf("staff without view permission could subscribe: %v", topics)
	}
}

func TestWsNotificationSessionOnlyReceivesOwnNotificationTopic(t *testing.T) {
	for _, hasView := range []bool{false, true} {
		name := "without view permission"
		if hasView {
			name = "with view permission"
		}
		t.Run(name, func(t *testing.T) {
			f := newWsAuthFixture(t, hasView)
			ws := newWsService()
			session := newWsAuthClient(f, "notification", realtimeRoleNotification)
			ownTopic := ws.notificationTopic(f.user.ID)
			if got := ws.defaultTopics(session); !slices.Equal(got, []string{ownTopic}) {
				t.Fatalf("unexpected notification defaults: %v", got)
			}
			requested := []string{ownTopic, ws.notificationTopic(f.user.ID + 1), realtimeTopicAdminAll, "conversation:100"}
			if got := ws.filterAllowedTopics(session, requested); !slices.Equal(got, []string{ownTopic}) {
				t.Fatalf("notification connection escaped its own topic: %v", got)
			}
			// Simulate stale or incorrectly registered membership. Delivery must
			// still reject conversation events, including when view is granted.
			ws.manager.Register(session, requested)
			ws.PublishToTopic("conversation:100", ws.newEvent("conversation:100", RealtimeConversationChangedEvent{Type: enums.IMRealtimeEventConversationUpdated}))
			ws.PublishToTopic(realtimeTopicAdminAll, ws.newEvent(realtimeTopicAdminAll, RealtimeConversationChangedEvent{Type: enums.IMRealtimeEventConversationUpdated}))
			if len(session.Send) != 0 {
				t.Fatal("notification connection received a conversation event")
			}
			ws.PublishToTopic(ownTopic, ws.newEvent(ownTopic, RealtimeNotificationCreatedEvent{}))
			if len(session.Send) != 1 || session.Closed.Load() {
				t.Fatal("valid notification connection did not receive its own notification")
			}
		})
	}
}

func TestWsStaleAdminAllMembershipDoesNotGrantDelivery(t *testing.T) {
	f := newWsAuthFixture(t, false)
	ws := newWsService()
	session := newWsAuthClient(f, "stale-member", realtimeRoleAdmin)
	ws.manager.Register(session, []string{realtimeTopicAdminAll})
	ws.PublishToTopic(realtimeTopicAdminAll, ws.newEvent(realtimeTopicAdminAll, RealtimeConversationChangedEvent{Type: enums.IMRealtimeEventConversationUpdated}))
	if len(session.Send) != 0 {
		t.Fatal("stale admin:all membership bypassed conversation-view permission")
	}
}

func TestWsStaffRevocationClosesBeforeDeliveryOrIncomingMessage(t *testing.T) {
	cases := []struct {
		name   string
		change func(*authSessionFixture) error
	}{
		{"revoked login", func(f *authSessionFixture) error {
			return f.db.Model(&f.session).Update("revoked_at", time.Now()).Error
		}},
		{"expired login", func(f *authSessionFixture) error {
			return f.db.Model(&f.session).Update("expired_at", time.Now().Add(-time.Second)).Error
		}},
		{"disabled user", func(f *authSessionFixture) error {
			return f.db.Model(&f.user).Update("status", enums.StatusDisabled).Error
		}},
		{"downgraded user type", func(f *authSessionFixture) error {
			return f.db.Model(&f.user).Update("user_type", enums.UserTypeUser).Error
		}},
		{"deleted user", func(f *authSessionFixture) error { return f.db.Model(&f.user).Update("deleted_at", time.Now()).Error }},
		{"removed role", func(f *authSessionFixture) error {
			return f.db.Where("user_id = ?", f.user.ID).Delete(&models.UserRole{}).Error
		}},
		{"disabled role", func(f *authSessionFixture) error {
			return f.db.Model(&f.role).Update("status", enums.StatusDisabled).Error
		}},
		{"removed permission", func(f *authSessionFixture) error {
			return f.db.Where("role_id = ?", f.role.ID).Delete(&models.RolePermission{}).Error
		}},
		{"disabled permission", func(f *authSessionFixture) error {
			return f.db.Model(&f.permission).Update("status", enums.StatusDisabled).Error
		}},
	}
	for _, tc := range cases {
		for _, trigger := range []string{"fanout", "incoming"} {
			t.Run(tc.name+"/"+trigger, func(t *testing.T) {
				f := newWsAuthFixture(t, true)
				ws := newWsService()
				session := newWsAuthClient(f, "staff", realtimeRoleAdmin)
				ws.manager.Register(session, ws.defaultTopics(session))
				if !ws.revalidateSession(session) {
					t.Fatal("expected valid staff session before access change")
				}
				if err := tc.change(f); err != nil {
					t.Fatal(err)
				}
				event := ws.newEvent(realtimeTopicAdminAll, RealtimeConversationChangedEvent{Type: enums.IMRealtimeEventConversationUpdated})
				if trigger == "fanout" {
					ws.PublishToTopic(realtimeTopicAdminAll, event)
				} else {
					ws.handleClientMessage(session, []byte(`{"type":"ping"}`))
				}
				if !session.Closed.Load() || ws.manager.HasTopic(realtimeTopicAdminAll) {
					t.Fatal("stale staff connection remained registered")
				}
				if len(session.Send) != 0 {
					t.Fatal("stale staff connection received an event or pong")
				}
				ws.handleClientMessage(session, []byte(`{"type":"subscribe","topics":["admin:all","conversation:100"]}`))
				ws.PublishToTopic(realtimeTopicAdminAll, event)
				if len(session.Send) != 0 || len(ws.manager.TopicList(session)) != 0 {
					t.Fatal("closed connection processed later messages or delivery")
				}
			})
		}
	}
}

// Existing realtime integration tests also need real login-session bindings now
// that delivery revalidates authorization instead of trusting topic membership.
func seedRealtimeStaffPrincipal(t *testing.T, db *gorm.DB) *dto.AuthPrincipal {
	t.Helper()
	if err := db.AutoMigrate(&models.Role{}, &models.Permission{}, &models.UserRole{}, &models.RolePermission{}, &models.UserPermission{}, &models.LoginSession{}); err != nil {
		t.Fatal(err)
	}
	user := models.User{ID: 1001, Username: "realtime-observer", UserType: enums.UserTypeEmployee, Status: enums.StatusOk}
	role := models.Role{Code: "realtime-observer", Status: enums.StatusOk}
	permission := models.Permission{Code: constants.PermissionConversationView.Code, Status: enums.StatusOk}
	for _, item := range []any{&user, &role, &permission} {
		if err := db.Create(item).Error; err != nil {
			t.Fatal(err)
		}
	}
	session := models.LoginSession{UserID: user.ID, Token: "ak_test_only_realtime", ExpiredAt: time.Now().Add(time.Hour)}
	for _, item := range []any{
		&session,
		&models.UserRole{UserID: user.ID, RoleID: role.ID},
		&models.RolePermission{RoleID: role.ID, PermissionID: permission.ID},
	} {
		if err := db.Create(item).Error; err != nil {
			t.Fatal(err)
		}
	}
	return buildAuthPrincipal(&user, &session, []string{role.Code}, []string{permission.Code})
}