easyai-ai-gateway/apps/api/internal/store/oidc_sessions.go

package store

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"
)

var ErrOIDCSessionNotFound = errors.New("OIDC session not found")

type OIDCSession struct {
	ID                   string
	SessionTokenHash     []byte
	GatewayUserID        string
	GatewayTenantID      string
	ExternalUserID       string
	UserStatus           string
	UserDeleted          bool
	TokenCiphertext      []byte
	AccessTokenExpiresAt time.Time
	LastSeenAt           time.Time
	IdleExpiresAt        time.Time
	AbsoluteExpiresAt    time.Time
	RefreshVersion       int64
	RefreshLockID        string
	RefreshLockUntil     time.Time
	CreatedAt            time.Time
	UpdatedAt            time.Time
}

type CreateOIDCSessionInput struct {
	SessionTokenHash     []byte
	GatewayUserID        string
	GatewayTenantID      string
	TokenCiphertext      []byte
	AccessTokenExpiresAt time.Time
	LastSeenAt           time.Time
	IdleExpiresAt        time.Time
	AbsoluteExpiresAt    time.Time
}

func (s *Store) CreateOIDCSession(ctx context.Context, input CreateOIDCSessionInput) (OIDCSession, error) {
	var id string
	err := s.pool.QueryRow(ctx, `
INSERT INTO gateway_oidc_sessions (
  session_token_hash, gateway_user_id, gateway_tenant_id, token_ciphertext,
  access_token_expires_at, last_seen_at, idle_expires_at, absolute_expires_at
)
VALUES ($1, $2::uuid, $3::uuid, $4, $5, $6, $7, $8)
RETURNING id::text`,
		input.SessionTokenHash, input.GatewayUserID, input.GatewayTenantID, input.TokenCiphertext,
		input.AccessTokenExpiresAt, input.LastSeenAt, input.IdleExpiresAt, input.AbsoluteExpiresAt,
	).Scan(&id)
	if err != nil {
		return OIDCSession{}, err
	}
	return s.FindOIDCSessionByHash(ctx, input.SessionTokenHash)
}

func (s *Store) FindOIDCSessionByHash(ctx context.Context, hash []byte) (OIDCSession, error) {
	item, err := scanOIDCSession(s.pool.QueryRow(ctx, `
SELECT `+oidcSessionColumns+`
FROM gateway_oidc_sessions s
JOIN gateway_users u ON u.id = s.gateway_user_id
WHERE s.session_token_hash = $1`, hash))
	if errors.Is(err, pgx.ErrNoRows) {
		return OIDCSession{}, ErrOIDCSessionNotFound
	}
	return item, err
}

func (s *Store) TouchOIDCSession(ctx context.Context, id string, lastSeenAt, idleExpiresAt time.Time) error {
	tag, err := s.pool.Exec(ctx, `
UPDATE gateway_oidc_sessions
SET last_seen_at = $2,
    idle_expires_at = LEAST($3, absolute_expires_at),
    updated_at = now()
WHERE id = $1::uuid
  AND idle_expires_at > $2
  AND absolute_expires_at > $2`, id, lastSeenAt, idleExpiresAt)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrOIDCSessionNotFound
	}
	return nil
}

func (s *Store) AcquireOIDCSessionRefreshLock(ctx context.Context, id string, version int64, lockID string, lockUntil, now time.Time) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
UPDATE gateway_oidc_sessions
SET refresh_lock_id = $3::uuid, refresh_lock_until = $4, updated_at = now()
WHERE id = $1::uuid
  AND refresh_version = $2
  AND (refresh_lock_until IS NULL OR refresh_lock_until <= $5)`, id, version, lockID, lockUntil, now)
	return tag.RowsAffected() == 1, err
}

func (s *Store) CompleteOIDCSessionRefresh(ctx context.Context, id string, version int64, lockID string, ciphertext []byte, accessTokenExpiresAt time.Time) error {
	tag, err := s.pool.Exec(ctx, `
UPDATE gateway_oidc_sessions
SET token_ciphertext = $4,
    access_token_expires_at = $5,
    refresh_version = refresh_version + 1,
    refresh_lock_id = NULL,
    refresh_lock_until = NULL,
    updated_at = now()
WHERE id = $1::uuid AND refresh_version = $2 AND refresh_lock_id = $3::uuid`,
		id, version, lockID, ciphertext, accessTokenExpiresAt)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrOIDCSessionNotFound
	}
	return nil
}

func (s *Store) DeleteOIDCSessionByHash(ctx context.Context, hash []byte) error {
	_, err := s.pool.Exec(ctx, `DELETE FROM gateway_oidc_sessions WHERE session_token_hash = $1`, hash)
	return err
}

func (s *Store) DeleteOIDCSessionByID(ctx context.Context, id string) error {
	_, err := s.pool.Exec(ctx, `DELETE FROM gateway_oidc_sessions WHERE id = $1::uuid`, id)
	return err
}

func (s *Store) CleanupExpiredOIDCSessions(ctx context.Context, now time.Time) (int64, error) {
	tag, err := s.pool.Exec(ctx, `
DELETE FROM gateway_oidc_sessions
WHERE idle_expires_at <= $1 OR absolute_expires_at <= $1`, now)
	return tag.RowsAffected(), err
}

const oidcSessionColumns = `
  s.id::text, s.session_token_hash, s.gateway_user_id::text, s.gateway_tenant_id::text,
  COALESCE(u.external_user_id, ''), u.status, u.deleted_at IS NOT NULL,
  s.token_ciphertext, s.access_token_expires_at, s.last_seen_at, s.idle_expires_at,
  s.absolute_expires_at, s.refresh_version, COALESCE(s.refresh_lock_id::text, ''),
  COALESCE(s.refresh_lock_until, 'epoch'::timestamptz), s.created_at, s.updated_at`

func scanOIDCSession(row pgx.Row) (OIDCSession, error) {
	var item OIDCSession
	err := row.Scan(
		&item.ID, &item.SessionTokenHash, &item.GatewayUserID, &item.GatewayTenantID,
		&item.ExternalUserID, &item.UserStatus, &item.UserDeleted, &item.TokenCiphertext,
		&item.AccessTokenExpiresAt, &item.LastSeenAt, &item.IdleExpiresAt, &item.AbsoluteExpiresAt,
		&item.RefreshVersion, &item.RefreshLockID, &item.RefreshLockUntil, &item.CreatedAt, &item.UpdatedAt,
	)
	return item, err
}