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
	OIDCUserBindingID    string
	ExternalUserID       string
	Issuer               string
	ApplicationID        string
	TenantID             string
	OIDCClientID         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
	OIDCUserBindingID    string
	OIDCClientID         string
	Issuer               string
	ApplicationID        string
	TenantID             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,
  oidc_user_binding_id, oidc_client_id
)
SELECT $1, u.id, $3::uuid, $4, $5, $6, $7, $8, NULLIF($9, '')::uuid, NULLIF($10, '')
FROM gateway_users u
LEFT JOIN gateway_oidc_user_bindings ub ON ub.id = NULLIF($9, '')::uuid
LEFT JOIN gateway_oidc_tenant_bindings tb ON tb.id = ub.tenant_binding_id
WHERE u.id = $2::uuid
  AND u.gateway_tenant_id = $3::uuid
  AND ($9 = '' OR (
    ub.gateway_user_id = u.id
    AND tb.gateway_tenant_id = u.gateway_tenant_id
    AND tb.issuer = $11
    AND tb.application_id = $12
    AND tb.external_tenant_id = $13
  ))
RETURNING id::text`,
		input.SessionTokenHash, input.GatewayUserID, input.GatewayTenantID, input.TokenCiphertext,
		input.AccessTokenExpiresAt, input.LastSeenAt, input.IdleExpiresAt, input.AbsoluteExpiresAt,
		input.OIDCUserBindingID, input.OIDCClientID,
		input.Issuer, input.ApplicationID, input.TenantID,
	).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
LEFT JOIN gateway_oidc_user_bindings ub ON ub.id = s.oidc_user_binding_id
LEFT JOIN gateway_oidc_tenant_bindings tb ON tb.id = ub.tenant_binding_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(s.oidc_user_binding_id::text, ''), COALESCE(ub.subject, u.external_user_id, ''),
  COALESCE(tb.issuer, ''), COALESCE(tb.application_id, ''), COALESCE(tb.external_tenant_id, ''),
  COALESCE(s.oidc_client_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.OIDCUserBindingID, &item.ExternalUserID, &item.Issuer, &item.ApplicationID,
		&item.TenantID, &item.OIDCClientID, &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
}