moteva/server/internal/modules/auth/postgres_store.go

package auth

import (
	"context"
	"errors"
	"strings"
	"time"

	"img_infinite_canvas/internal/domain/design"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
)

type PostgresStore struct {
	pool *pgxpool.Pool
}

func NewPostgresStore(ctx context.Context, dataSource string) (*PostgresStore, error) {
	pool, err := pgxpool.New(ctx, dataSource)
	if err != nil {
		return nil, err
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, err
	}
	return &PostgresStore{pool: pool}, nil
}

func (s *PostgresStore) FindUserByID(ctx context.Context, id string) (User, error) {
	return s.findUser(ctx, `SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status FROM users WHERE id = $1`, toPGUUID(id))
}

func (s *PostgresStore) FindUserByPhone(ctx context.Context, countryCode string, phone string) (User, error) {
	return s.findUser(ctx, `SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status FROM users WHERE phone_country_code = $1 AND phone = $2`, normalizeCountryCode(countryCode), normalizePhone(phone))
}

func (s *PostgresStore) FindUserByEmail(ctx context.Context, email string) (User, error) {
	return s.findUser(ctx, `SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status FROM users WHERE LOWER(email) = LOWER($1)`, normalizeEmail(email))
}

func (s *PostgresStore) FindUserByGoogleSubject(ctx context.Context, subject string) (User, error) {
	return s.findUser(ctx, `SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status FROM users WHERE google_subject = $1`, subject)
}

func (s *PostgresStore) FindUserByWechat(ctx context.Context, openID string, unionID string) (User, error) {
	if unionID != "" {
		return s.findUser(ctx, `
SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status
FROM users
WHERE wechat_union_id = $1 OR wechat_open_id = $2
ORDER BY CASE WHEN wechat_union_id = $1 THEN 0 ELSE 1 END
LIMIT 1`, unionID, openID)
	}
	return s.findUser(ctx, `SELECT id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status FROM users WHERE wechat_open_id = $1`, openID)
}

func (s *PostgresStore) UpsertUser(ctx context.Context, user User) (User, error) {
	now := time.Now()
	phone := normalizePhone(user.Phone)
	countryCode := ""
	if phone != "" {
		countryCode = normalizeCountryCode(user.CountryCode)
	}
	row := s.pool.QueryRow(ctx, `
INSERT INTO users (id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status, phone_verified_at, email_verified_at, last_login_at, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12,
        CASE WHEN $4 <> '' THEN $13::timestamptz ELSE NULL::timestamptz END,
        CASE WHEN $5 <> '' THEN $13::timestamptz ELSE NULL::timestamptz END,
        $13, $13, $13)
ON CONFLICT (id) DO UPDATE SET
    region = EXCLUDED.region,
    phone_country_code = CASE WHEN EXCLUDED.phone <> '' THEN EXCLUDED.phone_country_code ELSE users.phone_country_code END,
    phone = CASE WHEN EXCLUDED.phone <> '' THEN EXCLUDED.phone ELSE users.phone END,
    email = CASE WHEN EXCLUDED.email <> '' THEN EXCLUDED.email ELSE users.email END,
    name = CASE WHEN EXCLUDED.name <> '' THEN EXCLUDED.name ELSE users.name END,
    avatar_url = CASE WHEN EXCLUDED.avatar_url <> '' THEN EXCLUDED.avatar_url ELSE users.avatar_url END,
    password_hash = CASE WHEN EXCLUDED.password_hash <> '' THEN EXCLUDED.password_hash ELSE users.password_hash END,
    wechat_open_id = CASE WHEN EXCLUDED.wechat_open_id <> '' THEN EXCLUDED.wechat_open_id ELSE users.wechat_open_id END,
    wechat_union_id = CASE WHEN EXCLUDED.wechat_union_id <> '' THEN EXCLUDED.wechat_union_id ELSE users.wechat_union_id END,
    google_subject = CASE WHEN EXCLUDED.google_subject <> '' THEN EXCLUDED.google_subject ELSE users.google_subject END,
    status = EXCLUDED.status,
    phone_verified_at = CASE WHEN EXCLUDED.phone <> '' THEN EXCLUDED.phone_verified_at ELSE users.phone_verified_at END,
    email_verified_at = CASE WHEN EXCLUDED.email <> '' THEN EXCLUDED.email_verified_at ELSE users.email_verified_at END,
    last_login_at = EXCLUDED.last_login_at,
    updated_at = EXCLUDED.updated_at
RETURNING id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status
`, toPGUUID(user.ID), user.Region, countryCode, phone, normalizeEmail(user.Email), user.Name, user.AvatarURL, user.PasswordHash, user.WechatOpenID, user.WechatUnionID, user.GoogleSubject, activeStatus(user.Status), now)
	return scanUser(row)
}

func (s *PostgresStore) UpdateUserProfile(ctx context.Context, userID string, patch ProfilePatch) (User, error) {
	name := pgtype.Text{}
	if patch.Name != nil {
		name = pgtype.Text{String: strings.TrimSpace(*patch.Name), Valid: true}
	}
	avatarURL := pgtype.Text{}
	if patch.AvatarURL != nil {
		avatarURL = pgtype.Text{String: strings.TrimSpace(*patch.AvatarURL), Valid: true}
	}
	row := s.pool.QueryRow(ctx, `
UPDATE users
SET name = CASE WHEN $2::text IS NULL THEN name ELSE $2::text END,
    avatar_url = CASE WHEN $3::text IS NULL THEN avatar_url ELSE $3::text END,
    updated_at = NOW()
WHERE id = $1
RETURNING id, region, phone_country_code, phone, email, name, avatar_url, password_hash, wechat_open_id, wechat_union_id, google_subject, status
`, toPGUUID(userID), name, avatarURL)
	return scanUser(row)
}

func (s *PostgresStore) CreateDeviceSession(ctx context.Context, session DeviceSession, activeLimit int64) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := lockDeviceSessions(ctx, tx, session.UserID, session.DeviceType); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO auth_device_sessions (
    id, user_id, device_type, system, browser, client, user_agent,
    created_at, last_seen_at, expires_at, revoked_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
ON CONFLICT (id) DO NOTHING
`, session.ID, toPGUUID(session.UserID), session.DeviceType, session.System, session.Browser, session.Client, session.UserAgent,
		session.CreatedAt, session.LastSeenAt, session.ExpiresAt, session.RevokedAt); err != nil {
		return err
	}
	if _, err := enforceDeviceSessionLimitPostgres(ctx, tx, session.UserID, session.DeviceType, session.ID, activeLimit, session.CreatedAt); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *PostgresStore) FindDeviceSession(ctx context.Context, id string) (DeviceSession, error) {
	return scanDeviceSession(s.pool.QueryRow(ctx, `
SELECT id, user_id, device_type, system, browser, client, user_agent, created_at, last_seen_at, expires_at, revoked_at
FROM auth_device_sessions
WHERE id = $1
`, strings.TrimSpace(id)))
}

func (s *PostgresStore) ListDeviceSessions(ctx context.Context, userID string, now time.Time) ([]DeviceSession, error) {
	rows, err := s.pool.Query(ctx, `
SELECT id, user_id, device_type, system, browser, client, user_agent, created_at, last_seen_at, expires_at, revoked_at
FROM auth_device_sessions
WHERE user_id = $1 AND revoked_at IS NULL AND expires_at > $2
ORDER BY last_seen_at DESC, created_at DESC, id ASC
`, toPGUUID(userID), now)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	sessions := make([]DeviceSession, 0)
	for rows.Next() {
		session, scanErr := scanDeviceSession(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		sessions = append(sessions, session)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return sessions, nil
}

func (s *PostgresStore) TouchDeviceSession(ctx context.Context, session DeviceSession) error {
	result, err := s.pool.Exec(ctx, `
UPDATE auth_device_sessions
SET device_type = $3,
    system = $4,
    browser = $5,
    client = $6,
    user_agent = $7,
    last_seen_at = GREATEST(last_seen_at, $8)
WHERE id = $1 AND user_id = $2 AND revoked_at IS NULL AND expires_at > $8
`, session.ID, toPGUUID(session.UserID), session.DeviceType, session.System, session.Browser, session.Client, session.UserAgent, session.LastSeenAt)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return design.ErrNotFound
	}
	return nil
}

func (s *PostgresStore) EnforceDeviceSessionLimit(
	ctx context.Context,
	userID string,
	deviceType string,
	keepSessionID string,
	activeLimit int64,
	revokedAt time.Time,
) (int64, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := lockDeviceSessions(ctx, tx, userID, deviceType); err != nil {
		return 0, err
	}
	removed, err := enforceDeviceSessionLimitPostgres(ctx, tx, userID, deviceType, keepSessionID, activeLimit, revokedAt)
	if err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return removed, nil
}

func lockDeviceSessions(ctx context.Context, tx pgx.Tx, userID string, deviceType string) error {
	_, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, normalizeUUID(userID)+":"+deviceType)
	return err
}

func enforceDeviceSessionLimitPostgres(
	ctx context.Context,
	tx pgx.Tx,
	userID string,
	deviceType string,
	keepSessionID string,
	activeLimit int64,
	revokedAt time.Time,
) (int64, error) {
	if activeLimit < 1 {
		activeLimit = 1
	}
	result, err := tx.Exec(ctx, `
WITH overflow AS (
    SELECT id
    FROM auth_device_sessions
    WHERE user_id = $1 AND device_type = $2 AND revoked_at IS NULL AND expires_at > $4
    ORDER BY CASE WHEN id = $3 THEN 0 ELSE 1 END,
             created_at DESC,
             last_seen_at DESC,
             id DESC
    OFFSET $5
)
UPDATE auth_device_sessions
SET revoked_at = $4
WHERE id IN (SELECT id FROM overflow)
`, toPGUUID(userID), deviceType, strings.TrimSpace(keepSessionID), revokedAt, activeLimit)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected(), nil
}

func (s *PostgresStore) RevokeOtherDeviceSessions(ctx context.Context, userID string, currentSessionID string, revokedAt time.Time) (int64, error) {
	result, err := s.pool.Exec(ctx, `
UPDATE auth_device_sessions
SET revoked_at = $3
WHERE user_id = $1 AND id <> $2 AND revoked_at IS NULL AND expires_at > $3
`, toPGUUID(userID), strings.TrimSpace(currentSessionID), revokedAt)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected(), nil
}

func (s *PostgresStore) RevokeDeviceSession(ctx context.Context, userID string, sessionID string, revokedAt time.Time) (int64, error) {
	result, err := s.pool.Exec(ctx, `
UPDATE auth_device_sessions
SET revoked_at = $3
WHERE user_id = $1 AND id = $2 AND revoked_at IS NULL AND expires_at > $3
`, toPGUUID(userID), strings.TrimSpace(sessionID), revokedAt)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected(), nil
}

func (s *PostgresStore) CreateVerificationCode(ctx context.Context, code VerificationCode) error {
	_, err := s.pool.Exec(ctx, `
INSERT INTO auth_verification_codes (id, region, channel, target, purpose, code_hash, attempts, expires_at, created_at)
VALUES ($1, $2, $3, $4, $5, $6, 0, $7, NOW())
`, code.ID, code.Region, code.Channel, code.Target, code.Purpose, code.CodeHash, code.ExpiresAt)
	return err
}

func (s *PostgresStore) FindActiveVerificationCode(ctx context.Context, region string, channel string, target string, purpose string) (VerificationCode, error) {
	row := s.pool.QueryRow(ctx, `
SELECT id, region, channel, target, purpose, code_hash, attempts, expires_at, consumed_at
FROM auth_verification_codes
WHERE region = $1 AND channel = $2 AND target = $3 AND purpose = $4 AND consumed_at IS NULL
ORDER BY created_at DESC
LIMIT 1
`, region, channel, target, purpose)
	var code VerificationCode
	var consumedAt *time.Time
	if err := row.Scan(&code.ID, &code.Region, &code.Channel, &code.Target, &code.Purpose, &code.CodeHash, &code.Attempts, &code.ExpiresAt, &consumedAt); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return VerificationCode{}, design.ErrNotFound
		}
		return VerificationCode{}, err
	}
	code.ConsumedAt = consumedAt
	return code, nil
}

func (s *PostgresStore) MarkVerificationCodeConsumed(ctx context.Context, id string, consumedAt time.Time) error {
	_, err := s.pool.Exec(ctx, `UPDATE auth_verification_codes SET consumed_at = $2 WHERE id = $1`, id, consumedAt)
	return err
}

func (s *PostgresStore) IncrementVerificationCodeAttempts(ctx context.Context, id string) error {
	_, err := s.pool.Exec(ctx, `UPDATE auth_verification_codes SET attempts = attempts + 1 WHERE id = $1`, id)
	return err
}

func (s *PostgresStore) Close() error {
	if s == nil || s.pool == nil {
		return nil
	}
	s.pool.Close()
	return nil
}

func (s *PostgresStore) findUser(ctx context.Context, query string, args ...any) (User, error) {
	return scanUser(s.pool.QueryRow(ctx, query, args...))
}

type rowScanner interface {
	Scan(dest ...any) error
}

func scanUser(row rowScanner) (User, error) {
	var id pgtype.UUID
	var user User
	if err := row.Scan(&id, &user.Region, &user.CountryCode, &user.Phone, &user.Email, &user.Name, &user.AvatarURL, &user.PasswordHash, &user.WechatOpenID, &user.WechatUnionID, &user.GoogleSubject, &user.Status); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return User{}, design.ErrNotFound
		}
		return User{}, err
	}
	user.ID = fromPGUUID(id)
	return user, nil
}

func scanDeviceSession(row rowScanner) (DeviceSession, error) {
	var userID pgtype.UUID
	var session DeviceSession
	if err := row.Scan(
		&session.ID,
		&userID,
		&session.DeviceType,
		&session.System,
		&session.Browser,
		&session.Client,
		&session.UserAgent,
		&session.CreatedAt,
		&session.LastSeenAt,
		&session.ExpiresAt,
		&session.RevokedAt,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return DeviceSession{}, design.ErrNotFound
		}
		return DeviceSession{}, err
	}
	session.UserID = fromPGUUID(userID)
	return session, nil
}

func toPGUUID(id string) pgtype.UUID {
	parsed, err := uuid.Parse(normalizeUUID(id))
	if err != nil {
		parsed = uuid.MustParse(design.DefaultUserID)
	}
	return pgtype.UUID{Bytes: parsed, Valid: true}
}

func fromPGUUID(value pgtype.UUID) string {
	if !value.Valid {
		return design.DefaultUserID
	}
	return uuid.UUID(value.Bytes).String()
}