easyai-ai-gateway/apps/api/internal/oidcsession/service_test.go

package oidcsession

import (
	"bytes"
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestServiceStoresOnlyHashedSessionAndEncryptedTokens(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access-token": {ID: "subject-1", Source: "oidc", Roles: []string{"user"}, TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	service := newTestService(t, repository, verifier, &fakePublicClient{})
	service.now = func() time.Time { return now }
	localUser := &auth.User{ID: "subject-1", GatewayUserID: "11111111-1111-4111-8111-111111111111", GatewayTenantID: "22222222-2222-4222-8222-222222222222"}
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access-token", RefreshToken: "refresh-token", IDToken: "id-token"}, localUser)
	if err != nil {
		t.Fatal(err)
	}
	if len(raw) != 43 {
		t.Fatalf("opaque session length = %d, want 43", len(raw))
	}
	record := repository.snapshot()
	if bytes.Contains(record.TokenCiphertext, []byte("access-token")) || bytes.Contains(record.TokenCiphertext, []byte("refresh-token")) {
		t.Fatal("repository received plaintext token material")
	}
	hash, _ := sessionTokenHash(raw)
	if !bytes.Equal(hash, record.SessionTokenHash) || bytes.Equal([]byte(raw), record.SessionTokenHash) {
		t.Fatal("repository did not receive only the SHA-256 session hash")
	}
	user, err := service.Resolve(context.Background(), raw)
	if err != nil || user.ID != "subject-1" {
		t.Fatalf("resolve user=%#v err=%v", user, err)
	}
}

func TestServiceConcurrentExpiredRequestsRefreshExactlyOnce(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"old-access": {ID: "subject-1", Source: "oidc", Roles: []string{"user"}, TokenExpiresAt: now.Add(5 * time.Minute)},
		"new-access": {ID: "subject-1", Source: "oidc", Roles: []string{"user"}, TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	client := &fakePublicClient{refreshResponse: auth.OIDCTokenResponse{AccessToken: "new-access", RefreshToken: "rotated-refresh", ExpiresIn: 300}, delay: 40 * time.Millisecond}
	service := newTestService(t, repository, verifier, client)
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "old-access", RefreshToken: "old-refresh"}, &auth.User{
		ID: "subject-1", GatewayUserID: "11111111-1111-4111-8111-111111111111", GatewayTenantID: "22222222-2222-4222-8222-222222222222",
	})
	if err != nil {
		t.Fatal(err)
	}
	repository.expireAccessToken(now.Add(-time.Second))

	var wait sync.WaitGroup
	errorsFound := make(chan error, 20)
	for range 20 {
		wait.Add(1)
		go func() {
			defer wait.Done()
			user, resolveErr := service.Resolve(context.Background(), raw)
			if resolveErr != nil {
				errorsFound <- resolveErr
				return
			}
			if user.ID != "subject-1" {
				errorsFound <- errors.New("wrong resolved subject")
			}
		}()
	}
	wait.Wait()
	close(errorsFound)
	for err := range errorsFound {
		t.Errorf("concurrent resolve: %v", err)
	}
	if got := client.refreshCalls.Load(); got != 1 {
		t.Fatalf("refresh calls = %d, want exactly 1", got)
	}
	if repository.snapshot().RefreshVersion != 2 {
		t.Fatalf("refresh version = %d, want 2", repository.snapshot().RefreshVersion)
	}
}

func TestServiceDoesNotRefreshExpiredIdleSession(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)}}}
	client := &fakePublicClient{refreshResponse: auth.OIDCTokenResponse{AccessToken: "new", RefreshToken: "rotated"}}
	service := newTestService(t, repository, verifier, client)
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "refresh"}, &auth.User{
		ID: "subject-1", GatewayUserID: "11111111-1111-4111-8111-111111111111", GatewayTenantID: "22222222-2222-4222-8222-222222222222",
	})
	if err != nil {
		t.Fatal(err)
	}
	repository.expireIdle(now.Add(-time.Second))
	if _, err := service.Resolve(context.Background(), raw); !errors.Is(err, ErrSessionExpired) {
		t.Fatalf("Resolve() error = %v, want ErrSessionExpired", err)
	}
	if client.refreshCalls.Load() != 0 {
		t.Fatal("expired idle session attempted a refresh")
	}
}

func TestServiceKeepsExistingRefreshTokenWhenIssuerDoesNotRotate(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"old-access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
		"new-access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	client := &fakePublicClient{refreshResponse: auth.OIDCTokenResponse{AccessToken: "new-access"}}
	service := newTestService(t, repository, verifier, client)
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "old-access", RefreshToken: "existing-refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	repository.expireAccessToken(now.Add(-time.Second))
	if _, err := service.Resolve(context.Background(), raw); err != nil {
		t.Fatalf("Resolve() error = %v", err)
	}
	bundle, err := service.Delete(context.Background(), raw)
	if err != nil {
		t.Fatal(err)
	}
	if bundle.RefreshToken != "existing-refresh" || bundle.AccessToken != "new-access" {
		t.Fatal("refreshed bundle did not retain the existing refresh token")
	}
}

func TestServiceInvalidGrantDeletesSession(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	client := &fakePublicClient{refreshError: auth.ErrOIDCInvalidGrant}
	service := newTestService(t, repository, verifier, client)
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "revoked-refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	repository.expireAccessToken(now.Add(-time.Second))
	if _, err := service.Resolve(context.Background(), raw); !errors.Is(err, ErrSessionExpired) {
		t.Fatalf("Resolve() error = %v, want ErrSessionExpired", err)
	}
	if !repository.isDeleted() || client.refreshCalls.Load() != 1 {
		t.Fatal("invalid_grant did not delete the session after exactly one refresh")
	}
}

func TestServiceUsesStillValidAccessTokenWhenRefreshIsTemporarilyUnavailable(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(30 * time.Second)},
	}}
	client := &fakePublicClient{refreshError: context.DeadlineExceeded}
	service := newTestService(t, repository, verifier, client)
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	user, err := service.Resolve(context.Background(), raw)
	if err != nil || user.ID != "subject-1" {
		t.Fatalf("Resolve() user=%#v error=%v", user, err)
	}
	if client.refreshCalls.Load() != 1 || repository.isDeleted() {
		t.Fatal("temporary refresh failure did not fall back to the valid access token")
	}
}

func TestServiceReturnsUnavailableWhenExpiredTokenCannotRefresh(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	service := newTestService(t, repository, verifier, &fakePublicClient{refreshError: context.DeadlineExceeded})
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	repository.expireAccessToken(now.Add(-time.Second))
	if _, err := service.Resolve(context.Background(), raw); !errors.Is(err, ErrSessionRefreshUnavailable) {
		t.Fatalf("Resolve() error = %v, want ErrSessionRefreshUnavailable", err)
	}
}

func TestServiceRejectsWrongEncryptionKeyAndDisabledUser(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	service := newTestService(t, repository, verifier, &fakePublicClient{})
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	wrongCipher, err := NewCipher(bytes.Repeat([]byte{8}, 32))
	if err != nil {
		t.Fatal(err)
	}
	wrongKeyService, err := NewService(repository, wrongCipher, verifier, &fakePublicClient{}, Config{
		IdleTTL: 30 * time.Minute, AbsoluteTTL: 8 * time.Hour, RefreshBefore: time.Minute,
	})
	if err != nil {
		t.Fatal(err)
	}
	wrongKeyService.now = func() time.Time { return now }
	if _, err := wrongKeyService.Resolve(context.Background(), raw); !errors.Is(err, ErrSessionStoreUnavailable) {
		t.Fatalf("wrong-key Resolve() error = %v, want ErrSessionStoreUnavailable", err)
	}

	repository.disableUser()
	if _, err := service.Resolve(context.Background(), raw); !errors.Is(err, ErrGatewayUserDisabled) {
		t.Fatalf("disabled-user Resolve() error = %v, want ErrGatewayUserDisabled", err)
	}
}

func TestServiceDeletesSessionEvenWhenCiphertextCannotBeDecryptedDuringLogout(t *testing.T) {
	now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC)
	repository := newFakeRepository("subject-1")
	verifier := fakeVerifier{users: map[string]*auth.User{
		"access": {ID: "subject-1", Source: "oidc", TokenExpiresAt: now.Add(5 * time.Minute)},
	}}
	service := newTestService(t, repository, verifier, &fakePublicClient{})
	service.now = func() time.Time { return now }
	raw, err := service.Create(context.Background(), TokenBundle{AccessToken: "access", RefreshToken: "refresh"}, testLocalUser())
	if err != nil {
		t.Fatal(err)
	}
	repository.corruptCiphertext()
	if _, err := service.Delete(context.Background(), raw); err != nil {
		t.Fatalf("Delete() error = %v", err)
	}
	if !repository.isDeleted() {
		t.Fatal("logout left a session with unusable ciphertext in the store")
	}
}

func testLocalUser() *auth.User {
	return &auth.User{
		ID: "subject-1", GatewayUserID: "11111111-1111-4111-8111-111111111111",
		GatewayTenantID: "22222222-2222-4222-8222-222222222222",
	}
}

func newTestService(t *testing.T, repository Repository, verifier TokenVerifier, client PublicClient) *Service {
	t.Helper()
	cipher, err := NewCipher(bytes.Repeat([]byte{7}, 32))
	if err != nil {
		t.Fatal(err)
	}
	service, err := NewService(repository, cipher, verifier, client, Config{
		IdleTTL: 30 * time.Minute, AbsoluteTTL: 8 * time.Hour, RefreshBefore: time.Minute,
		RefreshLease: 5 * time.Second, RefreshWait: 2 * time.Second,
	})
	if err != nil {
		t.Fatal(err)
	}
	return service
}

type fakeVerifier struct{ users map[string]*auth.User }

func (f fakeVerifier) Verify(_ context.Context, token string) (*auth.User, error) {
	user := f.users[token]
	if user == nil {
		return nil, auth.ErrUnauthorized
	}
	copy := *user
	return &copy, nil
}

type fakePublicClient struct {
	refreshResponse auth.OIDCTokenResponse
	refreshError    error
	delay           time.Duration
	refreshCalls    atomic.Int64
}

func (f *fakePublicClient) Refresh(_ context.Context, _ string) (auth.OIDCTokenResponse, error) {
	f.refreshCalls.Add(1)
	if f.delay > 0 {
		time.Sleep(f.delay)
	}
	return f.refreshResponse, f.refreshError
}
func (f *fakePublicClient) RevokeRefreshToken(context.Context, string) error { return nil }
func (f *fakePublicClient) EndSessionURL(context.Context, string) (string, error) {
	return "https://gateway.example.com/", nil
}

type fakeRepository struct {
	mu       sync.Mutex
	record   store.OIDCSession
	deleted  bool
	external string
}

func newFakeRepository(external string) *fakeRepository { return &fakeRepository{external: external} }

func (f *fakeRepository) CreateOIDCSession(_ context.Context, input store.CreateOIDCSessionInput) (store.OIDCSession, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.record = store.OIDCSession{
		ID: "33333333-3333-4333-8333-333333333333", SessionTokenHash: append([]byte(nil), input.SessionTokenHash...),
		GatewayUserID: input.GatewayUserID, GatewayTenantID: input.GatewayTenantID, ExternalUserID: f.external,
		UserStatus: "active", TokenCiphertext: append([]byte(nil), input.TokenCiphertext...), AccessTokenExpiresAt: input.AccessTokenExpiresAt,
		LastSeenAt: input.LastSeenAt, IdleExpiresAt: input.IdleExpiresAt, AbsoluteExpiresAt: input.AbsoluteExpiresAt, RefreshVersion: 1,
	}
	return f.record, nil
}
func (f *fakeRepository) FindOIDCSessionByHash(_ context.Context, hash []byte) (store.OIDCSession, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.deleted || !bytes.Equal(hash, f.record.SessionTokenHash) {
		return store.OIDCSession{}, store.ErrOIDCSessionNotFound
	}
	item := f.record
	item.SessionTokenHash = append([]byte(nil), f.record.SessionTokenHash...)
	item.TokenCiphertext = append([]byte(nil), f.record.TokenCiphertext...)
	return item, nil
}
func (f *fakeRepository) TouchOIDCSession(_ context.Context, _ string, lastSeen, idle time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.deleted || !f.record.IdleExpiresAt.After(lastSeen) || !f.record.AbsoluteExpiresAt.After(lastSeen) {
		return store.ErrOIDCSessionNotFound
	}
	f.record.LastSeenAt, f.record.IdleExpiresAt = lastSeen, idle
	return nil
}
func (f *fakeRepository) AcquireOIDCSessionRefreshLock(_ context.Context, _ string, version int64, lockID string, until, now time.Time) (bool, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.deleted || f.record.RefreshVersion != version || f.record.RefreshLockID != "" && f.record.RefreshLockUntil.After(now) {
		return false, nil
	}
	f.record.RefreshLockID, f.record.RefreshLockUntil = lockID, until
	return true, nil
}
func (f *fakeRepository) CompleteOIDCSessionRefresh(_ context.Context, _ string, version int64, lockID string, ciphertext []byte, expires time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.deleted || f.record.RefreshVersion != version || f.record.RefreshLockID != lockID {
		return store.ErrOIDCSessionNotFound
	}
	f.record.TokenCiphertext = append([]byte(nil), ciphertext...)
	f.record.AccessTokenExpiresAt = expires
	f.record.RefreshVersion++
	f.record.RefreshLockID = ""
	f.record.RefreshLockUntil = time.Time{}
	return nil
}
func (f *fakeRepository) DeleteOIDCSessionByHash(context.Context, []byte) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.deleted = true
	return nil
}
func (f *fakeRepository) DeleteOIDCSessionByID(context.Context, string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.deleted = true
	return nil
}
func (f *fakeRepository) CleanupExpiredOIDCSessions(context.Context, time.Time) (int64, error) {
	return 0, nil
}
func (f *fakeRepository) snapshot() store.OIDCSession {
	f.mu.Lock()
	defer f.mu.Unlock()
	item := f.record
	item.TokenCiphertext = append([]byte(nil), item.TokenCiphertext...)
	return item
}
func (f *fakeRepository) expireAccessToken(expiry time.Time) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.record.AccessTokenExpiresAt = expiry
}
func (f *fakeRepository) expireIdle(expiry time.Time) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.record.IdleExpiresAt = expiry
}
func (f *fakeRepository) disableUser() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.record.UserStatus = "disabled"
}
func (f *fakeRepository) corruptCiphertext() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.record.TokenCiphertext = []byte("corrupt")
}
func (f *fakeRepository) isDeleted() bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.deleted
}