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

package oidcsession

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"strings"
	"time"

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

var (
	ErrSessionInvalid            = errors.New("OIDC session is invalid")
	ErrSessionExpired            = errors.New("OIDC session has expired")
	ErrSessionRefreshUnavailable = errors.New("OIDC session refresh is unavailable")
	ErrSessionStoreUnavailable   = errors.New("OIDC session store is unavailable")
	ErrGatewayUserDisabled       = errors.New("gateway user is disabled")
)

type Repository interface {
	CreateOIDCSession(context.Context, store.CreateOIDCSessionInput) (store.OIDCSession, error)
	FindOIDCSessionByHash(context.Context, []byte) (store.OIDCSession, error)
	TouchOIDCSession(context.Context, string, time.Time, time.Time) error
	AcquireOIDCSessionRefreshLock(context.Context, string, int64, string, time.Time, time.Time) (bool, error)
	CompleteOIDCSessionRefresh(context.Context, string, int64, string, []byte, time.Time) error
	DeleteOIDCSessionByHash(context.Context, []byte) error
	DeleteOIDCSessionByID(context.Context, string) error
	CleanupExpiredOIDCSessions(context.Context, time.Time) (int64, error)
}

type TokenVerifier interface {
	Verify(context.Context, string) (*auth.User, error)
}

type PublicClient interface {
	Refresh(context.Context, string) (auth.OIDCTokenResponse, error)
	RevokeRefreshToken(context.Context, string) error
	EndSessionURL(context.Context, string) (string, error)
}

type Config struct {
	IdleTTL       time.Duration
	AbsoluteTTL   time.Duration
	RefreshBefore time.Duration
	RefreshLease  time.Duration
	RefreshWait   time.Duration
}

type Service struct {
	repository Repository
	cipher     *Cipher
	verifier   TokenVerifier
	client     PublicClient
	config     Config
	now        func() time.Time
}

func NewService(repository Repository, cipher *Cipher, verifier TokenVerifier, client PublicClient, config Config) (*Service, error) {
	if repository == nil || cipher == nil || verifier == nil || client == nil {
		return nil, errors.New("OIDC session dependencies are required")
	}
	if config.IdleTTL <= 0 || config.AbsoluteTTL <= config.IdleTTL || config.RefreshBefore <= 0 {
		return nil, errors.New("OIDC session TTL configuration is invalid")
	}
	if config.RefreshLease <= 0 {
		config.RefreshLease = 5 * time.Second
	}
	if config.RefreshWait <= 0 {
		config.RefreshWait = 2 * time.Second
	}
	return &Service{repository: repository, cipher: cipher, verifier: verifier, client: client, config: config, now: time.Now}, nil
}

func (s *Service) Create(ctx context.Context, bundle TokenBundle, localUser *auth.User) (string, error) {
	if localUser == nil || localUser.GatewayUserID == "" || localUser.GatewayTenantID == "" || bundle.RefreshToken == "" {
		return "", ErrSessionInvalid
	}
	verified, err := s.verifier.Verify(ctx, bundle.AccessToken)
	if err != nil || verified == nil || verified.Source != "oidc" || verified.ID == "" || verified.ID != localUser.ID {
		return "", ErrSessionInvalid
	}
	now := s.now()
	if !verified.TokenExpiresAt.After(now) {
		return "", ErrSessionExpired
	}
	raw, hash, err := newSessionToken()
	if err != nil {
		return "", ErrSessionStoreUnavailable
	}
	aadSessionID := hex.EncodeToString(hash)
	ciphertext, err := s.cipher.EncryptBundle(bundle, aadSessionID, localUser.GatewayUserID)
	if err != nil {
		return "", ErrSessionStoreUnavailable
	}
	_, err = s.repository.CreateOIDCSession(ctx, store.CreateOIDCSessionInput{
		SessionTokenHash: hash, GatewayUserID: localUser.GatewayUserID, GatewayTenantID: localUser.GatewayTenantID,
		TokenCiphertext: ciphertext, AccessTokenExpiresAt: verified.TokenExpiresAt,
		LastSeenAt: now, IdleExpiresAt: now.Add(s.config.IdleTTL), AbsoluteExpiresAt: now.Add(s.config.AbsoluteTTL),
	})
	if err != nil {
		return "", ErrSessionStoreUnavailable
	}
	return raw, nil
}

func (s *Service) Resolve(ctx context.Context, raw string) (*auth.User, error) {
	hash, err := sessionTokenHash(raw)
	if err != nil {
		return nil, ErrSessionInvalid
	}
	record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
	if errors.Is(err, store.ErrOIDCSessionNotFound) {
		return nil, ErrSessionInvalid
	}
	if err != nil {
		return nil, ErrSessionStoreUnavailable
	}
	return s.resolveRecord(ctx, hash, record)
}

func (s *Service) resolveRecord(ctx context.Context, hash []byte, record store.OIDCSession) (*auth.User, error) {
	now := s.now()
	if record.UserDeleted || record.UserStatus != "active" {
		return nil, ErrGatewayUserDisabled
	}
	if !record.IdleExpiresAt.After(now) || !record.AbsoluteExpiresAt.After(now) {
		_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
		return nil, ErrSessionExpired
	}
	bundle, err := s.cipher.DecryptBundle(record.TokenCiphertext, hex.EncodeToString(hash), record.GatewayUserID)
	if err != nil {
		return nil, ErrSessionStoreUnavailable
	}
	if !record.AccessTokenExpiresAt.After(now.Add(s.config.RefreshBefore)) {
		return s.refresh(ctx, hash, record, bundle)
	}
	user, err := s.verifySessionUser(ctx, record, bundle.AccessToken)
	if err != nil {
		return nil, err
	}
	if err := s.touch(ctx, record, now); err != nil {
		return nil, err
	}
	return user, nil
}

func (s *Service) refresh(ctx context.Context, hash []byte, record store.OIDCSession, bundle TokenBundle) (*auth.User, error) {
	now := s.now()
	lockID, err := newUUID()
	if err != nil {
		return nil, ErrSessionStoreUnavailable
	}
	acquired, err := s.repository.AcquireOIDCSessionRefreshLock(ctx, record.ID, record.RefreshVersion, lockID, now.Add(s.config.RefreshLease), now)
	if err != nil {
		return nil, ErrSessionStoreUnavailable
	}
	if !acquired {
		return s.waitForRefresh(ctx, hash, record, bundle)
	}

	refreshed, err := s.client.Refresh(ctx, bundle.RefreshToken)
	if errors.Is(err, auth.ErrOIDCInvalidGrant) {
		_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
		return nil, ErrSessionExpired
	}
	if err != nil {
		// Keep the lease until it expires: an indeterminate refresh response must not
		// cause the same rotating refresh token to be replayed immediately.
		if record.AccessTokenExpiresAt.After(now) {
			user, verifyErr := s.verifySessionUser(ctx, record, bundle.AccessToken)
			if verifyErr == nil {
				if touchErr := s.touch(ctx, record, now); touchErr != nil {
					return nil, touchErr
				}
				return user, nil
			}
		}
		return nil, ErrSessionRefreshUnavailable
	}
	user, err := s.verifySessionUser(ctx, record, refreshed.AccessToken)
	if err != nil {
		_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
		return nil, ErrSessionExpired
	}
	if refreshed.RefreshToken == "" {
		// OAuth token responses may omit refresh_token when the issuer keeps the
		// existing token valid. Persist the old value unless a rotated one exists.
		refreshed.RefreshToken = bundle.RefreshToken
	}
	if refreshed.IDToken == "" {
		refreshed.IDToken = bundle.IDToken
	}
	newBundle := TokenBundle{AccessToken: refreshed.AccessToken, RefreshToken: refreshed.RefreshToken, IDToken: refreshed.IDToken}
	ciphertext, err := s.cipher.EncryptBundle(newBundle, hex.EncodeToString(hash), record.GatewayUserID)
	if err != nil {
		// The remote refresh succeeded, so the previous rotating refresh token may
		// already be invalid. Destroy the stale local session instead of replaying it.
		_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
		return nil, ErrSessionStoreUnavailable
	}
	if err := s.repository.CompleteOIDCSessionRefresh(ctx, record.ID, record.RefreshVersion, lockID, ciphertext, user.TokenExpiresAt); err != nil {
		return nil, ErrSessionStoreUnavailable
	}
	record.AccessTokenExpiresAt = user.TokenExpiresAt
	if err := s.touch(ctx, record, now); err != nil {
		return nil, err
	}
	return user, nil
}

func (s *Service) waitForRefresh(ctx context.Context, hash []byte, original store.OIDCSession, oldBundle TokenBundle) (*auth.User, error) {
	deadline := time.Now().Add(s.config.RefreshWait)
	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return nil, ErrSessionRefreshUnavailable
		case <-time.After(25 * time.Millisecond):
		}
		record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
		if errors.Is(err, store.ErrOIDCSessionNotFound) {
			return nil, ErrSessionExpired
		}
		if err != nil {
			return nil, ErrSessionStoreUnavailable
		}
		if record.RefreshVersion > original.RefreshVersion {
			return s.resolveRecord(ctx, hash, record)
		}
	}
	now := s.now()
	if original.AccessTokenExpiresAt.After(now) {
		user, err := s.verifySessionUser(ctx, original, oldBundle.AccessToken)
		if err == nil {
			if touchErr := s.touch(ctx, original, now); touchErr != nil {
				return nil, touchErr
			}
			return user, nil
		}
	}
	return nil, ErrSessionRefreshUnavailable
}

func (s *Service) verifySessionUser(ctx context.Context, record store.OIDCSession, accessToken string) (*auth.User, error) {
	user, err := s.verifier.Verify(ctx, accessToken)
	if err != nil || user == nil || user.Source != "oidc" || user.ID != record.ExternalUserID {
		return nil, ErrSessionInvalid
	}
	return user, nil
}

func (s *Service) touch(ctx context.Context, record store.OIDCSession, now time.Time) error {
	if err := s.repository.TouchOIDCSession(ctx, record.ID, now, now.Add(s.config.IdleTTL)); err != nil {
		if errors.Is(err, store.ErrOIDCSessionNotFound) {
			return ErrSessionExpired
		}
		return ErrSessionStoreUnavailable
	}
	return nil
}

func (s *Service) Delete(ctx context.Context, raw string) (TokenBundle, error) {
	hash, err := sessionTokenHash(raw)
	if err != nil {
		return TokenBundle{}, nil
	}
	record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
	if errors.Is(err, store.ErrOIDCSessionNotFound) {
		return TokenBundle{}, nil
	}
	if err != nil {
		return TokenBundle{}, ErrSessionStoreUnavailable
	}
	bundle, decryptErr := s.cipher.DecryptBundle(record.TokenCiphertext, hex.EncodeToString(hash), record.GatewayUserID)
	if err := s.repository.DeleteOIDCSessionByHash(ctx, hash); err != nil {
		return TokenBundle{}, ErrSessionStoreUnavailable
	}
	if decryptErr != nil {
		// The session row is already gone. Treat logout as successful even though
		// the unusable refresh token could not be revoked at the issuer.
		return TokenBundle{}, nil
	}
	return bundle, nil
}

func (s *Service) Cleanup(ctx context.Context) (int64, error) {
	count, err := s.repository.CleanupExpiredOIDCSessions(ctx, s.now())
	if err != nil {
		return 0, ErrSessionStoreUnavailable
	}
	return count, nil
}

func newSessionToken() (string, []byte, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", nil, err
	}
	encoded := base64.RawURLEncoding.EncodeToString(raw)
	hash := sha256.Sum256([]byte(encoded))
	return encoded, hash[:], nil
}

func sessionTokenHash(raw string) ([]byte, error) {
	raw = strings.TrimSpace(raw)
	decoded, err := base64.RawURLEncoding.DecodeString(raw)
	if err != nil || len(decoded) != 32 {
		return nil, ErrSessionInvalid
	}
	hash := sha256.Sum256([]byte(raw))
	return hash[:], nil
}

func newUUID() (string, error) {
	value := make([]byte, 16)
	if _, err := rand.Read(value); err != nil {
		return "", err
	}
	value[6] = value[6]&0x0f | 0x40
	value[8] = value[8]&0x3f | 0x80
	return hex.EncodeToString(value[0:4]) + "-" + hex.EncodeToString(value[4:6]) + "-" + hex.EncodeToString(value[6:8]) + "-" + hex.EncodeToString(value[8:10]) + "-" + hex.EncodeToString(value[10:16]), nil
}